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           1 :     AllocatorDesc(aclrtAllocator allocator,
      25             :                   aclrtAllocatorAllocFunc allocFunc,
      26             :                   aclrtAllocatorFreeFunc freeFunc,
      27             :                   aclrtAllocatorAllocAdviseFunc allocAdviseFunc,
      28             :                   aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc)
      29           1 :     {
      30           1 :         this->obj = allocator;
      31           1 :         this->allocFunc = allocFunc;
      32           1 :         this->freeFunc = freeFunc;
      33           1 :         this->allocAdviseFunc = allocAdviseFunc;
      34           1 :         this->getAddrFromBlockFunc = getAddrFromBlockFunc;
      35           1 :     }
      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           1 : aclrtAllocatorDesc aclrtAllocatorCreateDescImpl()
      48             : {
      49           2 :     ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorCreateDesc);
      50           1 :     ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
      51           1 :     ACL_LOG_INFO("Create allocator description.");
      52           1 :     AllocatorDesc *allocatorDesc = new(std::nothrow) AllocatorDesc;
      53           1 :     if (allocatorDesc == nullptr) {
      54           0 :         ACL_LOG_INNER_ERROR("alloc AllocatorDesc memory failed");
      55           0 :         return nullptr;
      56             :     }
      57           1 :     ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
      58           1 :     return static_cast<aclrtAllocatorDesc>(allocatorDesc);
      59             : }
      60             : 
      61           1 : aclError aclrtAllocatorDestroyDescImpl(aclrtAllocatorDesc allocatorDesc)
      62             : {
      63           2 :     ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorDestroyDesc);
      64           1 :     ACL_LOG_INFO("Destroy allocator description, allocatorDesc %p.", allocatorDesc);
      65           1 :     ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
      66           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
      67           1 :     delete static_cast<AllocatorDesc *>(allocatorDesc);
      68           1 :     allocatorDesc = nullptr;
      69           1 :     ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
      70           1 :     return ACL_SUCCESS;
      71             : }
      72             : 
      73           1 : aclError aclrtAllocatorSetObjToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocator allocator)
      74             : {
      75           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
      76           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocator);
      77           1 :     ACL_LOG_INFO("Set allocator to allocator description, allocatorDesc %p.", allocatorDesc);
      78           1 :     static_cast<AllocatorDesc *>(allocatorDesc)->obj = allocator;
      79           1 :     return ACL_SUCCESS;
      80             : }
      81             : 
      82           1 : aclError aclrtAllocatorSetAllocFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocFunc func)
      83             : {
      84           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
      85           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
      86           1 :     ACL_LOG_INFO("Set alloc function to allocator description, allocatorDesc %p.", allocatorDesc);
      87           1 :     static_cast<AllocatorDesc *>(allocatorDesc)->allocFunc = func;
      88           1 :     return ACL_SUCCESS;
      89             : }
      90             : 
      91           1 : aclError aclrtAllocatorSetFreeFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorFreeFunc func)
      92             : {
      93           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
      94           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
      95           1 :     ACL_LOG_INFO("Set free function to allocator description, allocatorDesc %p.", allocatorDesc);
      96           1 :     static_cast<AllocatorDesc *>(allocatorDesc)->freeFunc = func;
      97           1 :     return ACL_SUCCESS;
      98             : }
      99             : 
     100           1 : aclError aclrtAllocatorSetAllocAdviseFuncToDescImpl(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocAdviseFunc func)
     101             : {
     102           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
     103           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
     104           1 :     ACL_LOG_INFO("Set alloc advise function to allocator description, allocatorDesc %p.", allocatorDesc);
     105           1 :     static_cast<AllocatorDesc *>(allocatorDesc)->allocAdviseFunc = func;
     106           1 :     return ACL_SUCCESS;
     107             : }
     108             : 
     109           1 : aclError aclrtAllocatorSetGetAddrFromBlockFuncToDescImpl(aclrtAllocatorDesc allocatorDesc,
     110             :                                                      aclrtAllocatorGetAddrFromBlockFunc func)
     111             : {
     112           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
     113           1 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
     114           1 :     ACL_LOG_INFO("Set get_addr_from_block function to allocator description, allocatorDesc %p.", allocatorDesc);
     115           1 :     static_cast<AllocatorDesc *>(allocatorDesc)->getAddrFromBlockFunc = func;
     116           1 :     return ACL_SUCCESS;
     117             : }
     118             : 
     119           6 : aclError aclrtAllocatorRegisterImpl(aclrtStream stream, aclrtAllocatorDesc allocatorDesc)
     120             : {
     121             :     // stream must be not null when register external allocator
     122           6 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
     123           5 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
     124             : 
     125           5 :     AllocatorDesc *allocDesc = static_cast<AllocatorDesc *>(allocatorDesc);
     126           5 :     if (allocDesc->obj == nullptr) {
     127           1 :         ACL_LOG_INNER_ERROR("Should call aclrtAllocatorSetObjToDesc first.");
     128           1 :         return ACL_ERROR_INVALID_PARAM;
     129             :     }
     130           4 :     if (allocDesc->allocFunc == nullptr) {
     131           1 :         ACL_LOG_INNER_ERROR("Should call aclrtAllocatorSetAllocFuncToDesc first.");
     132           1 :         return ACL_ERROR_INVALID_PARAM;
     133             :     }
     134           3 :     if (allocDesc->freeFunc == nullptr) {
     135           1 :         ACL_LOG_INNER_ERROR("Should call aclrtAllocatorSetFreeFuncToDesc first.");
     136           1 :         return ACL_ERROR_INVALID_PARAM;
     137             :     }
     138           2 :     if (allocDesc->getAddrFromBlockFunc == nullptr) {
     139           1 :         ACL_LOG_INNER_ERROR("Should call aclrtAllocatorSetGetAddrFromBlockFuncToDesc first.");
     140           1 :         return ACL_ERROR_INVALID_PARAM;
     141             :     }
     142           1 :     AllocatorDesc allocDescCopy = AllocatorDesc(allocDesc->obj,
     143             :                                                 allocDesc->allocFunc,
     144             :                                                 allocDesc->freeFunc,
     145             :                                                 allocDesc->allocAdviseFunc,
     146           1 :                                                 allocDesc->getAddrFromBlockFunc);
     147           1 :     std::pair<aclrtAllocatorDesc, AllocatorDesc> allocatorDescPair(allocatorDesc, allocDescCopy);
     148           1 :     const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
     149           1 :     g_AllocatorDesMap[stream] = allocatorDescPair;
     150           1 :     ACL_LOG_INFO("Register external allocator success, stream %p, allocatorDesc %p.", stream, allocatorDesc);
     151           1 :     return ACL_SUCCESS;
     152             : }
     153             : 
     154           3 : aclError aclrtAllocatorGetByStreamImpl(aclrtStream stream,
     155             :                                    aclrtAllocatorDesc *allocatorDesc,
     156             :                                    aclrtAllocator *allocator,
     157             :                                    aclrtAllocatorAllocFunc *allocFunc,
     158             :                                    aclrtAllocatorFreeFunc *freeFunc,
     159             :                                    aclrtAllocatorAllocAdviseFunc *allocAdviseFunc,
     160             :                                    aclrtAllocatorGetAddrFromBlockFunc *getAddrFromBlockFunc)
     161             : {
     162           3 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
     163           6 :     const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
     164           3 :     const auto iter = g_AllocatorDesMap.find(stream);
     165           3 :     if (iter == g_AllocatorDesMap.end()) {
     166           0 :         return ACL_ERROR_INVALID_PARAM;
     167             :     }
     168           3 :     *allocatorDesc = iter->second.first;
     169           3 :     AllocatorDesc &desc = iter->second.second;
     170           3 :     if (allocator != nullptr) {
     171           2 :         *allocator = desc.obj;
     172             :     }
     173           3 :     if (allocFunc != nullptr) {
     174           2 :         *allocFunc = desc.allocFunc;
     175             :     }
     176           3 :     if (freeFunc != nullptr) {
     177           2 :         *freeFunc = desc.freeFunc;
     178             :     }
     179           3 :     if (allocAdviseFunc != nullptr) {
     180           2 :         *allocAdviseFunc = desc.allocAdviseFunc;
     181             :     }
     182           3 :     if (getAddrFromBlockFunc != nullptr) {
     183           2 :         *getAddrFromBlockFunc = desc.getAddrFromBlockFunc;
     184             :     }
     185           3 :     ACL_LOG_INFO("Get allocator By Stream success, stream %p.", stream);
     186           3 :     return ACL_SUCCESS;
     187             : }
     188             : 
     189           2 : aclError aclrtAllocatorUnregisterImpl(aclrtStream stream)
     190             : {
     191           2 :     ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
     192           1 :     const std::unique_lock<std::mutex> lk(g_AllocatorDescMutex);
     193           1 :     g_AllocatorDesMap.erase(stream);
     194           1 :     ACL_LOG_INFO("Unregister external allocator success, stream %p.", stream);
     195           1 :     return ACL_SUCCESS;
     196             : }