1 /**
2 * Copyright 2019 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/gpu/gpu_memory_allocator.h"
19 #include "runtime/device/gpu/cuda_driver.h"
20 #include "utils/log_adapter.h"
21 #include "utils/ms_context.h"
22 #include "utils/convert_utils_base.h"
23
24 namespace mindspore {
25 namespace device {
26 namespace gpu {
27 const size_t kGBToByte = 1024 << 20;
28
Init()29 bool GPUMemoryAllocator::Init() {
30 size_t total_size = total_mem_size();
31 size_t free_size = CudaDriver::free_mem_size();
32 auto context_ptr = MsContext::GetInstance();
33 MS_EXCEPTION_IF_NULL(context_ptr);
34 limited_device_memory_ = context_ptr->get_param<float>(MS_CTX_MAX_DEVICE_MEMORY);
35 available_device_memory_ = FloatToSize(limited_device_memory_ * kGBToByte);
36 if (total_size > 0 && free_size > 0 && available_device_memory_ > 0) {
37 MS_LOG(INFO) << "GPU device total memory size " << total_size << ", current free memory size " << free_size
38 << ", set max available memory size " << available_device_memory_ << ".";
39 } else {
40 MS_LOG(EXCEPTION)
41 << "The total size or free size or max_device_memory size of GPU memory can't be zero, total memory size "
42 << total_size << ", current free memory size " << free_size << ", set max available memory size "
43 << available_device_memory_ << ".";
44 }
45 return true;
46 }
47
CheckMaxDeviceMemory() const48 void GPUMemoryAllocator::CheckMaxDeviceMemory() const {
49 auto context_ptr = MsContext::GetInstance();
50 MS_EXCEPTION_IF_NULL(context_ptr);
51 auto max_device_memory = context_ptr->get_param<float>(MS_CTX_MAX_DEVICE_MEMORY);
52 // Currently not support modifying the max device memory.
53 if (limited_device_memory_ != max_device_memory) {
54 MS_LOG(EXCEPTION)
55 << "Can't change or set context param max_device_memory during running, currently effective max_device_memory("
56 << limited_device_memory_ << "GB), set new max_device_memory(" << max_device_memory << "GB) failed.";
57 }
58 }
59
Finalize()60 bool GPUMemoryAllocator::Finalize() {
61 if (buffer_q_addr_ != nullptr) {
62 if (!CudaDriver::FreeDeviceMem(buffer_q_addr_)) {
63 MS_LOG(ERROR) << "Could not free buffer queue memory.";
64 return false;
65 }
66 }
67 return true;
68 }
69
AllocBufferQueueMem(size_t size,DeviceMemPtr * addr)70 bool GPUMemoryAllocator::AllocBufferQueueMem(size_t size, DeviceMemPtr *addr) {
71 auto alloc_size = AllocDeviceMem(size, addr);
72 buffer_q_addr_ = *addr;
73 // Buffer queue needs to ensure that the alloc_size and size is equal.
74 return (alloc_size == size) ? true : false;
75 }
76
AllocDeviceMem(size_t size,DeviceMemPtr * addr)77 size_t GPUMemoryAllocator::AllocDeviceMem(size_t size, DeviceMemPtr *addr) {
78 if (size == 0) {
79 MS_LOG(EXCEPTION) << "The memory alloc size is 0.";
80 }
81 auto free_size = free_mem_size();
82 if (size > free_size) {
83 MS_LOG(EXCEPTION) << "Memory not enough: current free memory size[" << free_size
84 << "] is smaller than required size[" << size << "].";
85 }
86
87 auto alloc_size = CudaDriver::AllocDeviceMem(size, addr);
88 if (alloc_size == 0) {
89 MS_LOG(EXCEPTION) << "Alloc device memory[" << size << "] failed.";
90 }
91 total_used_device_memory_ += alloc_size;
92 available_device_memory_ -= alloc_size;
93 MS_LOG(INFO) << "Current free memory size[" << free_size - alloc_size << "], current alloc size[" << alloc_size
94 << "], total used size[" << total_used_device_memory_ << "].";
95 return alloc_size;
96 }
97
FreeDeviceMem(const DeviceMemPtr & addr)98 bool GPUMemoryAllocator::FreeDeviceMem(const DeviceMemPtr &addr) { return CudaDriver::FreeDeviceMem(addr); }
99
free_mem_size()100 size_t GPUMemoryAllocator::free_mem_size() { return std::min(CudaDriver::free_mem_size(), available_device_memory_); }
101
total_mem_size()102 size_t GPUMemoryAllocator::total_mem_size() { return CudaDriver::total_mem_size(); }
103 } // namespace gpu
104 } // namespace device
105 } // namespace mindspore
106