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