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 "plugin/device/gpu/hal/device/gpu_memory_allocator.h"
19 #include "plugin/device/gpu/hal/device/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 constexpr float kReservedMemoryRatio = 0.0625; // 1/16
29 static const size_t MEM_ALIGN_SIZE = 512;
30
Init()31 bool GPUMemoryAllocator::Init() {
32 size_t total_size = CudaDriver::total_mem_size();
33 size_t free_size = CudaDriver::free_mem_size();
34 auto context_ptr = MsContext::GetInstance();
35 MS_EXCEPTION_IF_NULL(context_ptr);
36 limited_device_memory_ = context_ptr->get_param<float>(MS_CTX_MAX_DEVICE_MEMORY);
37 available_device_memory_ = FloatToSize(limited_device_memory_ * kGBToByte);
38 if (total_size > 0 && free_size > 0 && available_device_memory_ > 0) {
39 MS_LOG(INFO) << "GPU device total memory size " << total_size << ", current free memory size " << free_size
40 << ", set max available memory size " << available_device_memory_ << ".";
41 } else {
42 MS_LOG(EXCEPTION) << "#umsg#GPU memory error:#umsg#The total size or free size or max_device_memory size of GPU "
43 "memory can't be zero, total memory size "
44 << total_size << ", current free memory size " << free_size << ", set max available memory size "
45 << available_device_memory_ << ".";
46 }
47 // In gpu mode, recommend 1/16 reserved for other cuda functions
48 if (available_device_memory_ > total_size) {
49 size_t recommend_mem_size_for_others = FloatToSize(total_size * kReservedMemoryRatio);
50 SetMemPoolBlockSize(std::min(available_device_memory_, total_size - recommend_mem_size_for_others));
51 } else {
52 SetMemPoolBlockSize(std::min(available_device_memory_, total_size));
53 }
54 return true;
55 }
56
CheckMaxDeviceMemory() const57 void GPUMemoryAllocator::CheckMaxDeviceMemory() const {
58 auto context_ptr = MsContext::GetInstance();
59 MS_EXCEPTION_IF_NULL(context_ptr);
60 auto max_device_memory = context_ptr->get_param<float>(MS_CTX_MAX_DEVICE_MEMORY);
61 // Currently not support modifying the max device memory.
62 if (!common::IsFloatEqual(limited_device_memory_, max_device_memory)) {
63 MS_LOG(EXCEPTION) << "#umsg#Can't change or set context param max_device_memory during running:#umsg#Currently "
64 "effective max_device_memory("
65 << limited_device_memory_ << "GB), set new max_device_memory(" << max_device_memory
66 << "GB) failed.";
67 }
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;
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) << "#umsg#GPU memory error:#umsg#The memory alloc size is 0.";
80 }
81 auto free_size = free_mem_size();
82 if (size > free_size) {
83 MS_LOG(EXCEPTION) << "#umsg#Memory not enough:#umsg#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) << "#umsg#Memory not enough:#umsg#Alloc device memory[" << size << "] failed.";
90 }
91 total_used_device_memory_ += alloc_size;
92 available_device_memory_ -= alloc_size;
93 MS_LOG(INFO) << "Cuda current free memory size[" << free_size << "], alloc size[" << alloc_size
94 << "], left free memory size[" << free_size - alloc_size << "]"
95 << ".Total used size[" << total_used_device_memory_ << "].";
96 return alloc_size;
97 }
98
FreeDeviceMem(const DeviceMemPtr & addr)99 bool GPUMemoryAllocator::FreeDeviceMem(const DeviceMemPtr &addr) { return CudaDriver::FreeDeviceMem(addr); }
100
free_mem_size()101 size_t GPUMemoryAllocator::free_mem_size() { return std::min(CudaDriver::free_mem_size(), available_device_memory_); }
102
AlignMemorySize(size_t size) const103 size_t GPUMemoryAllocator::AlignMemorySize(size_t size) const {
104 if (size == 0) {
105 return MEM_ALIGN_SIZE;
106 }
107 return ((size + MEM_ALIGN_SIZE - 1) / MEM_ALIGN_SIZE) * MEM_ALIGN_SIZE;
108 }
109 } // namespace gpu
110 } // namespace device
111 } // namespace mindspore
112