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              : #ifndef AICPU_TS_CHANNEL_HELPER_H
      11              : #define AICPU_TS_CHANNEL_HELPER_H
      12              : 
      13              : #include <memory>
      14              : #include <mutex>
      15              : #include <vector>
      16              : #include "hccl/hccl_types.h"
      17              : #include "hcomm_res_defs.h"
      18              : #include "mem_device_pub.h"
      19              : 
      20              : static const uint32_t HCOMM_AICPU_CHANNEL_CTX_MAGIC_WORD = 0x0fcf0f4fU;
      21              : static const uint32_t HCOMM_AICPU_CHANNEL_CTX_VERSION = 0U;
      22              : 
      23              : typedef struct {
      24              :     CommAbiHeader abiHeader{HCOMM_AICPU_CHANNEL_CTX_VERSION, HCOMM_AICPU_CHANNEL_CTX_MAGIC_WORD, 0, 0};
      25              :     CommEngine engine{COMM_ENGINE_RESERVED};
      26              :     CommProtocol protocol{COMM_PROTOCOL_RESERVED};
      27              :     void *deviceChannel{nullptr};
      28              : } HcommAicpuChannelCtx;
      29              : 
      30              : /**
      31              :  * @brief AICPU 专用辅助类，持有 ctx 内存等 AICPU 专属资源，提供 AICPU 批量建链状态机接口。
      32              :  */
      33              : class AicpuTsChannelHelper {
      34              : public:
      35              :     // === per-channel 实例方法 ===
      36            2 :     HcclResult PreAllocCtx(ChannelHandle &outHandle)
      37              :     {
      38            2 :         ctxMem_ = std::make_shared<hccl::DeviceMem>(hccl::DeviceMem::alloc(sizeof(HcommAicpuChannelCtx)));
      39            2 :         CHK_PTR_NULL(ctxMem_->ptr());
      40            2 :         outHandle = reinterpret_cast<ChannelHandle>(ctxMem_->ptr());
      41            2 :         HCCL_INFO("[AicpuTsChannelHelper] pre-alloc ctx success, ctxPtr[%p]", ctxMem_->ptr());
      42            2 :         return HCCL_SUCCESS;
      43              :     }
      44              :     void SetCtxMem(std::shared_ptr<hccl::DeviceMem> mem)
      45              :     {
      46              :         ctxMem_ = std::move(mem);
      47              :     }
      48            2 :     void *GetCtxPtr() const
      49              :     {
      50            2 :         return ctxMem_ ? ctxMem_->ptr() : nullptr;
      51              :     }
      52              : 
      53              :     // === 批量静态方法 ===
      54              :     static HcclResult HandleStatus(const ChannelHandle *channelList, uint32_t listNum, CommEngine engine,
      55              :         const HcommChannelDesc *channelDescs, const std::vector<int32_t> &linkStatusList, int32_t *statusList);
      56              :     static HcclResult PreAllocChannels(ChannelHandle *targetChannels, ChannelHandle *userChannels,
      57              :         HcommChannelDesc *channelDescs, uint32_t channelNum);
      58              :     static HcclResult EnsureKernelBinLoaded(CommEngine engine);
      59           34 :     static aclrtBinHandle GetBinHandle()
      60              :     {
      61           34 :         return g_BinHandle;
      62              :     }
      63              : 
      64              :     static HcclResult TryFillCtxList(ChannelHandle *hostChannelHandles, uint32_t listNum,
      65              :         const hccl::DeviceMem &deviceChannelList, void *&outCtxList, bool &isCtxMode);
      66              : 
      67              : private:
      68              :     static HcclResult LaunchKernel(const ChannelHandle *channelList, uint32_t listNum, CommEngine engine,
      69              :         const HcommChannelDesc *channelDescs, aclrtBinHandle binHandle);
      70              : 
      71              :     std::shared_ptr<hccl::DeviceMem> ctxMem_;
      72              :     static aclrtBinHandle g_BinHandle;
      73              :     static std::mutex g_BinHandleMtx;
      74              : };
      75              : 
      76           59 : inline HcclResult UnwrapChannelHandle(ChannelHandle &handle)
      77              : {
      78           59 :     if (handle == 0) {
      79            8 :         HCCL_ERROR("[%s] handle is 0.", __func__);
      80            8 :         return HCCL_E_PTR;
      81              :     }
      82           51 :     const auto *ctx = reinterpret_cast<const HcommAicpuChannelCtx *>(static_cast<uintptr_t>(handle));
      83           51 :     if (ctx->abiHeader.version == HCOMM_AICPU_CHANNEL_CTX_VERSION
      84            2 :         && ctx->abiHeader.magicWord == HCOMM_AICPU_CHANNEL_CTX_MAGIC_WORD) {
      85            1 :         if (ctx->deviceChannel == nullptr) {
      86            0 :             HCCL_ERROR("[%s] ctx deviceChannel is null.", __func__);
      87            0 :             return HCCL_E_INTERNAL;
      88              :         }
      89            1 :         handle = reinterpret_cast<ChannelHandle>(ctx->deviceChannel);
      90              :     }
      91           51 :     return HCCL_SUCCESS;
      92              : }
      93              : 
      94              : #endif // AICPU_TS_CHANNEL_HELPER_H