Line data Source code
1 : /**
2 : * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 :
11 : #include <memory>
12 : #include <mutex>
13 : #include <map>
14 : #include "common/prof_reporter.h"
15 : #include "common/log_inner.h"
16 : #include "acl_rt_impl.h"
17 : #include "common/resource_statistics.h"
18 :
19 : namespace {
20 : class AllocatorDesc {
21 : public:
22 : AllocatorDesc() = default;
23 : ~AllocatorDesc() = default;
24 2 : AllocatorDesc(aclrtAllocator allocator,
25 : aclrtAllocatorAllocFunc allocFunc,
26 : aclrtAllocatorFreeFunc freeFunc,
27 : aclrtAllocatorAllocAdviseFunc allocAdviseFunc,
28 : aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc)
29 2 : {
30 2 : this->obj = allocator;
31 2 : this->allocFunc = allocFunc;
32 2 : this->freeFunc = freeFunc;
33 2 : this->allocAdviseFunc = allocAdviseFunc;
34 2 : this->getAddrFromBlockFunc = getAddrFromBlockFunc;
35 2 : }
36 : aclrtAllocator obj;
37 : aclrtAllocatorAllocFunc allocFunc;
38 : aclrtAllocatorFreeFunc freeFunc;
39 : aclrtAllocatorAllocAdviseFunc allocAdviseFunc;
40 : aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc;
41 : };
42 : std::mutex g_AllocatorDescMutex;
43 : // The first aclrtAllocatorDesc is created by the user, while the second AllocatorDesc is a saved copy.
44 : std::map<aclrtStream, std::pair<aclrtAllocatorDesc, AllocatorDesc>> g_AllocatorDesMap;
45 : }
46 :
47 : #ifdef __cplusplus
48 : extern "C" {
49 : #endif
50 :
51 2 : aclrtAllocatorDesc aclrtAllocatorCreateDescImpl()
52 : {
53 2 : ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorCreateDesc);
54 2 : ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
55 2 : ACL_LOG_INFO("Create allocator description.");
56 2 : AllocatorDesc *allocatorDesc = new(std::nothrow) AllocatorDesc;
57 2 : ACL_CHECK_MALLOC_RESULT_REPORT_RET(allocatorDesc, sizeof(AllocatorDesc), "new", nullptr);
58 2 : ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
59 2 : return static_cast<aclrtAllocatorDesc>(allocatorDesc);
60 2 : }
61 :
62 2 : aclError aclrtAllocatorDestroyDescImpl(aclrtAllocatorDesc allocatorDesc)
63 : {
64 2 : ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorDestroyDesc);
65 2 : ACL_LOG_INFO("Destroy allocator description, allocatorDesc %p.", allocatorDesc);
66 2 : ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
67 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
68 2 : delete static_cast<AllocatorDesc *>(allocatorDesc);
69 2 : allocatorDesc = nullptr;
70 2 : ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
71 2 : return ACL_SUCCESS;
72 2 : }
73 :
74 2 : aclError aclrtAllocatorSetObjToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocator allocator)
75 : {
76 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
77 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocator);
78 2 : ACL_LOG_INFO("Set allocator to allocator description, allocatorDesc %p.", allocatorDesc);
79 2 : static_cast<AllocatorDesc *>(allocatorDesc)->obj = allocator;
80 2 : return ACL_SUCCESS;
81 : }
82 :
83 2 : aclError aclrtAllocatorSetAllocFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocFunc func)
84 : {
85 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
86 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
87 2 : ACL_LOG_INFO("Set alloc function to allocator description, allocatorDesc %p.", allocatorDesc);
88 2 : static_cast<AllocatorDesc *>(allocatorDesc)->allocFunc = func;
89 2 : return ACL_SUCCESS;
90 : }
91 :
92 2 : aclError aclrtAllocatorSetFreeFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorFreeFunc func)
93 : {
94 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
95 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
96 2 : ACL_LOG_INFO("Set free function to allocator description, allocatorDesc %p.", allocatorDesc);
97 2 : static_cast<AllocatorDesc *>(allocatorDesc)->freeFunc = func;
98 2 : return ACL_SUCCESS;
99 : }
100 :
101 2 : aclError aclrtAllocatorSetAllocAdviseFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocAdviseFunc func)
102 : {
103 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
104 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
105 2 : ACL_LOG_INFO("Set alloc advise function to allocator description, allocatorDesc %p.", allocatorDesc);
106 2 : static_cast<AllocatorDesc *>(allocatorDesc)->allocAdviseFunc = func;
107 2 : return ACL_SUCCESS;
108 : }
109 :
110 2 : aclError aclrtAllocatorSetGetAddrFromBlockFuncToDescImpl(aclrtAllocatorDesc allocatorDesc,
111 : aclrtAllocatorGetAddrFromBlockFunc func)
112 : {
113 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
114 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
115 2 : ACL_LOG_INFO("Set get_addr_from_block function to allocator description, allocatorDesc %p.", allocatorDesc);
116 2 : static_cast<AllocatorDesc *>(allocatorDesc)->getAddrFromBlockFunc = func;
117 2 : return ACL_SUCCESS;
118 : }
119 :
120 7 : aclError aclrtAllocatorRegisterImpl(aclrtStream stream, aclrtAllocatorDesc allocatorDesc)
121 : {
122 : // stream must be not null when register external allocator
123 7 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
124 6 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
125 :
126 6 : AllocatorDesc *allocDesc = static_cast<AllocatorDesc *>(allocatorDesc);
127 10 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->obj, "allocatorDesc->obj");
128 9 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->allocFunc, "allocatorDesc->allocFunc");
129 8 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->freeFunc, "allocatorDesc->freeFunc");
130 7 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->getAddrFromBlockFunc, "allocatorDesc->getAddrFromBlockFunc");
131 :
132 : AllocatorDesc allocDescCopy = AllocatorDesc(allocDesc->obj,
133 : allocDesc->allocFunc,
134 : allocDesc->freeFunc,
135 : allocDesc->allocAdviseFunc,
136 2 : allocDesc->getAddrFromBlockFunc);
137 2 : std::pair<aclrtAllocatorDesc, AllocatorDesc> allocatorDescPair(allocatorDesc, allocDescCopy);
138 2 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
139 2 : g_AllocatorDesMap[stream] = allocatorDescPair;
140 2 : ACL_LOG_INFO("Register external allocator success, stream %p, allocatorDesc %p.", stream, allocatorDesc);
141 2 : return ACL_SUCCESS;
142 2 : }
143 :
144 5 : aclError aclrtAllocatorGetByStreamImpl(aclrtStream stream,
145 : aclrtAllocatorDesc *allocatorDesc,
146 : aclrtAllocator *allocator,
147 : aclrtAllocatorAllocFunc *allocFunc,
148 : aclrtAllocatorFreeFunc *freeFunc,
149 : aclrtAllocatorAllocAdviseFunc *allocAdviseFunc,
150 : aclrtAllocatorGetAddrFromBlockFunc *getAddrFromBlockFunc)
151 : {
152 5 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
153 5 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
154 5 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
155 5 : const auto iter = g_AllocatorDesMap.find(stream);
156 5 : if (iter == g_AllocatorDesMap.end()) {
157 1 : std::string funcName = acl::AclErrorLogManager::GetFuncNameWithoutImplSuffix(__func__);
158 1 : acl::AclErrorLogManager::ReportInputError(acl::INVALID_PARAM_NO_VALUE_MSG,
159 2 : std::vector<const char *>({"func", "param", "reason"}),
160 2 : std::vector<const char *>({funcName.c_str(), "stream", "The stream is not registered with any allocator"}));
161 1 : return ACL_ERROR_INVALID_PARAM;
162 1 : }
163 4 : *allocatorDesc = iter->second.first;
164 4 : AllocatorDesc &desc = iter->second.second;
165 4 : if (allocator != nullptr) {
166 2 : *allocator = desc.obj;
167 : }
168 4 : if (allocFunc != nullptr) {
169 2 : *allocFunc = desc.allocFunc;
170 : }
171 4 : if (freeFunc != nullptr) {
172 2 : *freeFunc = desc.freeFunc;
173 : }
174 4 : if (allocAdviseFunc != nullptr) {
175 2 : *allocAdviseFunc = desc.allocAdviseFunc;
176 : }
177 4 : if (getAddrFromBlockFunc != nullptr) {
178 2 : *getAddrFromBlockFunc = desc.getAddrFromBlockFunc;
179 : }
180 4 : ACL_LOG_INFO("Get allocator By Stream success, stream %p.", stream);
181 4 : return ACL_SUCCESS;
182 5 : }
183 :
184 3 : aclError aclrtAllocatorUnregisterImpl(aclrtStream stream)
185 : {
186 3 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
187 2 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
188 2 : g_AllocatorDesMap.erase(stream);
189 2 : ACL_LOG_INFO("Unregister external allocator success, stream %p.", stream);
190 2 : return ACL_SUCCESS;
191 2 : }
192 : #ifdef __cplusplus
193 : }
194 : #endif
|
