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