From a580fbf1cb4b69fc0bc606a396884efe27c15a77 Mon Sep 17 00:00:00 2001
From: zhaoguochun1995 <zhaoguochun1995@163.com>
Date: Tue, 9 Apr 2024 06:54:15 +0000
Subject: [PATCH] add get_env_or_default func

---
 .../core/allocator/DIPUBFCachingAllocator.cpp | 23 ++++++++-----------
 .../core/allocator/DIPUCachingAllocator.cpp   | 11 +++------
 dipu/torch_dipu/csrc_dipu/utils/env.hpp       | 18 +++++++++++++++
 3 files changed, 30 insertions(+), 22 deletions(-)
 create mode 100644 dipu/torch_dipu/csrc_dipu/utils/env.hpp

diff --git a/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp b/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp
index 48b471f01..b0993eaa8 100644
--- 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
index 0c557bacb..54dc8b898 100644
--- a/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.cpp
+++ b/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.cpp
@@ -12,6 +12,7 @@
 
 #include "csrc_dipu/base/basedef.h"
 #include "csrc_dipu/runtime/devproxy/deviceproxy.h"
+#include "csrc_dipu/utils/env.hpp"
 
 namespace dipu {
 
@@ -19,14 +20,8 @@ namespace dipu {
 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
new file mode 100644
index 000000000..e561a429c
--- /dev/null
+++ 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