add get_env_or_default func

DeepLink-org · Apr 9, 2024 · a580fbf · a580fbf
1 parent 4128475
commit a580fbf
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 22 deletions.
diff --git a/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp b/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp
@@ -7,21 +7,15 @@
 #include <utility>
 #include <vector>
 
+#include "csrc_dipu/utils/env.hpp"
+
 #include "DIPUCachingAllocator.h"
 #include "DIPUSpinMutex.h"
 
 namespace dipu {
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-const size_t kMaxExtendSize = []() {
-  size_t maxExtendSize = 1024;
-  const char* env = std::getenv("DIPU_MAX_EXTEND_SIZE");
-  if (env != nullptr) {
-    maxExtendSize = std::atoi(env);
-  }
-  maxExtendSize = maxExtendSize << 20U;
-  return maxExtendSize;
-}();
+const size_t kMaxExtendSize = get_env_or_default("DIPU_MAX_EXTEND_SIZE", 1024);
 
 class BFCachingAllocatorImpl {
  public:
@@ -434,14 +428,15 @@ class BFCachingAllocator : public CacheAllocator {
   }
 
   void empty_resource_pool() const {
+    using namespace std::chrono_literals;
     std::lock_guard<mutex_t> lk(resource_pool_mutex_);
-    auto start = std::chrono::system_clock::now();
+    auto start = std::chrono::steady_clock::now();
+    constexpr auto maxWaitTime = 32us;
     while (!async_mem_pool()->empty()) {
       if (!async_mem_pool()->ready()) {
-        auto elasped = std::chrono::duration_cast<std::chrono::microseconds>(
-                           std::chrono::system_clock::now() - start)
-                           .count();
-        if (elasped < 32) {
+        auto now = std::chrono::steady_clock::now();
+        auto elasped = now - start;
+        if (elasped < maxWaitTime) {
           std::this_thread::yield();
           continue;
         } else {

diff --git a/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.cpp b/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.cpp
@@ -12,21 +12,16 @@
 
 #include "csrc_dipu/base/basedef.h"
 #include "csrc_dipu/runtime/devproxy/deviceproxy.h"
+#include "csrc_dipu/utils/env.hpp"
 
 namespace dipu {
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 std::mutex DIPURawDeviceAllocator::mutex_;
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-size_t kMaxAsyncResourcePoolLength = []() {
-  size_t maxAsyncResourcePoolLength = 64;
-  const char* env = std::getenv("DIPU_MAX_ASYNC_RESOURCE_POOL_LENGTH");
-  if (env != nullptr) {
-    maxAsyncResourcePoolLength = std::atoi(env);
-  }
-  return maxAsyncResourcePoolLength;
-}();
+size_t kMaxAsyncResourcePoolLength =
+    get_env_or_default("DIPU_MAX_ASYNC_RESOURCE_POOL_LENGTH", 64);
 
 namespace {
 

diff --git a/dipu/torch_dipu/csrc_dipu/utils/env.hpp b/dipu/torch_dipu/csrc_dipu/utils/env.hpp
@@ -0,0 +1,18 @@
+// Copyright (c) 2024, DeepLink.
+#pragma once
+#include <sstream>
+
+namespace dipu {
+
+template <typename T>
+T get_env_or_default(const char* env_name, const T& defalut_value) {
+  const char* env = std::getenv(env_name);
+  if (env == nullptr) {
+    return defalut_value;
+  }
+  T value;
+  std::istringstream(env) >> value;
+  return value;
+}
+
+}  // namespace dipu