1 /**
2 * Copyright 2023 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 #include "plugin/device/ascend/hal/device/ascend_gmem_adapter.h"
17 #include <pthread.h>
18 #include <sys/types.h>
19 #include <fcntl.h>
20 #include <sys/mman.h>
21 #include <unistd.h>
22
23 #include <tuple>
24
25 #include "plugin/device/ascend/hal/device/ascend_stream_manager.h"
26 #include "utils/convert_utils_base.h"
27
28 namespace mindspore {
29 namespace device {
30 namespace ascend {
31 static constexpr const char kGMemLibName[] = "libgmem.so";
32 static constexpr const char kMsEnableGmem[] = "MS_ENABLE_GMEM";
33 constexpr uint64_t kAscendMmapAlignSize = 1 << 21;
34 constexpr int kMapPeerShared = 0x8000000;
35
GetRoundUpAlignSize(size_t input_size) const36 const size_t AscendGmemAdapter::GetRoundUpAlignSize(size_t input_size) const {
37 return (input_size + kAscendMmapAlignSize - 1) & ~(kAscendMmapAlignSize - 1);
38 }
39
GetRoundDownAlignSize(size_t input_size) const40 const size_t AscendGmemAdapter::GetRoundDownAlignSize(size_t input_size) const {
41 return input_size & ~(kAscendMmapAlignSize - 1);
42 }
43
AllocDeviceMem(size_t size,DeviceMemPtr * addr) const44 size_t AscendGmemAdapter::AllocDeviceMem(size_t size, DeviceMemPtr *addr) const {
45 size_t align_size = GetRoundUpAlignSize(size);
46 uint8_t *alloc_addr = MmapMemory(align_size, nullptr);
47 if (alloc_addr == nullptr) {
48 MS_LOG(WARNING) << "Malloc memory failed.";
49 return 0;
50 }
51 *addr = alloc_addr;
52 return align_size;
53 }
54
EagerFreeDeviceMem(const DeviceMemPtr addr,const size_t size) const55 size_t AscendGmemAdapter::EagerFreeDeviceMem(const DeviceMemPtr addr, const size_t size) const {
56 MS_LOG(DEBUG) << "Enter ascend eager free device mem, addr : " << addr << ", size : " << size << ".";
57 if (size <= 0) {
58 MS_LOG(WARNING) << "Size is non positive.";
59 return 0;
60 }
61 size_t addr_size_t = reinterpret_cast<size_t>(addr);
62 // Adjust addr -> round up addr, size -> round down size.
63 size_t from_addr = GetRoundUpAlignSize(addr_size_t);
64 size_t end_addr = GetRoundDownAlignSize(addr_size_t + size);
65 if (end_addr <= from_addr) {
66 MS_LOG(DEBUG) << "End addr : " << end_addr << " is not bigger than from_addr : " << from_addr << ".";
67 return 0;
68 }
69 size_t real_size = end_addr - from_addr;
70 int ret = free_eager_(from_addr, SizeToUlong(real_size), nullptr);
71 return ret != 0 ? 0 : real_size;
72 }
73
MmapMemory(size_t size,void * addr) const74 uint8_t *AscendGmemAdapter::MmapMemory(size_t size, void *addr) const {
75 MS_LOG(DEBUG) << "Enter mmap memory, size : " << size << ".";
76 if (size <= 0) {
77 MS_LOG(ERROR) << "Size : " << size << " is non positive.";
78 return nullptr;
79 }
80
81 int flags = MAP_PRIVATE | MAP_ANONYMOUS | kMapPeerShared;
82 int prot = PROT_READ | PROT_WRITE;
83 void *mapped_addr = mmap(addr, size, prot, flags, -1, 0);
84 if (mapped_addr == MAP_FAILED) {
85 MS_LOG(EXCEPTION) << "Mmap failed.";
86 }
87 return static_cast<uint8_t *>(mapped_addr);
88 }
89
MunmapMemory(void * addr,const size_t size) const90 bool AscendGmemAdapter::MunmapMemory(void *addr, const size_t size) const {
91 MS_LOG(DEBUG) << "Enter munmap memory, addr : " << addr << ", size : " << size << ".";
92 auto ret = munmap(addr, size);
93 return ret != -1;
94 }
95
LoadGMemLib()96 void AscendGmemAdapter::LoadGMemLib() noexcept {
97 if (common::GetEnv(kMsEnableGmem) != "1") {
98 return;
99 }
100 MS_LOG(INFO) << "MS_ENABLE_GMEM is set, try to open gmem.";
101 gmem_handle_ = dlopen(kGMemLibName, RTLD_NOW);
102 if (gmem_handle_ != nullptr) {
103 MS_LOG(WARNING) << "Open GMem lib success, mindspore will use gmem to optimize memory usage.";
104 LIB_FUNC(GMEM_FREE_EAGER) gmem_free_eager = DlsymFuncObj(gmemFreeEager, gmem_handle_);
105 if (gmem_free_eager != nullptr) {
106 is_eager_free_enabled_ = true;
107 // Load gmem lib success, and then enable stream callback.
108 AscendStreamMng::GetInstance().enable_callback(is_eager_free_enabled_);
109 free_eager_ = gmem_free_eager;
110 } else {
111 MS_LOG(WARNING) << "Load gmem free eager failed.";
112 if (dlclose(gmem_handle_) != 0) {
113 MS_LOG(ERROR) << "Close GMem lib failed, detail : " << dlerror() << ".";
114 }
115 }
116 } else {
117 MS_LOG(DEBUG) << "Open GMem lib failed.";
118 }
119 }
120
UnloadGMemLib()121 void AscendGmemAdapter::UnloadGMemLib() noexcept {
122 if (gmem_handle_ != nullptr) {
123 MS_LOG(INFO) << "Close GMem lib.";
124 if (dlclose(gmem_handle_) != 0) {
125 MS_LOG(ERROR) << "Close GMem lib failed, detail : " << dlerror() << ".";
126 }
127 gmem_handle_ = nullptr;
128 }
129 }
130 } // namespace ascend
131 } // namespace device
132 } // namespace mindspore
133