mem pools expands from high addr, dynamic mem expands from low addr

mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.cc

@@ -24,7 +24,6 @@ namespace ascend {
 constexpr uint64_t kAscendDeviceMemGB = 30;
 constexpr uint64_t kMemSizeGB = 30;
 constexpr uint64_t kAscendDeviceMemSize = (kAscendDeviceMemGB << kMemSizeGB);
-constexpr uint64_t kReservedMemorySize = 10 * 1024 * 1024;
 
 void AscendMemoryManager::MallocDeviceMemory() {
   auto context_mem = GetDeviceMemSizeFromContext();
@@ -35,8 +34,9 @@ void AscendMemoryManager::MallocDeviceMemory() {
     MS_EXCEPTION(DeviceProcessError) << "rtMalloc mem size[" << device_mem_size_ << "] fail, ret[" << ret << "]";
   }
 
-  dynamic_mem_offset_ = device_mem_size_ - kReservedMemorySize;
+  AscendMemoryPool::GetInstance().set_device_mem_size(device_mem_size_);
   AscendMemoryPool::GetInstance().set_device_mem_pool_base(device_mem_base_);
+  AscendMemoryPool::GetInstance().set_device_mem_pool_offset(device_mem_size_);
   AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_);
 }
 
@@ -69,18 +69,11 @@ void AscendMemoryManager::FreeDeviceMemory() {
     }
     device_mem_base_ = nullptr;
   }
-  if (device_mem_pool_base_ != nullptr) {
-    auto ret = rtFree(device_mem_pool_base_);
-    if (ret != RT_ERROR_NONE) {
-      MS_LOG(ERROR) << "rtFree mem size[" << device_mem_pool_size_ << "] fail, ret[" << ret << "]";
-    }
-    device_mem_pool_base_ = nullptr;
-  }
 }
 
 void AscendMemoryManager::ResetDynamicMemory() {
   total_dynamic_size_ = 0;
-  dynamic_mem_offset_ = device_mem_size_ - kReservedMemorySize;
+  dynamic_mem_offset_ = 0;
   AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_);
 }
 
@@ -99,7 +92,7 @@ uint8_t *AscendMemoryManager::MallocStaticMem(size_t size, bool communication_me
 
   auto device_mem_pool_offset = AscendMemoryPool::GetInstance().device_mem_pool_offset();
   MS_LOG(INFO) << "Malloc Memory: Static, total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_
-               << "] memory pool[" << device_mem_pool_offset << "])"
+               << "] memory pool[" << device_mem_size_ - device_mem_pool_offset << "])"
                << " malloc [" << align_size << "] communication_mem: " << communication_mem;
 
   if (communication_mem) {
@@ -123,15 +116,11 @@ uint8_t *AscendMemoryManager::MallocDynamicMem(size_t size, bool communication_m
   MS_LOG(INFO) << "Malloc Memory: Dynamic, total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_
                << "] memory pool[" << device_mem_pool_offset << "])"
                << " malloc [" << align_size << "] communication_mem: " << communication_mem;
-
-  if (dynamic_mem_offset_ < align_size) {
-    MS_LOG(EXCEPTION) << "Out of memory!!! total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_
-                      << "]) malloc [" << align_size << "] failed!";
-  }
-  auto new_offset = dynamic_mem_offset_ - align_size;
-  if (new_offset <= device_mem_pool_offset) {
+  auto offset = dynamic_mem_offset_;
+  auto new_offset = dynamic_mem_offset_ + align_size;
+  if (new_offset >= device_mem_pool_offset) {
     MS_LOG(EXCEPTION) << "Out of memory!!! total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_
-                      << "] memory pool[" << device_mem_pool_offset << "])"
+                      << "] memory pool[" << device_mem_size_ - device_mem_pool_offset << "])"
                       << " malloc [" << align_size << "] failed!";
   }
   total_dynamic_size_ += align_size;
@@ -139,9 +128,9 @@ uint8_t *AscendMemoryManager::MallocDynamicMem(size_t size, bool communication_m
   AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_);
   if (communication_mem) {
     // create protect area [kMemAlignSize -- data -- kMemAlignSize]
-    return device_mem_base_ + dynamic_mem_offset_ + kMemAlignSize;
+    return device_mem_base_ + offset + kMemAlignSize;
   } else {
-    return device_mem_base_ + dynamic_mem_offset_;
+    return device_mem_base_ + offset;
   }
 }
 }  // namespace ascend

mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc

@@ -25,13 +25,13 @@ size_t AscendMemoryPool::AllocDeviceMem(size_t size, DeviceMemPtr *addr) {
   if (size == 0) {
     MS_LOG(EXCEPTION) << "Failed to alloc memory pool resource, the size is zero!";
   }
-  if (device_mem_pool_offset_ + size >= graph_dynamic_mem_offset_) {
+  if (device_mem_pool_offset_ - size <= graph_dynamic_mem_offset_) {
     MS_LOG(EXCEPTION) << "Failed to alloc memory pool memory, the current device_mem_pool_offset_ ["
                       << device_mem_pool_offset_ << "], current graph_dynamic_mem_offset_ " << graph_dynamic_mem_offset_
                       << "], need memory size [" << size << "]";
   }
+  device_mem_pool_offset_ -= size;
   *addr = device_mem_pool_base_ + device_mem_pool_offset_;
-  device_mem_pool_offset_ += size;
   if (*addr == nullptr) {
     MS_LOG(EXCEPTION) << "Alloc device memory pool address is nullptr, failed to alloc memory pool resource!";
   }
@@ -57,6 +57,12 @@ void AscendMemoryPool::set_device_mem_pool_base(uint8_t *device_mem_pool_base) {
   device_mem_pool_base_ = device_mem_pool_base;
 }
 
+void AscendMemoryPool::set_device_mem_size(uint64_t device_mem_size) { device_mem_size_ = device_mem_size; }
+
+void AscendMemoryPool::set_device_mem_pool_offset(uint64_t device_mem_pool_offset) {
+  device_mem_pool_offset_ = device_mem_pool_offset;
+}
+
 void AscendMemoryPool::set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset) {
   graph_dynamic_mem_offset_ = graph_dynamic_mem_offset;
 }
@@ -64,14 +70,14 @@ void AscendMemoryPool::set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_o
 uint64_t AscendMemoryPool::device_mem_pool_offset() const { return device_mem_pool_offset_; }
 
 size_t AscendMemoryPool::free_mem_size() {
-  if (graph_dynamic_mem_offset_ <= device_mem_pool_offset_) {
+  if (graph_dynamic_mem_offset_ >= device_mem_pool_offset_) {
     MS_LOG(EXCEPTION) << "graph dynamic mem offset [" << graph_dynamic_mem_offset_
                       << "] less than or equal to device mem pool offset [" << device_mem_pool_offset_ << "]!";
   }
-  return graph_dynamic_mem_offset_ - device_mem_pool_offset_;
+  return device_mem_pool_offset_ - graph_dynamic_mem_offset_;
 }
 
-size_t AscendMemoryPool::total_mem_size() { return graph_dynamic_mem_offset_ == 0 ? 0 : graph_dynamic_mem_offset_ - 1; }
+size_t AscendMemoryPool::total_mem_size() { return device_mem_size_ - graph_dynamic_mem_offset_; }
 }  // namespace ascend
 }  // namespace device
 }  // namespace mindspore

mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h

@@ -31,7 +31,9 @@ class AscendMemoryPool : public DynamicMemPoolBestFit {
 
   size_t AllocDeviceMem(size_t size, DeviceMemPtr *addr) override;
   bool FreeDeviceMem(const DeviceMemPtr &addr) override;
+  void set_device_mem_size(uint64_t device_mem_size);
   void set_device_mem_pool_base(uint8_t *device_mem_pool_base);
+  void set_device_mem_pool_offset(uint64_t device_mem_pool_offset);
   void set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset);
 
   uint64_t device_mem_pool_offset() const;
@@ -52,6 +54,7 @@ class AscendMemoryPool : public DynamicMemPoolBestFit {
  private:
   AscendMemoryPool() = default;
   uint8_t *device_mem_pool_base_{nullptr};
+  uint64_t device_mem_size_{0};
   uint64_t device_mem_pool_offset_{0};
   uint64_t graph_dynamic_mem_offset_{0};
 };