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
17 #include "src/extendrt/kernel/ascend/plugin/ascend_allocator_plugin.h"
18 #include <memory>
19 #if !defined(_WIN32)
20 #include "src/extendrt/cxx_api/dlutils.h"
21 #endif
22
23 namespace mindspore::kernel {
24 namespace {
25 constexpr auto kAscendkernelPluginSoNmae = "libascend_kernel_plugin.so";
26 constexpr auto kFunCreateAscendAllocatorPluginImpl = "CreateAclAllocator";
27 #if !defined(_WIN32)
28 std::mutex mutex_;
29 #endif
30 } // namespace
31
32 AscendAllocatorPlugin::AscendAllocatorPlugin() = default;
33
~AscendAllocatorPlugin()34 AscendAllocatorPlugin::~AscendAllocatorPlugin() {
35 #if !defined(_WIN32)
36 std::lock_guard<std::mutex> l(mutex_);
37 MS_LOG(INFO) << "AscendAllocatorPlugin::~AscendAllocatorPlugin() begin.";
38 ascend_allocator_plugin_impl_ = nullptr;
39 DLSoClose(handle_);
40 handle_ = nullptr;
41 MS_LOG(INFO) << "AscendAllocatorPlugin::~AscendAllocatorPlugin() end.";
42 #endif
43 }
44
GetInstance()45 AscendAllocatorPlugin &AscendAllocatorPlugin::GetInstance() {
46 #if !defined(_WIN32)
47 std::lock_guard<std::mutex> l(mutex_);
48 #endif
49 static AscendAllocatorPlugin instance;
50 return instance;
51 }
52
Register()53 bool AscendAllocatorPlugin::Register() {
54 #if !defined(_WIN32)
55 std::lock_guard<std::mutex> l(mutex_);
56 if (is_registered_) {
57 return true;
58 }
59 MS_LOG(INFO) << "AscendAllocatorPlugin Register.";
60 auto ret = DLSoPath({"libmindspore-lite.so", "_c_lite", "tools/converter/lib", "libmindspore_converter.so"},
61 kAscendkernelPluginSoNmae, &plugin_path_);
62 if (ret != kSuccess) {
63 MS_LOG(ERROR) << "get real path of " << kAscendkernelPluginSoNmae << " failed.";
64 return false;
65 }
66 MS_LOG(INFO) << "find ascend allocator plugin so path: " << plugin_path_;
67 void *func = nullptr;
68 ret = DLSoOpen(plugin_path_, kFunCreateAscendAllocatorPluginImpl, &handle_, &func);
69 if (ret != kSuccess) {
70 MS_LOG(ERROR) << "DLSoOpen failed, so path: " << plugin_path_
71 << " , func name: " << kFunCreateAscendAllocatorPluginImpl << ", err: " << ret.ToString();
72 return false;
73 }
74 auto create_plugin_impl_func = reinterpret_cast<AscendAllocatorPluginImpl *(*)(void)>(func);
75 if (create_plugin_impl_func == nullptr) {
76 MS_LOG(ERROR) << "cast " << kFunCreateAscendAllocatorPluginImpl << " failed.";
77 return false;
78 }
79 ascend_allocator_plugin_impl_ = std::shared_ptr<AscendAllocatorPluginImpl>(create_plugin_impl_func());
80 if (ascend_allocator_plugin_impl_ == nullptr) {
81 MS_LOG(ERROR) << "create ascend allocator plugin impl failed.";
82 return false;
83 }
84 is_registered_ = true;
85 MS_LOG(INFO) << "register ascend allocator success.";
86 #endif
87 return true;
88 }
89
GetCurrentDeviceId()90 int AscendAllocatorPlugin::GetCurrentDeviceId() {
91 #if !defined(_WIN32)
92 if (!is_registered_) {
93 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
94 return -1;
95 }
96 if (ascend_allocator_plugin_impl_ == nullptr) {
97 return -1;
98 }
99 auto device_data = ascend_allocator_plugin_impl_->GetCurrentDeviceId();
100 return device_data;
101 #endif
102 return -1;
103 }
104
Malloc(size_t size,int device_id)105 void *AscendAllocatorPlugin::Malloc(size_t size, int device_id) {
106 #if !defined(_WIN32)
107 if (!is_registered_) {
108 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
109 return nullptr;
110 }
111 if (device_id < -1) {
112 MS_LOG(ERROR) << "device id must more than 0";
113 return nullptr;
114 }
115 if (ascend_allocator_plugin_impl_ == nullptr) {
116 MS_LOG(ERROR) << "ascend_allocator_plugin_impl_ is nullptr.";
117 return nullptr;
118 }
119 auto device_data = ascend_allocator_plugin_impl_->Malloc(size, device_id);
120 return device_data;
121 #endif
122 return nullptr;
123 }
124
Free(void * device_data,int device_id)125 void AscendAllocatorPlugin::Free(void *device_data, int device_id) {
126 #if !defined(_WIN32)
127 if (!is_registered_) {
128 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
129 return;
130 }
131 if (ascend_allocator_plugin_impl_ == nullptr) {
132 MS_LOG(ERROR) << "ascend_allocator_plugin_impl_ is nullptr.";
133 return;
134 }
135 if (device_data == nullptr) {
136 MS_LOG(ERROR) << "device data is nullptr.";
137 return;
138 }
139 ascend_allocator_plugin_impl_->Free(device_data, device_id);
140 #endif
141 return;
142 }
143
MallocHost(size_t size)144 void *AscendAllocatorPlugin::MallocHost(size_t size) {
145 #if !defined(_WIN32)
146 if (!is_registered_) {
147 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
148 return nullptr;
149 }
150 if (ascend_allocator_plugin_impl_ == nullptr) {
151 MS_LOG(ERROR) << "ascend_allocator_plugin_impl_ is nullptr.";
152 return nullptr;
153 }
154 auto device_data = ascend_allocator_plugin_impl_->MallocHost(size);
155 return device_data;
156 #endif
157 return nullptr;
158 }
159
FreeHost(void * host_data)160 void AscendAllocatorPlugin::FreeHost(void *host_data) {
161 #if !defined(_WIN32)
162 if (!is_registered_) {
163 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
164 return;
165 }
166 if (ascend_allocator_plugin_impl_ == nullptr) {
167 MS_LOG(ERROR) << "ascend_allocator_plugin_impl_ is nullptr.";
168 return;
169 }
170 if (host_data == nullptr) {
171 MS_LOG(ERROR) << "host data is nullptr.";
172 return;
173 }
174 ascend_allocator_plugin_impl_->FreeHost(host_data);
175 #endif
176 return;
177 }
178
CopyDeviceDataToHost(void * device_data,void * host_data,size_t data_size,int device_id)179 Status AscendAllocatorPlugin::CopyDeviceDataToHost(void *device_data, void *host_data, size_t data_size,
180 int device_id) {
181 #if !defined(_WIN32)
182 if (!is_registered_) {
183 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
184 return kLiteMemoryFailed;
185 }
186 if (device_data == nullptr) {
187 MS_LOG(INFO) << "device data is nullptr.";
188 return kLiteMemoryFailed;
189 }
190 if (ascend_allocator_plugin_impl_ == nullptr) {
191 return kLiteMemoryFailed;
192 }
193 return ascend_allocator_plugin_impl_->CopyDeviceDataToHost(device_data, host_data, data_size, device_id);
194 #endif
195 return kSuccess;
196 }
197
CopyDeviceDataToDevice(void * src_device,void * dst_device,size_t src_data_size,size_t dst_data_size,int src_device_id,int dst_device_id)198 Status AscendAllocatorPlugin::CopyDeviceDataToDevice(void *src_device, void *dst_device, size_t src_data_size,
199 size_t dst_data_size, int src_device_id, int dst_device_id) {
200 #if !defined(_WIN32)
201 if (!is_registered_) {
202 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
203 return kLiteMemoryFailed;
204 }
205 if (src_device_id < -1 || dst_device_id < -1) {
206 MS_LOG(ERROR) << "device id is wrong, src device id: " << src_device_id << ", dst device id: " << dst_device_id;
207 return kLiteError;
208 }
209 if (dst_device == nullptr || src_device == nullptr) {
210 MS_LOG(INFO) << "device data is nullptr.";
211 return kLiteMemoryFailed;
212 }
213 if (ascend_allocator_plugin_impl_ == nullptr) {
214 return kLiteMemoryFailed;
215 }
216 return ascend_allocator_plugin_impl_->CopyDeviceDataToDevice(src_device, dst_device, src_data_size, dst_data_size,
217 src_device_id, dst_device_id);
218 #endif
219 return kSuccess;
220 }
221
CopyHostDataToDevice(void * host_data,void * device_data,size_t data_size)222 Status AscendAllocatorPlugin::CopyHostDataToDevice(void *host_data, void *device_data, size_t data_size) {
223 #if !defined(_WIN32)
224 if (!is_registered_) {
225 MS_LOG(ERROR) << "AscendAllocatorPlugin is not registered.";
226 return kLiteMemoryFailed;
227 }
228 if (device_data == nullptr) {
229 MS_LOG(INFO) << "device data is nullptr.";
230 return kLiteMemoryFailed;
231 }
232 if (ascend_allocator_plugin_impl_ == nullptr) {
233 return kLiteMemoryFailed;
234 }
235 return ascend_allocator_plugin_impl_->CopyHostDataToDevice(host_data, device_data, data_size);
236 #endif
237 return kSuccess;
238 }
239 } // namespace mindspore::kernel
240