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

pull/4105/head
liangzelang 5 years ago
parent b0b4fa08b3
commit fe1f36ea5c

@ -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

@ -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

@ -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};
};

Loading…
Cancel
Save