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 "engine_ctxs.h"
      12              : #include "log.h"
      13              : #include "hcomm_c_adpt.h"
      14              : #include "comm_engine_utils.h"
      15              : 
      16              : namespace hccl {
      17          165 : EngineCtxs::EngineCtxs() {}
      18              : 
      19          165 : EngineCtxs::~EngineCtxs() {}
      20              : 
      21            7 : HcclResult EngineCtxs::CreateCommEngineCtx(const std::string& tag, CommEngine engine, uint64_t size, void** ctx)
      22              : {
      23            7 :     std::lock_guard<std::mutex> lock(mutex_);
      24              :     // 阻止重复创建
      25            7 :     if (contextMap_.find(tag) != contextMap_.end()) {
      26            1 :         auto engineCtxMap = contextMap_[tag];
      27            1 :         CHK_PRT_RET(
      28              :             engineCtxMap.find(engine) != engineCtxMap.end(),
      29              :             HCCL_ERROR(
      30              :                 "[%s] already exist a context with same key, tag[%s], engine[%s]", __func__, tag.c_str(),
      31              :                 GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str()),
      32              :             HCCL_E_PARA);
      33            1 :     }
      34              : 
      35            7 :     CHK_RET(static_cast<HcclResult>(HcommEngineCtxCreate(engine, size, ctx)));
      36            7 :     contextMap_[tag][engine] = {HCCL_MEM_TYPE_NUM, *ctx, size}; // type不需要使用
      37            7 :     HCCL_INFO(
      38              :         "[%s]create context success, tag[%s], engine[%s]", __func__, tag.c_str(),
      39              :         GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
      40            7 :     return HCCL_SUCCESS;
      41            7 : }
      42              : 
      43            1 : HcclResult EngineCtxs::GetCommEngineCtx(const std::string& tag, CommEngine engine, void** ctx, uint64_t* size)
      44              : {
      45            1 :     std::lock_guard<std::mutex> lock(mutex_);
      46              :     // Ctx未创建返回
      47            1 :     const auto& tagIter = contextMap_.find(tag);
      48            1 :     if (tagIter == contextMap_.end()) {
      49            0 :         HCCL_INFO("[%s] not exist a context with tag[%s]", __func__, tag.c_str());
      50            0 :         return HCCL_E_NOT_FOUND;
      51              :     }
      52              : 
      53            1 :     const auto& engineCtxMap = tagIter->second;
      54            1 :     const auto& engineIter = engineCtxMap.find(engine);
      55            1 :     if (engineIter == engineCtxMap.end()) {
      56            0 :         HCCL_INFO(
      57              :             "[%s] not exist a context with tag[%s], engine[%s]", __func__, tag.c_str(),
      58              :             GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
      59            0 :         return HCCL_E_NOT_FOUND;
      60              :     }
      61              : 
      62            1 :     const auto& ctxRes = engineIter->second;
      63            1 :     *ctx = ctxRes.addr;
      64            1 :     *size = ctxRes.size;
      65            1 :     HCCL_INFO(
      66              :         "[%s] get context success, tag[%s], engine[%s]", __func__, tag.c_str(),
      67              :         GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
      68            1 :     return HCCL_SUCCESS;
      69            1 : }
      70              : 
      71            1 : HcclResult EngineCtxs::CopyCommEngineCtx(
      72              :     const std::string& tag, CommEngine engine, const void* srcCtx, uint64_t size, uint64_t dstCtxOffset)
      73              : {
      74              :     void* dstCtx;
      75            1 :     uint64_t dstSize = 0;
      76            1 :     CHK_RET(GetCommEngineCtx(tag, engine, &dstCtx, &dstSize));
      77            1 :     CHK_PRT_RET(
      78              :         dstCtxOffset + size > dstSize,
      79              :         HCCL_ERROR(
      80              :             "[%s]Copy engine ctx failed: buffer overflow detected. tag[%s], engine[%s], "
      81              :             "dstSize[%llu], dstCtxOffset[%llu], copySize[%llu]",
      82              :             __func__, tag.c_str(), GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str(), dstSize, dstCtxOffset,
      83              :             size),
      84              :         HCCL_E_PARA);
      85            0 :     CHK_RET(static_cast<HcclResult>(HcommEngineCtxCopy(
      86              :         engine, reinterpret_cast<uint8_t*>(dstCtx) + dstCtxOffset, srcCtx, size))); // 增加大小判断，增加强转
      87            0 :     HCCL_INFO(
      88              :         "[%s]copy engine ctx success, tag[%s], engine[%s]", __func__, tag.c_str(),
      89              :         GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
      90            0 :     return HCCL_SUCCESS;
      91              : }
      92              : 
      93            9 : HcclResult EngineCtxs::DestroyEngineCtx(const std::string& tag, CommEngine engine)
      94              : {
      95            9 :     std::lock_guard<std::mutex> lock(mutex_);
      96              :     // Ctx不存在返回错误
      97            9 :     if (contextMap_.find(tag) == contextMap_.end()) {
      98            1 :         HCCL_ERROR("[%s] not exist a context with tag[%s]", __func__, tag.c_str());
      99            1 :         return HCCL_E_PARA;
     100              :     }
     101            8 :     auto& engineCtxMap = contextMap_[tag];
     102            8 :     if (engineCtxMap.find(engine) == engineCtxMap.end()) {
     103            1 :         HCCL_ERROR(
     104              :             "[%s] not exist a context with tag[%s], engine[%s]", __func__, tag.c_str(),
     105              :             GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
     106            1 :         return HCCL_E_PARA;
     107              :     }
     108              :     // 获取内存信息
     109            7 :     HcclMem& memInfo = engineCtxMap[engine];
     110            7 :     CHK_RET(static_cast<HcclResult>(HcommEngineCtxDestroy(engine, memInfo.addr)));
     111              :     // 从映射中移除
     112            7 :     engineCtxMap.erase(engine);
     113            7 :     if (engineCtxMap.empty()) {
     114            6 :         contextMap_.erase(tag);
     115              :     }
     116              : 
     117            7 :     HCCL_INFO(
     118              :         "[%s]destroy context success, tag[%s], engine[%s]", __func__, tag.c_str(),
     119              :         GetEnumToString(GetCommEngineStatusStrMap(), engine).c_str());
     120            7 :     return HCCL_SUCCESS;
     121            9 : }
     122              : } // namespace hccl