1 /**
2 * Copyright 2020 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <algorithm>
18 #include "runtime/device/ascend/ascend_memory_pool.h"
19 #include "runtime/mem.h"
20 #include "runtime/device/ascend/ascend_kernel_runtime.h"
21 #include "utils/log_adapter.h"
22
23 namespace mindspore {
24 namespace device {
25 namespace ascend {
26 // The minimum unit size (256MB) of memory block used for dynamic extend.
27 static const size_t ASCEND_DYNAMIC_MEM_ALLOC_UNIT_SIZE = 256 << 20;
28 // The minimum unit size (8MB) of memory block used for dynamic extend in graph mode.
29 static const size_t ASCEND_DYNAMIC_MEM_ALLOC_UNIT_SIZE_FOR_GRAPH = 8 << 20;
30
Init(uint8_t * device_mem_base,uint64_t device_mem_size,uint64_t dynamic_mem_offset)31 void AscendMemoryPool::Init(uint8_t *device_mem_base, uint64_t device_mem_size, uint64_t dynamic_mem_offset) {
32 static bool initialized = false;
33 if (initialized) {
34 return;
35 }
36
37 MS_EXCEPTION_IF_NULL(device_mem_base);
38 set_device_mem_pool_base(device_mem_base);
39
40 if (dynamic_mem_offset > device_mem_size) {
41 MS_LOG(EXCEPTION) << "Dynamic memory offset: " << dynamic_mem_offset
42 << " exceed the device memory size: " << device_mem_size;
43 }
44 set_device_mem_size(device_mem_size);
45 set_device_mem_pool_offset(device_mem_size);
46 set_graph_dynamic_mem_offset(dynamic_mem_offset);
47 initialized = true;
48 }
49
CalMemBlockAllocSize(size_t size)50 size_t AscendMemoryPool::CalMemBlockAllocSize(size_t size) {
51 auto device_free_mem_size = free_mem_size();
52 if (device_free_mem_size < size) {
53 MS_LOG(EXCEPTION) << "Memory not enough: current free memory size[" << device_free_mem_size
54 << "] is smaller than required size[" << size << "], dynamic offset ["
55 << graph_dynamic_mem_offset_ << "] memory pool offset["
56 << device_mem_size_ - device_mem_pool_offset_ << "])";
57 return 0;
58 }
59 auto alloc_mem_size = ASCEND_DYNAMIC_MEM_ALLOC_UNIT_SIZE;
60 auto ms_context = MsContext::GetInstance();
61 MS_EXCEPTION_IF_NULL(ms_context);
62 const bool pynative_mode = (ms_context->get_param<int>(MS_CTX_EXECUTION_MODE) == kPynativeMode);
63 if (pynative_mode) {
64 // Growing at twice of alloc size
65 constexpr size_t kDouble = 2;
66 while (alloc_mem_size < size) {
67 alloc_mem_size = alloc_mem_size * kDouble;
68 }
69 } else {
70 alloc_mem_size = ASCEND_DYNAMIC_MEM_ALLOC_UNIT_SIZE_FOR_GRAPH;
71 while (alloc_mem_size < size) {
72 alloc_mem_size = alloc_mem_size + ASCEND_DYNAMIC_MEM_ALLOC_UNIT_SIZE_FOR_GRAPH;
73 }
74 }
75 alloc_mem_size = std::min(alloc_mem_size, device_free_mem_size);
76 return alloc_mem_size;
77 }
78
AllocDeviceMem(size_t size,DeviceMemPtr * addr)79 size_t AscendMemoryPool::AllocDeviceMem(size_t size, DeviceMemPtr *addr) {
80 MS_LOG(INFO) << "Malloc Memory: Pool, total[" << device_mem_size_ << "] (dynamic[" << graph_dynamic_mem_offset_
81 << "] memory pool[" << device_mem_size_ - device_mem_pool_offset_ << "])"
82 << " malloc [" << size << "]";
83
84 if (size == 0) {
85 MS_LOG(EXCEPTION) << "Failed to alloc memory pool resource, the size is zero!";
86 }
87
88 if (device_mem_pool_offset_ - size < graph_dynamic_mem_offset_) {
89 MS_LOG(EXCEPTION) << "Failed to alloc memory pool memory, the current device_mem_pool_offset_ ["
90 << device_mem_pool_offset_ << "], current graph_dynamic_mem_offset_ " << graph_dynamic_mem_offset_
91 << "], need memory size [" << size << "]";
92 }
93 device_mem_pool_offset_ -= size;
94 *addr = device_mem_pool_base_ + device_mem_pool_offset_;
95 if (*addr == nullptr) {
96 MS_LOG(EXCEPTION) << "Alloc device memory pool address is nullptr, failed to alloc memory pool resource!";
97 }
98 return size;
99 }
100
FreeDeviceMem(const DeviceMemPtr & addr)101 bool AscendMemoryPool::FreeDeviceMem(const DeviceMemPtr &addr) {
102 MS_EXCEPTION_IF_NULL(addr);
103 return true;
104 }
105
ResetIdleMemBuf()106 void AscendMemoryPool::ResetIdleMemBuf() {
107 auto idle_mem_buf_map = DynamicMemPoolBestFit::global_idle_mem_buf_map();
108 for (auto &it : idle_mem_buf_map) {
109 MS_EXCEPTION_IF_NULL(it.second);
110 (void)rtMemset(it.second->device_addr_, it.first, 0, it.first);
111 }
112 }
113
AlignMemorySize(size_t size) const114 size_t AscendMemoryPool::AlignMemorySize(size_t size) const {
115 if (size == 0) {
116 MS_LOG(EXCEPTION) << "The align memory size is a zero !";
117 }
118 return size;
119 }
120
set_device_mem_pool_base(uint8_t * device_mem_pool_base)121 void AscendMemoryPool::set_device_mem_pool_base(uint8_t *device_mem_pool_base) {
122 MS_EXCEPTION_IF_NULL(device_mem_pool_base);
123 device_mem_pool_base_ = device_mem_pool_base;
124 }
125
set_device_mem_size(uint64_t device_mem_size)126 void AscendMemoryPool::set_device_mem_size(uint64_t device_mem_size) { device_mem_size_ = device_mem_size; }
127
set_device_mem_pool_offset(uint64_t device_mem_pool_offset)128 void AscendMemoryPool::set_device_mem_pool_offset(uint64_t device_mem_pool_offset) {
129 device_mem_pool_offset_ = device_mem_pool_offset;
130 }
131
set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset)132 void AscendMemoryPool::set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset) {
133 graph_dynamic_mem_offset_ = graph_dynamic_mem_offset;
134 }
135
device_mem_pool_offset() const136 uint64_t AscendMemoryPool::device_mem_pool_offset() const { return device_mem_pool_offset_; }
137
free_mem_size()138 size_t AscendMemoryPool::free_mem_size() {
139 if (graph_dynamic_mem_offset_ >= device_mem_pool_offset_) {
140 MS_LOG(EXCEPTION) << "graph dynamic mem offset [" << graph_dynamic_mem_offset_
141 << "] less than or equal to device mem pool offset [" << device_mem_pool_offset_ << "]!";
142 }
143 return device_mem_pool_offset_ - graph_dynamic_mem_offset_;
144 }
145
total_mem_size()146 size_t AscendMemoryPool::total_mem_size() { return device_mem_size_ - graph_dynamic_mem_offset_; }
147 } // namespace ascend
148 } // namespace device
149 } // namespace mindspore
150