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 4 : AllocatorDesc() = default;
23 : ~AllocatorDesc() = default;
24 2 : AllocatorDesc(
25 : aclrtAllocator allocator, aclrtAllocatorAllocFunc allocFunc, aclrtAllocatorFreeFunc freeFunc,
26 : aclrtAllocatorAllocAdviseFunc allocAdviseFunc, aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc)
27 2 : {
28 2 : this->obj = allocator;
29 2 : this->allocFunc = allocFunc;
30 2 : this->freeFunc = freeFunc;
31 2 : this->allocAdviseFunc = allocAdviseFunc;
32 2 : this->getAddrFromBlockFunc = getAddrFromBlockFunc;
33 2 : }
34 : aclrtAllocator obj = nullptr;
35 : aclrtAllocatorAllocFunc allocFunc;
36 : aclrtAllocatorFreeFunc freeFunc;
37 : aclrtAllocatorAllocAdviseFunc allocAdviseFunc;
38 : aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc;
39 : };
40 : std::mutex g_AllocatorDescMutex;
41 : // The first aclrtAllocatorDesc is created by the user, while the second AllocatorDesc is a saved copy.
42 : std::map<aclrtStream, std::pair<aclrtAllocatorDesc, AllocatorDesc>> g_AllocatorDesMap;
43 : } // namespace
44 :
45 : #ifdef __cplusplus
46 : extern "C" {
47 : #endif
48 :
49 2 : aclrtAllocatorDesc aclrtAllocatorCreateDescImpl()
50 : {
51 2 : ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorCreateDesc);
52 2 : ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
53 2 : ACL_LOG_INFO("Create allocator description.");
54 2 : AllocatorDesc* allocatorDesc = new (std::nothrow) AllocatorDesc;
55 2 : ACL_CHECK_MALLOC_RESULT_REPORT_RET(allocatorDesc, sizeof(AllocatorDesc), "new", nullptr);
56 2 : ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
57 2 : return static_cast<aclrtAllocatorDesc>(allocatorDesc);
58 2 : }
59 :
60 2 : aclError aclrtAllocatorDestroyDescImpl(aclrtAllocatorDesc allocatorDesc)
61 : {
62 2 : ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorDestroyDesc);
63 2 : ACL_LOG_INFO("Destroy allocator description, allocatorDesc %p.", allocatorDesc);
64 2 : ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
65 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
66 2 : delete static_cast<AllocatorDesc*>(allocatorDesc);
67 2 : allocatorDesc = nullptr;
68 2 : ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
69 2 : return ACL_SUCCESS;
70 2 : }
71 :
72 2 : aclError aclrtAllocatorSetObjToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocator allocator)
73 : {
74 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
75 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocator);
76 2 : ACL_LOG_INFO("Set allocator to allocator description, allocatorDesc %p.", allocatorDesc);
77 2 : static_cast<AllocatorDesc*>(allocatorDesc)->obj = allocator;
78 2 : return ACL_SUCCESS;
79 : }
80 :
81 2 : aclError aclrtAllocatorSetAllocFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocFunc func)
82 : {
83 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
84 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
85 2 : ACL_LOG_INFO("Set alloc function to allocator description, allocatorDesc %p.", allocatorDesc);
86 2 : static_cast<AllocatorDesc*>(allocatorDesc)->allocFunc = func;
87 2 : return ACL_SUCCESS;
88 : }
89 :
90 2 : aclError aclrtAllocatorSetFreeFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorFreeFunc func)
91 : {
92 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
93 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
94 2 : ACL_LOG_INFO("Set free function to allocator description, allocatorDesc %p.", allocatorDesc);
95 2 : static_cast<AllocatorDesc*>(allocatorDesc)->freeFunc = func;
96 2 : return ACL_SUCCESS;
97 : }
98 :
99 2 : aclError aclrtAllocatorSetAllocAdviseFuncToDescImpl(
100 : aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocAdviseFunc func)
101 : {
102 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
103 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
104 2 : ACL_LOG_INFO("Set alloc advise function to allocator description, allocatorDesc %p.", allocatorDesc);
105 2 : static_cast<AllocatorDesc*>(allocatorDesc)->allocAdviseFunc = func;
106 2 : return ACL_SUCCESS;
107 : }
108 :
109 2 : aclError aclrtAllocatorSetGetAddrFromBlockFuncToDescImpl(
110 : aclrtAllocatorDesc allocatorDesc, aclrtAllocatorGetAddrFromBlockFunc func)
111 : {
112 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
113 2 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
114 2 : ACL_LOG_INFO("Set get_addr_from_block function to allocator description, allocatorDesc %p.", allocatorDesc);
115 2 : static_cast<AllocatorDesc*>(allocatorDesc)->getAddrFromBlockFunc = func;
116 2 : return ACL_SUCCESS;
117 : }
118 :
119 7 : aclError aclrtAllocatorRegisterImpl(aclrtStream stream, aclrtAllocatorDesc allocatorDesc)
120 : {
121 : // stream must be not null when register external allocator
122 7 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
123 6 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
124 :
125 6 : AllocatorDesc* allocDesc = static_cast<AllocatorDesc*>(allocatorDesc);
126 10 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->obj, "allocatorDesc->obj");
127 9 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->allocFunc, "allocatorDesc->allocFunc");
128 8 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(allocDesc->freeFunc, "allocatorDesc->freeFunc");
129 7 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT_WITH_PRAM_NAME(
130 : allocDesc->getAddrFromBlockFunc, "allocatorDesc->getAddrFromBlockFunc");
131 :
132 : AllocatorDesc allocDescCopy = AllocatorDesc(
133 : allocDesc->obj, allocDesc->allocFunc, allocDesc->freeFunc, allocDesc->allocAdviseFunc,
134 2 : allocDesc->getAddrFromBlockFunc);
135 2 : std::pair<aclrtAllocatorDesc, AllocatorDesc> allocatorDescPair(allocatorDesc, allocDescCopy);
136 2 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
137 2 : g_AllocatorDesMap[stream] = allocatorDescPair;
138 2 : ACL_LOG_INFO("Register external allocator success, stream %p, allocatorDesc %p.", stream, allocatorDesc);
139 2 : return ACL_SUCCESS;
140 2 : }
141 :
142 5 : aclError aclrtAllocatorGetByStreamImpl(
143 : aclrtStream stream, aclrtAllocatorDesc* allocatorDesc, aclrtAllocator* allocator,
144 : aclrtAllocatorAllocFunc* allocFunc, aclrtAllocatorFreeFunc* freeFunc,
145 : aclrtAllocatorAllocAdviseFunc* allocAdviseFunc, aclrtAllocatorGetAddrFromBlockFunc* getAddrFromBlockFunc)
146 : {
147 5 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
148 5 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
149 5 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
150 5 : const auto iter = g_AllocatorDesMap.find(stream);
151 5 : if (iter == g_AllocatorDesMap.end()) {
152 1 : std::string funcName = acl::AclErrorLogManager::GetFuncNameWithoutImplSuffix(__func__);
153 1 : acl::AclErrorLogManager::ReportInputError(
154 2 : acl::INVALID_PARAM_NO_VALUE_MSG, std::vector<const char*>({"func", "param", "reason"}),
155 2 : std::vector<const char*>({funcName.c_str(), "stream", "The stream is not registered with any allocator"}));
156 1 : return ACL_ERROR_INVALID_PARAM;
157 1 : }
158 4 : *allocatorDesc = iter->second.first;
159 4 : AllocatorDesc& desc = iter->second.second;
160 4 : if (allocator != nullptr) {
161 2 : *allocator = desc.obj;
162 : }
163 4 : if (allocFunc != nullptr) {
164 2 : *allocFunc = desc.allocFunc;
165 : }
166 4 : if (freeFunc != nullptr) {
167 2 : *freeFunc = desc.freeFunc;
168 : }
169 4 : if (allocAdviseFunc != nullptr) {
170 2 : *allocAdviseFunc = desc.allocAdviseFunc;
171 : }
172 4 : if (getAddrFromBlockFunc != nullptr) {
173 2 : *getAddrFromBlockFunc = desc.getAddrFromBlockFunc;
174 : }
175 4 : ACL_LOG_INFO("Get allocator By Stream success, stream %p.", stream);
176 4 : return ACL_SUCCESS;
177 5 : }
178 :
179 3 : aclError aclrtAllocatorUnregisterImpl(aclrtStream stream)
180 : {
181 3 : ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
182 2 : const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
183 2 : g_AllocatorDesMap.erase(stream);
184 2 : ACL_LOG_INFO("Unregister external allocator success, stream %p.", stream);
185 2 : return ACL_SUCCESS;
186 2 : }
187 : #ifdef __cplusplus
188 : }
189 : #endif
|
