diff --git a/utils/inc/kernel/moe_distribute_base.h b/utils/inc/kernel/moe_distribute_base.h index b26bdc6..2883988 100644 --- a/utils/inc/kernel/moe_distribute_base.h +++ b/utils/inc/kernel/moe_distribute_base.h @@ -1,12 +1,12 @@ /** - * This program is free software, you can redistribute it and/or modify. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ /*! * \file moe_distribute_base.h @@ -16,350 +16,606 @@ #ifndef MOE_DISTRIBUTE_BASE_H #define MOE_DISTRIBUTE_BASE_H +#if __has_include("../common/op_kernel/mc2_kernel_utils.h") +#include "../common/op_kernel/mc2_kernel_utils.h" +#else +#include "../../common/op_kernel/mc2_kernel_utils.h" +#endif + +#if ASC_DEVKIT_MAJOR >= 9 +#include "basic_api/kernel_basic_intf.h" +#else #include "kernel_operator.h" +#endif +#include "adv_api/hccl/hccl.h" + +constexpr uint32_t MAX_RANK_NUM = 64U; // 最大卡数 +constexpr uint32_t MAX_OP_NUM = 8U; // MC2最大通信算子数 +constexpr uint32_t WRITE_SQE_SIZE = 64U; +constexpr uint32_t WRITE_WITH_NOTIFY_SQE_SIZE = 96U; +constexpr uint64_t WIN_STATE_OFFSET = 350U * 1024U; +constexpr uint64_t STATE_WIN_OFFSET = 950U * 1024U; +constexpr uint64_t WIN_PICI_OFFSET = 1024U * 1024U; +constexpr uint64_t PICI_WIN_SIZE = 512UL; +constexpr uint32_t NORMAL_CQE_SIZE = 64U; +constexpr uint32_t CQ_DEPTH_256 = + 256U; // 为cqeBuf申请256*32B空间，初始化HGM上的CQ空间时，如果cqDepth>256，则循环多次DataCopy +constexpr uint32_t UB_ALIGN = 32U; // UB按32字节对齐 +constexpr uint32_t WIN_SQPI_OFFSET = 0U; +constexpr uint32_t WIN_SQCI_OFFSET = 1U; +constexpr uint32_t WIN_CQPI_OFFSET = 2U; +constexpr uint32_t WIN_CQCI_OFFSET = 3U; +constexpr uint32_t WIN_SQPILINEAR_OFFSET = 4U; +constexpr uint32_t WIN_CQCILINEAR_OFFSET = 5U; +constexpr uint32_t WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET = 6U; +constexpr uint32_t UINT8_BITS_OFFSET = 8U; + +// WQ 32bit offset +constexpr uint32_t WQ_JETTYID_OFFSET = 0U; +constexpr uint32_t WQ_WQESIZE_OFFSET = 4U; +constexpr uint32_t WQ_SQDEPTH_OFFSET = 5U; +constexpr uint32_t WQ_TP_ID_OFFSET = 12U; +constexpr uint32_t WQ_RMTEID_0_3_OFFSET = 13U; +constexpr uint32_t WQ_RMTEID_4_7_OFFSET = 14U; +constexpr uint32_t WQ_RMTEID_8_11_OFFSET = 15U; +constexpr uint32_t WQ_RMTEID_12_15_OFFSET = 16U; +constexpr uint32_t WQ_RMTOBJID_OFFSET = 17U; +constexpr uint32_t WQ_RMTTOKENVALUE_OFFSET = 18U; +constexpr uint32_t WQ_LOCALTOKENID_OFFSET = 19U; +// WQ 64bit offset +constexpr uint32_t WQ_SQVA_OFFSET = 1U; +constexpr uint32_t WQ_DBADDR_OFFSET = 5U; + +// CQ 32bit offset +constexpr uint32_t CQ_JFCID_OFFSET = 0U; +constexpr uint32_t CQ_CQESIZE_OFFSET = 4U; +constexpr uint32_t CQ_CQDEPTH_OFFSET = 5U; +// CQ 64bit offset +constexpr uint32_t CQ_CQVA_OFFSET = 1U; +constexpr uint32_t CQ_DBADDR_OFFSET = 5U; + +// URMA protocol 8bit offset +constexpr uint32_t SQE_COMMON_UINT8_OFFSET_2 = 2U; // udf_flg:1 inline_en:1 cqe:1 se:1 fence:1 odr:3 +constexpr uint32_t SQE_COMMON_UINT8_OFFSET_3 = 3U; // owner:1 rmt_jetty_type:2 token_en:1 nf:1 rsv:3 +constexpr uint32_t SQE_COMMON_TARGET_HINT_OFFSET = 4U; +constexpr uint32_t SQE_COMMON_OPCODE_OFFSET = 5U; +constexpr uint32_t SQE_COMMON_SGE_NUM_OFFSET = 11U; +constexpr uint32_t CQE_STATUS_OFFSET = 3U; +constexpr uint32_t CQE_SUBSTATUS_OFFSET = 2U; +constexpr uint32_t CQE_ENTRY_IDX_HIGH_OFFSET = 5U; +constexpr uint32_t CQE_ENTRY_IDX_LOW_OFFSET = 4U; + +// URMA protocol 32bit offset +constexpr uint32_t SQE_COMMON_UINT32_OFFSET_2 = 2U; // sge_num:8 tp_id:24 +constexpr uint32_t SQE_COMMON_RMT_JETTY_OR_SEG_ID_OFFSET = 3U; +constexpr uint32_t SQE_COMMON_RMT_EID_31_0_OFFSET = 4U; +constexpr uint32_t SQE_COMMON_RMT_EID_63_32_OFFSET = 5U; +constexpr uint32_t SQE_COMMON_RMT_EID_95_64_OFFSET = 6U; +constexpr uint32_t SQE_COMMON_RMT_EID_127_96_OFFSET = 7U; +constexpr uint32_t SQE_COMMON_RMT_TOKEN_VALUE_OFFSET = 8U; +constexpr uint32_t SQE_UDF_OFFSET = 9U; +constexpr uint32_t SQE_WITH_NOTIFY_UDF_OFFSET = SQE_UDF_OFFSET; +constexpr uint32_t SQE_LENGTH_OFFSET = 12U; +constexpr uint32_t SQE_WITH_NOTIFY_LENGTH_OFFSET = 20U; +constexpr uint32_t SQE_TOKEN_ID_OFFSET = 13U; // rsv:12 token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_TOKEN_ID_OFFSET = 21U; // rsv:12 token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_TOKEN_ID_OFFSET = 12U; // rsv:12 notify_token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_TOKEN_VALUE_OFFSET = 13U; + +// URMA protocol 64bit offset +constexpr uint32_t SQE_RMT_ADDR_OFFSET = 5U; +constexpr uint32_t SQE_WITH_NOTIFY_RMT_ADDR_OFFSET = SQE_RMT_ADDR_OFFSET; +constexpr uint32_t SQE_DATA_ADDR_OFFSET = 7U; +constexpr uint32_t SQE_WITH_NOTIFY_DATA_ADDR_OFFSET = 11U; +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_ADDR_OFFSET = 7U; +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_DATA_OFFSET = 8U; -constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64; -constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19U; -constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2; -constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024; -constexpr uint32_t TIME_CYCLE = 50; // 系统cycle数转换成时间的基准单位，固定为50 - -struct HcclSignalInfo { - uint64_t resId; // 在代表event时为eventid，notify时为notifyid - uint64_t addr; - uint32_t devId; - uint32_t tsId; - uint32_t rankId; - uint32_t flag; +struct HcclAiRMAWQ { + uint32_t jettyId; + uint64_t sqVA; // SQE在HBM上起始地址 + uint32_t wqeSize; // 一个WQEBB占用内存大小（64B） + uint32_t sqDepth; // 可用的WQEBB个数 + uint64_t headAddr; // AIV无依赖 + uint64_t tailAddr; // AIV无依赖 + uint64_t dbAddr; // JFSDoorBell地址 + uint32_t tp_id; + uint8_t rmtEid[16]; + uint32_t rmtObjId; // rmtTokenID + uint32_t rmtTokenValue; + uint32_t localTokenId; }; -struct ListCommon { - uint64_t nextHost; - uint64_t preHost; - uint64_t nextDevice; - uint64_t preDevice; +struct HcclAiRMACQ { + uint32_t jfcId; + uint64_t cqVA; // CQE在HBM上起始地址 + uint32_t cqeSize; // 一个CQE占用内存大小（64B） + uint32_t cqDepth; // 可用的CQE个数 + uint64_t headAddr; + uint64_t tailAddr; + uint64_t dbAddr; // JFCDoorBell地址 }; -struct HcclStreamInfo { - int32_t streamIds; - uint32_t sqIds; - uint32_t cqIds; // 记录物理cqId - uint32_t logicCqids; // 记录逻辑cqId +struct HcclCombinOpParam { + uint64_t workSpace; // client和server之间通信的地址 + uint64_t workSpaceSize; // client和server之间通信的空间大小 + uint32_t rankId; // 当前卡rankId + uint32_t rankDim; // 总卡数 + uint64_t winSize; // ccu不使用 + uint64_t windowsIn[MAX_RANK_NUM]; // ccu不使用 + uint64_t windowsOut[MAX_RANK_NUM]; // ccu不使用 + + // for ccu + uint64_t xnAddr; // Xn寄存器其实地址 + uint64_t ckeAddr; // CKE寄存器其实地址 + uint64_t msAddr; // MS地址，预留 + uint64_t msSize; // 可写的MS个数，预留 + + uint32_t opType[MAX_OP_NUM]; + uint8_t algorithmType[MAX_OP_NUM]; + + HcclAiRMAWQ sqs[MAX_RANK_NUM]; + HcclAiRMACQ cqs[MAX_RANK_NUM]; }; -struct LocalResInfoV2 { - uint32_t streamNum; - uint32_t signalNum; - HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM]; - HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM]; - HcclStreamInfo mainStreamInfo; - HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM]; // 集合通信AICPU展开资源 - ListCommon nextTagRes; // HccltagLocalResV2 -}; +// URMA通信接口 +__aicore__ inline GM_ADDR GetURMASqInfoGMAddr(GM_ADDR hcclContext, uint32_t dstRankId) +{ + ascendc_assert((hcclContext != nullptr && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or dstRankId >= MAX_RANK_NUM"); -enum class rtFloatOverflowMode_t { - RT_OVERFLOW_MODE_SATURATION = 0, - RT_OVERFLOW_MODE_INFNAN, - RT_OVERFLOW_MODE_UNDEF, -}; + return (GM_ADDR)((__gm__ HcclAiRMAWQ *)(((__gm__ HcclCombinOpParam *)hcclContext)->sqs) + dstRankId); +} -struct AlgoTopoInfo { - uint32_t userRank; // 通信域 RankID - uint32_t userRankSize; // 通信域的Rank数量 - int32_t deviceLogicId; - bool isSingleMeshAggregation; - uint32_t deviceNumPerAggregation; // 每个Module中的Device数量 - uint32_t superPodNum; // 集群中总的超节点数 - uint32_t devicePhyId; - uint32_t topoType; // TopoType - uint32_t deviceType; - uint32_t serverNum; - uint32_t meshAggregationRankSize; - uint32_t multiModuleDiffDeviceNumMode; - uint32_t multiSuperPodDiffServerNumMode; - uint32_t realUserRank; - bool isDiffDeviceModule; - bool isDiffDeviceType; - uint32_t gcdDeviceNumPerAggregation; - uint32_t moduleNum; - uint32_t isUsedRdmaRankPairNum; - uint64_t isUsedRdmaRankPair; - uint32_t pairLinkCounterNum; - uint64_t pairLinkCounter; - uint32_t nicNum; - uint64_t nicList; // niclist数组指针 - uint64_t complanRankLength; // complanRank占用的字节数 - uint64_t complanRank; // 指针 - uint64_t bridgeRankNum; // bridgeRank占用的个数 - uint64_t bridgeRank; // 指针 - uint64_t serverAndsuperPodRankLength; // serverAndsuperPodRank占用的字节数 - uint64_t serverAndsuperPodRank; // 指针 -}; +__aicore__ inline GM_ADDR GetURMACqInfoGMAddr(GM_ADDR hcclContext, uint32_t dstRankId) +{ + ascendc_assert((hcclContext != nullptr && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or dstRankId >= MAX_RANK_NUM"); -struct HcclOpConfig { - uint8_t deterministic; //确定性计算开关 - uint8_t retryEnable; // 是否重执行 - uint8_t highPerfEnable; - uint8_t padding[5]; // 大小需要64By对齐，未来添加参数时减小padding - uint8_t linkTimeOut[8]; // 发送超时时长 - uint64_t notifyWaitTime; // 超时时长，同HCCL_EXEC_TIMEOUT - uint32_t retryHoldTime; - uint32_t retryIntervalTime; - bool interHccsDisable = false; //使能rdma开关 - rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF; - uint32_t multiQpThreshold = 512; // 多QP每个QP分担数据量最小阈值 -}; + return (GM_ADDR)((__gm__ HcclAiRMACQ *)(((__gm__ HcclCombinOpParam *)hcclContext)->cqs) + dstRankId); +} -struct HcclMC2WorkSpace { - uint64_t workSpace; - uint64_t workSpaceSize; -}; +__aicore__ inline void GetURMASqInfoTensor(const AscendC::LocalTensor &sqInfoTensor, GM_ADDR hcclContext, + uint32_t dstRankId) +{ + GM_ADDR sqInfoGMAddr = GetURMASqInfoGMAddr(hcclContext, dstRankId); + ascendc_assert((sqInfoGMAddr != nullptr), "GetURMASqInfoGMAddr failed"); + + AscendC::GlobalTensor sqInfoGlobalTensor; + sqInfoGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoGMAddr)); + AscendC::DataCopyExtParams sqInfoParams = {1U, static_cast(sizeof(HcclAiRMAWQ)), 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams sqInfoPadParams{false, 0U, 0U, 0U}; + AscendC::DataCopyPad(sqInfoTensor, sqInfoGlobalTensor, sqInfoParams, sqInfoPadParams); +} -struct RemoteResPtr { - uint64_t nextHostPtr; - uint64_t nextDevicePtr; -}; +__aicore__ inline void GetURMACqInfoTensor(const AscendC::LocalTensor &cqInfoTensor, GM_ADDR hcclContext, + uint32_t dstRankId) +{ + GM_ADDR cqInfoGMAddr = GetURMACqInfoGMAddr(hcclContext, dstRankId); + ascendc_assert((cqInfoGMAddr != nullptr), "GetURMACqInfoGMAddr failed"); + + AscendC::GlobalTensor cqInfoGlobalTensor; + cqInfoGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoGMAddr)); + AscendC::DataCopyExtParams cqInfoParams = {1U, static_cast(sizeof(HcclAiRMACQ)), 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams cqInfoPadParams{false, 0U, 0U, 0U}; + AscendC::DataCopyPad(cqInfoTensor, cqInfoGlobalTensor, cqInfoParams, cqInfoPadParams); +} -struct HDCommunicateParams { - uint64_t hostAddr { 0 }; - uint64_t deviceAddr { 0 }; - uint64_t readCacheAddr { 0 }; - uint32_t devMemSize{ 0 }; - uint32_t buffLen{ 0 }; - uint32_t flag{ 0 }; -}; +__aicore__ inline void GenerateCommWriteSQE(const AscendC::LocalTensor &sqeTensor, uint32_t sqeCount = 1) +{ + AscendC::Duplicate(sqeTensor, 0, WRITE_SQE_SIZE); // 初始化为全0 + AscendC::SyncFunc(); + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) = 0b001; // odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_3) = 0b10110000; // owner=1, rmt_jetty_type=0b01, token_en=1, nf=0 + sqeTensor(SQE_COMMON_OPCODE_OFFSET) = 0x3; // opcode=0x3 + sqeTensor(SQE_COMMON_SGE_NUM_OFFSET) = 1; // sge_num=1 +} -struct HcclRankRelationResV2 { - uint32_t remoteUsrRankId; - uint32_t remoteWorldRank; - uint64_t windowsIn; - uint64_t windowsOut; - uint64_t windowsExp; - ListCommon nextTagRes; -}; +__aicore__ inline void GenerateCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, + uint32_t sqeCount = 1) +{ + AscendC::Duplicate(sqeTensor, 0, WRITE_WITH_NOTIFY_SQE_SIZE); // 初始化为全0 + AscendC::SyncFunc(); + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) = 0b001; // odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_3) = 0b10110000; // owner=1, rmt_jetty_type=0b01, token_en=1, nf=0 + sqeTensor(SQE_COMMON_OPCODE_OFFSET) = 0x5; // opcode=0x5，使用write with notify SQE + sqeTensor(SQE_COMMON_SGE_NUM_OFFSET) = 1; // sge_num=1 +} -struct HcclOpResParam { - // 本地资源 - HcclMC2WorkSpace mc2WorkSpace; - uint32_t localUsrRankId; // usrrankid - uint32_t rankSize; // 通信域内total rank个数 - uint64_t winSize; // 每个win大小，静态图时，可能是0，如果通信域内也有动态图，则可能为非0 - uint64_t localWindowsIn; // 全F为无效值 - uint64_t localWindowsOut; // 全F为无效值 - char hcomId[128]; - // aicore识别remote window - uint64_t winExpSize; - uint64_t localWindowsExp; - uint32_t rWinStart; // 为HcclRankRelationRes起始位置 - uint32_t rWinOffset; // 为HcclRemoteRes的大小 - uint64_t version; - LocalResInfoV2 localRes; - AlgoTopoInfo topoInfo; - - // 外部配置参数 - HcclOpConfig config; - uint64_t hostStateInfo; - uint64_t aicpuStateInfo; - uint64_t lockAddr; - uint32_t rsv[16]; - uint32_t notifysize; // RDMA场景使用，910B/910_93为4B，其余芯片为8B - uint32_t remoteResNum; // 有效的remoteResNum - RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM]; //数组指针，指向HcclRankRelationResV2，下标为remoteUserRankId - - // communicate retry - HDCommunicateParams kfcControlTransferH2DParams; - HDCommunicateParams kfcStatusTransferD2HParams; - uint64_t tinyMem; // for all2all - uint64_t tinyMemSize; - // 零拷贝场景使用 - uint64_t zeroCopyHeadPtr; - uint64_t zeroCopyTailPtr; - uint64_t zeroCopyRingBuffer; - uint64_t zeroCopyIpcPtrs[16]; // 保存集合通信时每个对端的输入输出内存地址 - uint32_t zeroCopyDevicePhyId[16]; // 保存每个rank对应的物理卡Id - - bool utraceStatusFlag; -}; +__aicore__ inline void UpdateCommonSQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + // 更新tp_id(24b), jetty_id(20b), rmt_eid(128b), rmt_token_value(32b) + + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_COMMON_UINT32_OFFSET_2) &= 0xff000000; // 保留sge_num + templateSqeU32(SQE_COMMON_UINT32_OFFSET_2) |= (sqInfoU32(WQ_TP_ID_OFFSET) & 0x00ffffff); // tp_id + templateSqeU32(SQE_COMMON_RMT_JETTY_OR_SEG_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_jetty_or_seg_id + templateSqeU32(SQE_COMMON_RMT_EID_31_0_OFFSET) = sqInfoU32(WQ_RMTEID_0_3_OFFSET); // rmt_eid + templateSqeU32(SQE_COMMON_RMT_EID_63_32_OFFSET) = sqInfoU32(WQ_RMTEID_4_7_OFFSET); + templateSqeU32(SQE_COMMON_RMT_EID_95_64_OFFSET) = sqInfoU32(WQ_RMTEID_8_11_OFFSET); + templateSqeU32(SQE_COMMON_RMT_EID_127_96_OFFSET) = sqInfoU32(WQ_RMTEID_12_15_OFFSET); + templateSqeU32(SQE_COMMON_RMT_TOKEN_VALUE_OFFSET) = sqInfoU32(WQ_RMTTOKENVALUE_OFFSET); // rmt_token_value +} -// Transport 内存类型 -enum class HcclAiRMAMemType : uint32_t { - LOCAL_INPUT = 0, - REMOTE_INPUT, +__aicore__ inline void UpdateCommWriteSQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + UpdateCommonSQE(sqeTensor, sqInfoTensor); - LOCAL_OUTPUT, - REMOTE_OUTPUT, + // 更新rmt_token_id(20b) + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_token_id +} - // 可透传更多的内存，可在MAX_NUM之前追加，例如： - // LOCAL_EXP, - // REMOTE_EXP, - MAX_NUM -}; +__aicore__ inline void UpdateCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + UpdateCommonSQE(sqeTensor, sqInfoTensor); + + // 更新rmt_token_id(20b), notify_token_value(32b), notify_token_id(20b) + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_WITH_NOTIFY_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_token_id + templateSqeU32(SQE_WITH_NOTIFY_NOTIFY_TOKEN_VALUE_OFFSET) = + sqInfoU32(WQ_RMTTOKENVALUE_OFFSET); // notify_token_value + templateSqeU32(SQE_WITH_NOTIFY_NOTIFY_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // notify_token_id +} -// Transport 内存信息 -struct HcclAiRMAMemInfo { - uint32_t memMaxNum{0}; // 最大内存数量，等于 HcclAiRMAMemType::MAX_NUM - uint32_t sizeOfMemDetails{0}; // sizeof(MemDetails)，用于内存校验和偏移计算 - uint64_t memDetailPtr{0}; // MemDetails数组首地址, 个数: HcclAiRMAMemType::MAX_NUM - // 可往后追加字段 -}; +__aicore__ inline void SetCommWriteSQE(const AscendC::LocalTensor &sqeTensor, uint64_t dataAddr, + uint64_t rmtAddr, uint32_t length, uint8_t cqe) +{ + // 设置cqe(1 bit), odr(3 bits), data_addr(64 bits), rmt_addr(64 bits), length(32 bits) + + AscendC::LocalTensor sqeLocalU32 = sqeTensor.ReinterpretCast(); + AscendC::LocalTensor sqeLocalU64 = sqeTensor.ReinterpretCast(); + + sqeLocalU64(SQE_RMT_ADDR_OFFSET) = rmtAddr; + sqeLocalU64(SQE_DATA_ADDR_OFFSET) = dataAddr; + sqeLocalU32(SQE_LENGTH_OFFSET) = length; + if (cqe > 0U) { + // cqe=1, odr=0b010 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00100010; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11111010; + } else { + // cqe=0, odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11011001; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00000001; + } +} -// 全部 Transport QP/Mem 信息 -struct HcclAiRMAInfo { - uint32_t curRankId{0}; // 当前rankId - uint32_t rankNum{0}; // rank数量 - uint32_t qpNum{0}; // 单个Transport的QP数量 - - uint32_t sizeOfAiRMAWQ{0}; // sizeof(HcclAiRMAWQ) - uint32_t sizeOfAiRMACQ{0}; // sizeof(HcclAiRMACQ) - uint32_t sizeOfAiRMAMem{0}; // sizeof(HcclAiRMAMemInfo) - - // HcclAiRMAWQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取SQ指针：sqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ - // 0 <= qpIndex < qpNum - uint64_t sqPtr{0}; - - // HcclAiRMACQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取SCQ指针：scqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ - // 0 <= qpIndex < qpNum - uint64_t scqPtr{0}; - - // HcclAiRMAWQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取RQ指针：rqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ - // 0 <= qpIndex < qpNum - uint64_t rqPtr{0}; - - // HcclAiRMACQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取RCQ指针: rcqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ - // 0 <= qpIndex < qpNum - uint64_t rcqPtr{0}; - - // HcclAivMemInfo一维数组 - // 内存信息个数: rankNum - // 计算偏移获取内存信息指针: memPtr + rankId * sizeOfAiRMAMem - // srcRankId 获取自身内存信息，dstRankId 获取 Transport 内存信息 - uint64_t memPtr{0}; - // 可往后追加字段 -}; -struct CombinedCapability { - uint64_t dataplaneModeBitmap; -}; +__aicore__ inline void SetCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, uint64_t dataAddr, + uint64_t rmtAddr, uint32_t length, uint64_t notifyAddr, + uint64_t notifyData, uint8_t cqe) +{ + // 设置cqe(1 bit), odr(3 bits), data_addr(64 bits), rmt_addr(64 bits), length(32 bits), notify_addr(64 bits), + // notify_data(64 bits) + + AscendC::LocalTensor sqeLocalU32 = sqeTensor.ReinterpretCast(); + AscendC::LocalTensor sqeLocalU64 = sqeTensor.ReinterpretCast(); + + sqeLocalU64(SQE_WITH_NOTIFY_RMT_ADDR_OFFSET) = rmtAddr; + sqeLocalU64(SQE_WITH_NOTIFY_DATA_ADDR_OFFSET) = dataAddr; + sqeLocalU32(SQE_WITH_NOTIFY_LENGTH_OFFSET) = length; + sqeLocalU64(SQE_WITH_NOTIFY_NOTIFY_ADDR_OFFSET) = notifyAddr; // notify_addr + sqeLocalU64(SQE_WITH_NOTIFY_NOTIFY_DATA_OFFSET) = notifyData; // notify_data + + if (cqe > 0U) { + // cqe=1, odr=0b010 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00100010; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11111010; + } else { + // cqe=0, odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11011001; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00000001; + } +} -struct HcclA2CombineOpParam { - uint64_t workSpace; // Address for communication between client and server, - // hccl requests and clears - uint64_t workSpaceSize; // Space for communication between client and server - uint32_t rankId; // id of this rank - uint32_t rankNum; // num of ranks in this comm group - uint64_t winSize; // size of each windows memory - uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM]; // windows address for input, windowsIn[rankId] corresponds - // to the local card address, - // and others are cross-card mapping addresses. - uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM]; // windows address for output, windowsOut[rankId] corresponds - // to the local card address, - // and others are cross-card mapping addresses. - uint8_t res[8328]; - uint8_t multiFlag; - __gm__ AscendC::IbVerbsData *data; - uint64_t dataSize; - // 追加字段 - uint64_t sizeOfAiRMAInfo; // sizeof(HcclAiRMAInfo) - uint64_t aiRMAInfo; // HcclAiRMAInfo* 单个结构体指针 - - CombinedCapability* capability; // address of the communication capability information structure on the Device - uint64_t capabilitySize; // size of the communication capability information structure -}; -enum class DataplaneMode : uint32_t { - HOST = 0, - AICPU = 1, - AIV = 2, -}; +__aicore__ inline void SendJFSDoorBell(const AscendC::LocalTensor &jfsDoorBellTensor, + const AscendC::LocalTensor &sqInfoTensor, uint32_t sqPi) +{ + // 敲JFS DoorBell，更新硬件的sqPi + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + st_dev(sqPi, (__gm__ uint32_t *)(sqInfoU64(WQ_DBADDR_OFFSET)), 0); // Scalar操作 +} -enum class DBMode : int32_t { - INVALID_DB = -1, - HW_DB = 0, - SW_DB -}; +__aicore__ inline void SendJFCDoorBell(const AscendC::LocalTensor &jfcDoorBellTensor, + const AscendC::LocalTensor &cqInfoTensor, uint32_t cqCi) +{ + // 敲JFC DoorBell，更新硬件的cqCi + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU64 = cqInfoTensor.ReinterpretCast(); -struct HcclAiRMAWQ { - uint32_t wqn{0}; - uint64_t bufAddr{0}; - uint32_t wqeSize{0}; - uint32_t depth{0}; - uint64_t headAddr{0}; - uint64_t tailAddr{0}; - DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw - uint64_t dbAddr{0}; - uint32_t sl{0}; -}; + uint64_t jfcDbValue = (static_cast(cqInfoU32(CQ_JFCID_OFFSET)) << 32) + static_cast(cqCi); + st_dev(jfcDbValue, (__gm__ uint64_t *)(cqInfoU64(CQ_DBADDR_OFFSET)), 0); // Scalar操作 +} -struct HcclAiRMACQ { - uint32_t cqn{0}; - uint64_t bufAddr{0}; - uint32_t cqeSize{0}; - uint32_t depth{0}; - uint64_t headAddr{0}; - uint64_t tailAddr{0}; - DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw - uint64_t dbAddr{0}; -}; +__aicore__ inline bool CheckCqeStatus(const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &cqeTensor, uint32_t &outCqCi, + uint32_t &outCqCiLinear, uint32_t &newestCompletedPi) +{ + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU64 = cqInfoTensor.ReinterpretCast(); + + uint32_t cqeSize = cqInfoU32(CQ_CQESIZE_OFFSET); + AscendC::GlobalTensor cqGlobalTensor; + cqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoU64(CQ_CQVA_OFFSET))); + AscendC::DataCopyExtParams cqeParams = {1U, 8U, 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams cqePadParams{false, 0U, 0U, 0U}; + + AscendC::DataCopyPad(cqeTensor, cqGlobalTensor[cqeSize * outCqCi], cqeParams, cqePadParams); + AscendC::SyncFunc(); + if (cqeTensor(CQE_STATUS_OFFSET) == 0xff) { + return true; + } else if (cqeTensor(CQE_STATUS_OFFSET) != 0) { + // 退出kernel并报错 + ascendc_assert(false, "CQE status is abnormal! status is %d, substatus is %d.\n", cqeTensor(CQE_STATUS_OFFSET), + cqeTensor(CQE_SUBSTATUS_OFFSET)); + } -struct hns_roce_rc_sq_wqe { - uint32_t byte_4; - uint32_t msg_len; - uint32_t immtdata; - uint32_t byte_16; - uint32_t byte_20; - uint32_t rkey; - uint64_t remoteVA; -}; + // status==0，处理当前CQE，从entry_idx获取对应WQE的sqPi + newestCompletedPi = (static_cast(cqeTensor(CQE_ENTRY_IDX_HIGH_OFFSET)) << UINT8_BITS_OFFSET) + + static_cast(cqeTensor(CQE_ENTRY_IDX_LOW_OFFSET)); + // 把CQE的status设置为无效值，并写回CQ + cqeTensor(CQE_STATUS_OFFSET) = 0xff; + AscendC::SyncFunc(); + AscendC::DataCopyPad(cqGlobalTensor[cqeSize * outCqCi], cqeTensor, cqeParams); + AscendC::SyncFunc(); // 防止cqeTensor被下一轮加载覆盖 + return false; +} -struct hns_roce_lite_wqe_data_seg { - uint32_t len; - uint32_t lkey; - uint64_t localVA; -}; +__aicore__ inline void PollCommCQUpdateSQCI(const AscendC::LocalTensor &sqInfoTensor, + const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t &outSqCi, + uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + uint32_t cqDepth = cqInfoU32(CQ_CQDEPTH_OFFSET); + + uint32_t pollTimes = 0; + uint32_t newestCompletedPi = 0; + while (pollTimes < cqDepth - 1) { + if (CheckCqeStatus(cqInfoTensor, cqeTensor, outCqCi, outCqCiLinear, newestCompletedPi)) { + break; + } + + // 递增本地的outCqCi + outCqCi = (outCqCi + 1) % cqDepth; + ++outCqCiLinear; + ++pollTimes; + } + + // 更新本地的outSqCi + if (pollTimes > 0) { + outSqCi = (newestCompletedPi + 1) % sqDepth; + } + AscendC::SyncFunc(); // 等cqe status写回无效值 + // 通过敲JFC DoorBell，更新硬件的cqCi + SendJFCDoorBell(jfcDoorBellTensor, cqInfoTensor, outCqCiLinear); +} + +__aicore__ inline void InvalidateCqeStatus(const AscendC::LocalTensor &cqeInfoTensor, + const AscendC::LocalTensor &cqeTensor) +{ + // 为cqeBuf_申请256*32B空间，初始化HGM上的CQ空间时，如果cqDepth>256，则循环多次DataCopy + AscendC::LocalTensor cqInfoU64 = cqeInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU32 = cqeInfoTensor.ReinterpretCast(); + + AscendC::GlobalTensor cqGlobalTensor; + cqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoU64(CQ_CQVA_OFFSET))); + uint32_t cqDepth = cqInfoU32(CQ_CQDEPTH_OFFSET); + uint32_t cqeNum = (cqDepth > CQ_DEPTH_256) ? CQ_DEPTH_256 : cqDepth; + + AscendC::Duplicate(cqeTensor, 0xff, UB_ALIGN * cqeNum); // 初始化为全1 + AscendC::SyncFunc(); + + AscendC::DataCopyExtParams dataCopyParams = {static_cast(cqeNum), UB_ALIGN, 0U, + NORMAL_CQE_SIZE - UB_ALIGN, 0U}; + uint32_t loopTimes = AscendC::Ceil(cqDepth, CQ_DEPTH_256); + if (loopTimes > 1) { + for (uint32_t i = 0; i < loopTimes - 1; ++i) { + AscendC::DataCopyPad(cqGlobalTensor[NORMAL_CQE_SIZE * CQ_DEPTH_256 * i], cqeTensor, dataCopyParams); + } + uint32_t tailBlockCqeNum = cqDepth - CQ_DEPTH_256 * (loopTimes - 1); + if (tailBlockCqeNum != CQ_DEPTH_256) { + dataCopyParams = {static_cast(tailBlockCqeNum), UB_ALIGN, 0U, NORMAL_CQE_SIZE - UB_ALIGN, 0U}; + } + } + AscendC::DataCopyPad(cqGlobalTensor[NORMAL_CQE_SIZE * CQ_DEPTH_256 * (loopTimes - 1)], cqeTensor, dataCopyParams); +} + +__aicore__ inline uint32_t GetAvailableSpace(const AscendC::LocalTensor &sqInfoTensor, uint32_t &outSqPi, + uint32_t &outSqCi) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + + // 根据本地的outSqPi outSqCi 计算SQ可用空间大小（SQ不会用满，会预留一个WQE） + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + + if (outSqPi == outSqCi) { + return (sqDepth - 1); + } else if ((outSqPi + 1) % sqDepth == outSqCi) { + return 0; + } else if (outSqPi > outSqCi) { + return (sqDepth - (outSqPi - outSqCi) - 1); + } + return (outSqCi - outSqPi - 1); +} + +__aicore__ inline void +PutCommSQE(const AscendC::LocalTensor &sqInfoTensor, const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &sqeTensor, const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t sqeCount, uint32_t &outSqPi, + uint32_t &outSqPiLinear, uint32_t &outSqCi, uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + uint32_t sqeSize = sqInfoU32(WQ_WQESIZE_OFFSET); // 单个WQEBB占用大小（64B），由HCCL传递 + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + ascendc_assert(sqeCount < sqDepth, "too many SQE! SQE num[%d] should less than sqDepth[%d].", sqeCount, sqDepth); + + // 计算SQ可用空间大小 + uint32_t availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + while (availableSpace < sqeCount) { + // 可用空间不足时，轮询CQ，更新本地的cqCi_ sqCi_和硬件的cqCi + PollCommCQUpdateSQCI(sqInfoTensor, cqInfoTensor, cqeTensor, jfcDoorBellTensor, outSqCi, outCqCi, outCqCiLinear); + availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + } + + // 可用空间足够时，把WQE拷贝到HBM上的SQ中 + AscendC::GlobalTensor sqGlobalTensor; + sqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoU64(WQ_SQVA_OFFSET))); -__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) { - __gm__ uint8_t* start = - (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); - __gm__ uint8_t* end = - (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); - AscendC::GlobalTensor global; - global.SetGlobalBuffer(start); - for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) { - AscendC::DataCacheCleanAndInvalid(global[i]); + if (likely(sqDepth - outSqPi >= sqeCount)) { + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, WRITE_SQE_SIZE * sqeCount); + } else { + uint32_t firstPartSize = sqDepth - outSqPi; + uint32_t secondPartSize = sqeCount - firstPartSize; + + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, WRITE_SQE_SIZE * firstPartSize); + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_SQE_SIZE * firstPartSize], WRITE_SQE_SIZE * secondPartSize); } + + // 更新本地的outSqPi + outSqPiLinear += sqeCount; + outSqPi = outSqPiLinear % sqDepth; } -__aicore__ inline DataplaneMode GetDataplaneMode(GM_ADDR contextGM0) { - __gm__ HcclA2CombineOpParam *winContext_ = (__gm__ HcclA2CombineOpParam *)contextGM0; - CombinedCapability* capability = winContext_->capability; - uint64_t capabilitySize = winContext_->capabilitySize; - DataplaneMode dataplaneMode = DataplaneMode::AICPU; - if (capability == 0) { - return dataplaneMode; + +__aicore__ inline void +PutCommNotifySQE(const AscendC::LocalTensor &sqInfoTensor, const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &sqeTensor, const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t sqeCount, uint32_t &outSqPi, + uint32_t &outSqPiLinear, uint32_t &outSqCi, uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + uint32_t sqeSize = sqInfoU32(WQ_WQESIZE_OFFSET); // 单个WQEBB占用大小（64B），由HCCL传递 + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + ascendc_assert((sqeCount << 1) < sqDepth, "too many SQE! SQE num[%d] should less than sqDepth[%d].", + (sqeCount << 1), sqDepth); + + // 计算SQ可用空间大小 + uint32_t availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + while (availableSpace < (sqeCount << 1)) { + // 可用空间不足时，轮询CQ，更新本地的cqCi_ sqCi_和硬件的cqCi + PollCommCQUpdateSQCI(sqInfoTensor, cqInfoTensor, cqeTensor, jfcDoorBellTensor, outSqCi, outCqCi, outCqCiLinear); + availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); } - uint64_t dataplaneModeBitmap = capability->dataplaneModeBitmap; - if ((dataplaneModeBitmap & 0x04) == 0x04) { - dataplaneMode = DataplaneMode::AIV; + + // 可用空间足够时，把WQE拷贝到HBM上的SQ中 + AscendC::GlobalTensor sqGlobalTensor; + sqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoU64(WQ_SQVA_OFFSET))); + + if (likely((sqDepth - outSqPi) >= (sqeCount << 1))) { + AscendC::DataCopyParams intriParams{static_cast(sqeCount), WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, intriParams); + } else { + uint32_t firstPartSize = (sqDepth - outSqPi) >> 1; + uint32_t secondPartSize = sqeCount - firstPartSize; + AscendC::DataCopyParams firstPartParams{static_cast(firstPartSize), WRITE_WITH_NOTIFY_SQE_SIZE >> 5, + 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, firstPartParams); + + if (((sqDepth - outSqPi) & 1) == 0) { + AscendC::DataCopyParams secondPartParams{static_cast(secondPartSize), + WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize], secondPartParams); + } else { + // 如果剩余奇数个WQEBB，其中一个WQE会被切块分在尾和头 + AscendC::DataCopy(sqGlobalTensor[sqeSize * (sqDepth - 1)], + sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize], sqeSize); + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize + sqeSize], + WRITE_WITH_NOTIFY_SQE_SIZE - sqeSize); + + AscendC::DataCopyParams secondPartParams{static_cast(secondPartSize - 1), + WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize], sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * (firstPartSize + 1)], + secondPartParams); + } } - return dataplaneMode; + + // 更新本地的outSqPi + outSqPiLinear += (sqeCount << 1); + outSqPi = outSqPiLinear % sqDepth; } -__aicore__ inline int64_t GetCurrentTimestampUs() +__aicore__ inline void GetPICI(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, uint32_t &outSqPi, + uint32_t &outSqCi, uint32_t &outCqPi, uint32_t &outCqCi, uint32_t &outSqPiLinear, + uint32_t &outCqCiLinear) { - return AscendC::GetSystemCycle() / TIME_CYCLE; + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); + + outSqPi = piCiGlobalTensor(WIN_SQPI_OFFSET); + outSqCi = piCiGlobalTensor(WIN_SQCI_OFFSET); + outCqPi = piCiGlobalTensor(WIN_CQPI_OFFSET); + outCqCi = piCiGlobalTensor(WIN_CQCI_OFFSET); + outSqPiLinear = piCiGlobalTensor(WIN_SQPILINEAR_OFFSET); + outCqCiLinear = piCiGlobalTensor(WIN_CQCILINEAR_OFFSET); } -__aicore__ inline void RecordRankCommDuration(AscendC::LocalTensor performanceInfoU32Tensor, uint32_t rankId, int64_t startTime) +__aicore__ inline void GetIsFirstInComm(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, bool &outIsFirst) { - int64_t endTime = GetCurrentTimestampUs(); - int32_t duration = static_cast(endTime - startTime); // int32_t可以表示2^31(us)，约35min在实际场景下满足需要 - performanceInfoU32Tensor.SetValue(rankId * sizeof(int64_t) / sizeof(int32_t), duration); // 使用int32_t是因为atomicAdd不支持int64_t类型，这里只赋值到int64_t的低32位。 + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); + + outIsFirst = (piCiGlobalTensor(WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET) > 0); +} + +__aicore__ inline void UpdatePICI(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, uint32_t sqPi, + uint32_t sqCi, uint32_t cqPi, uint32_t cqCi, uint32_t sqPiLinear, uint32_t cqCiLinear) +{ + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + + piCiGlobalTensor(WIN_SQPI_OFFSET) = sqPi; + piCiGlobalTensor(WIN_SQCI_OFFSET) = sqCi; + piCiGlobalTensor(WIN_CQPI_OFFSET) = cqPi; + piCiGlobalTensor(WIN_CQCI_OFFSET) = cqCi; + piCiGlobalTensor(WIN_SQPILINEAR_OFFSET) = sqPiLinear; + piCiGlobalTensor(WIN_CQCILINEAR_OFFSET) = cqCiLinear; + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); } -#endif // MOE_DISTRIBUTE_BASE_H \ No newline at end of file + +__aicore__ inline void UpdateIsFirstInComm(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, bool isFirst) +{ + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + + piCiGlobalTensor(WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET) = isFirst ? 1 : 0; + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); +} + +#endif // MOE_DISTRIBUTE_BASE_H diff --git a/xllm_ops/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h b/xllm_ops/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h index 785eb31..2883988 100644 --- a/xllm_ops/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h +++ b/xllm_ops/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h @@ -1,12 +1,12 @@ /** - * This program is free software, you can redistribute it and/or modify. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ /*! * \file moe_distribute_base.h @@ -16,350 +16,606 @@ #ifndef MOE_DISTRIBUTE_BASE_H #define MOE_DISTRIBUTE_BASE_H +#if __has_include("../common/op_kernel/mc2_kernel_utils.h") +#include "../common/op_kernel/mc2_kernel_utils.h" +#else +#include "../../common/op_kernel/mc2_kernel_utils.h" +#endif + +#if ASC_DEVKIT_MAJOR >= 9 +#include "basic_api/kernel_basic_intf.h" +#else #include "kernel_operator.h" +#endif +#include "adv_api/hccl/hccl.h" + +constexpr uint32_t MAX_RANK_NUM = 64U; // 最大卡数 +constexpr uint32_t MAX_OP_NUM = 8U; // MC2最大通信算子数 +constexpr uint32_t WRITE_SQE_SIZE = 64U; +constexpr uint32_t WRITE_WITH_NOTIFY_SQE_SIZE = 96U; +constexpr uint64_t WIN_STATE_OFFSET = 350U * 1024U; +constexpr uint64_t STATE_WIN_OFFSET = 950U * 1024U; +constexpr uint64_t WIN_PICI_OFFSET = 1024U * 1024U; +constexpr uint64_t PICI_WIN_SIZE = 512UL; +constexpr uint32_t NORMAL_CQE_SIZE = 64U; +constexpr uint32_t CQ_DEPTH_256 = + 256U; // 为cqeBuf申请256*32B空间，初始化HGM上的CQ空间时，如果cqDepth>256，则循环多次DataCopy +constexpr uint32_t UB_ALIGN = 32U; // UB按32字节对齐 +constexpr uint32_t WIN_SQPI_OFFSET = 0U; +constexpr uint32_t WIN_SQCI_OFFSET = 1U; +constexpr uint32_t WIN_CQPI_OFFSET = 2U; +constexpr uint32_t WIN_CQCI_OFFSET = 3U; +constexpr uint32_t WIN_SQPILINEAR_OFFSET = 4U; +constexpr uint32_t WIN_CQCILINEAR_OFFSET = 5U; +constexpr uint32_t WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET = 6U; +constexpr uint32_t UINT8_BITS_OFFSET = 8U; + +// WQ 32bit offset +constexpr uint32_t WQ_JETTYID_OFFSET = 0U; +constexpr uint32_t WQ_WQESIZE_OFFSET = 4U; +constexpr uint32_t WQ_SQDEPTH_OFFSET = 5U; +constexpr uint32_t WQ_TP_ID_OFFSET = 12U; +constexpr uint32_t WQ_RMTEID_0_3_OFFSET = 13U; +constexpr uint32_t WQ_RMTEID_4_7_OFFSET = 14U; +constexpr uint32_t WQ_RMTEID_8_11_OFFSET = 15U; +constexpr uint32_t WQ_RMTEID_12_15_OFFSET = 16U; +constexpr uint32_t WQ_RMTOBJID_OFFSET = 17U; +constexpr uint32_t WQ_RMTTOKENVALUE_OFFSET = 18U; +constexpr uint32_t WQ_LOCALTOKENID_OFFSET = 19U; +// WQ 64bit offset +constexpr uint32_t WQ_SQVA_OFFSET = 1U; +constexpr uint32_t WQ_DBADDR_OFFSET = 5U; + +// CQ 32bit offset +constexpr uint32_t CQ_JFCID_OFFSET = 0U; +constexpr uint32_t CQ_CQESIZE_OFFSET = 4U; +constexpr uint32_t CQ_CQDEPTH_OFFSET = 5U; +// CQ 64bit offset +constexpr uint32_t CQ_CQVA_OFFSET = 1U; +constexpr uint32_t CQ_DBADDR_OFFSET = 5U; + +// URMA protocol 8bit offset +constexpr uint32_t SQE_COMMON_UINT8_OFFSET_2 = 2U; // udf_flg:1 inline_en:1 cqe:1 se:1 fence:1 odr:3 +constexpr uint32_t SQE_COMMON_UINT8_OFFSET_3 = 3U; // owner:1 rmt_jetty_type:2 token_en:1 nf:1 rsv:3 +constexpr uint32_t SQE_COMMON_TARGET_HINT_OFFSET = 4U; +constexpr uint32_t SQE_COMMON_OPCODE_OFFSET = 5U; +constexpr uint32_t SQE_COMMON_SGE_NUM_OFFSET = 11U; +constexpr uint32_t CQE_STATUS_OFFSET = 3U; +constexpr uint32_t CQE_SUBSTATUS_OFFSET = 2U; +constexpr uint32_t CQE_ENTRY_IDX_HIGH_OFFSET = 5U; +constexpr uint32_t CQE_ENTRY_IDX_LOW_OFFSET = 4U; + +// URMA protocol 32bit offset +constexpr uint32_t SQE_COMMON_UINT32_OFFSET_2 = 2U; // sge_num:8 tp_id:24 +constexpr uint32_t SQE_COMMON_RMT_JETTY_OR_SEG_ID_OFFSET = 3U; +constexpr uint32_t SQE_COMMON_RMT_EID_31_0_OFFSET = 4U; +constexpr uint32_t SQE_COMMON_RMT_EID_63_32_OFFSET = 5U; +constexpr uint32_t SQE_COMMON_RMT_EID_95_64_OFFSET = 6U; +constexpr uint32_t SQE_COMMON_RMT_EID_127_96_OFFSET = 7U; +constexpr uint32_t SQE_COMMON_RMT_TOKEN_VALUE_OFFSET = 8U; +constexpr uint32_t SQE_UDF_OFFSET = 9U; +constexpr uint32_t SQE_WITH_NOTIFY_UDF_OFFSET = SQE_UDF_OFFSET; +constexpr uint32_t SQE_LENGTH_OFFSET = 12U; +constexpr uint32_t SQE_WITH_NOTIFY_LENGTH_OFFSET = 20U; +constexpr uint32_t SQE_TOKEN_ID_OFFSET = 13U; // rsv:12 token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_TOKEN_ID_OFFSET = 21U; // rsv:12 token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_TOKEN_ID_OFFSET = 12U; // rsv:12 notify_token_id:20 +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_TOKEN_VALUE_OFFSET = 13U; + +// URMA protocol 64bit offset +constexpr uint32_t SQE_RMT_ADDR_OFFSET = 5U; +constexpr uint32_t SQE_WITH_NOTIFY_RMT_ADDR_OFFSET = SQE_RMT_ADDR_OFFSET; +constexpr uint32_t SQE_DATA_ADDR_OFFSET = 7U; +constexpr uint32_t SQE_WITH_NOTIFY_DATA_ADDR_OFFSET = 11U; +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_ADDR_OFFSET = 7U; +constexpr uint32_t SQE_WITH_NOTIFY_NOTIFY_DATA_OFFSET = 8U; -constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64; -constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19U; -constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2; -constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024; -constexpr uint32_t TIME_CYCLE = 50; // 系统cycle数转换成时间的基准单位，固定为50 - -struct HcclSignalInfo { - uint64_t resId; // 在代表event时为eventid，notify时为notifyid - uint64_t addr; - uint32_t devId; - uint32_t tsId; - uint32_t rankId; - uint32_t flag; +struct HcclAiRMAWQ { + uint32_t jettyId; + uint64_t sqVA; // SQE在HBM上起始地址 + uint32_t wqeSize; // 一个WQEBB占用内存大小（64B） + uint32_t sqDepth; // 可用的WQEBB个数 + uint64_t headAddr; // AIV无依赖 + uint64_t tailAddr; // AIV无依赖 + uint64_t dbAddr; // JFSDoorBell地址 + uint32_t tp_id; + uint8_t rmtEid[16]; + uint32_t rmtObjId; // rmtTokenID + uint32_t rmtTokenValue; + uint32_t localTokenId; }; -struct ListCommon { - uint64_t nextHost; - uint64_t preHost; - uint64_t nextDevice; - uint64_t preDevice; +struct HcclAiRMACQ { + uint32_t jfcId; + uint64_t cqVA; // CQE在HBM上起始地址 + uint32_t cqeSize; // 一个CQE占用内存大小（64B） + uint32_t cqDepth; // 可用的CQE个数 + uint64_t headAddr; + uint64_t tailAddr; + uint64_t dbAddr; // JFCDoorBell地址 }; -struct HcclStreamInfo { - int32_t streamIds; - uint32_t sqIds; - uint32_t cqIds; // 记录物理cqId - uint32_t logicCqids; // 记录逻辑cqId +struct HcclCombinOpParam { + uint64_t workSpace; // client和server之间通信的地址 + uint64_t workSpaceSize; // client和server之间通信的空间大小 + uint32_t rankId; // 当前卡rankId + uint32_t rankDim; // 总卡数 + uint64_t winSize; // ccu不使用 + uint64_t windowsIn[MAX_RANK_NUM]; // ccu不使用 + uint64_t windowsOut[MAX_RANK_NUM]; // ccu不使用 + + // for ccu + uint64_t xnAddr; // Xn寄存器其实地址 + uint64_t ckeAddr; // CKE寄存器其实地址 + uint64_t msAddr; // MS地址，预留 + uint64_t msSize; // 可写的MS个数，预留 + + uint32_t opType[MAX_OP_NUM]; + uint8_t algorithmType[MAX_OP_NUM]; + + HcclAiRMAWQ sqs[MAX_RANK_NUM]; + HcclAiRMACQ cqs[MAX_RANK_NUM]; }; -struct LocalResInfoV2 { - uint32_t streamNum; - uint32_t signalNum; - HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM]; - HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM]; - HcclStreamInfo mainStreamInfo; - HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM]; // 集合通信AICPU展开资源 - ListCommon nextTagRes; // HccltagLocalResV2 -}; +// URMA通信接口 +__aicore__ inline GM_ADDR GetURMASqInfoGMAddr(GM_ADDR hcclContext, uint32_t dstRankId) +{ + ascendc_assert((hcclContext != nullptr && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or dstRankId >= MAX_RANK_NUM"); -enum class rtFloatOverflowMode_t { - RT_OVERFLOW_MODE_SATURATION = 0, - RT_OVERFLOW_MODE_INFNAN, - RT_OVERFLOW_MODE_UNDEF, -}; + return (GM_ADDR)((__gm__ HcclAiRMAWQ *)(((__gm__ HcclCombinOpParam *)hcclContext)->sqs) + dstRankId); +} -struct AlgoTopoInfo { - uint32_t userRank; // 通信域 RankID - uint32_t userRankSize; // 通信域的Rank数量 - int32_t deviceLogicId; - bool isSingleMeshAggregation; - uint32_t deviceNumPerAggregation; // 每个Module中的Device数量 - uint32_t superPodNum; // 集群中总的超节点数 - uint32_t devicePhyId; - uint32_t topoType; // TopoType - uint32_t deviceType; - uint32_t serverNum; - uint32_t meshAggregationRankSize; - uint32_t multiModuleDiffDeviceNumMode; - uint32_t multiSuperPodDiffServerNumMode; - uint32_t realUserRank; - bool isDiffDeviceModule; - bool isDiffDeviceType; - uint32_t gcdDeviceNumPerAggregation; - uint32_t moduleNum; - uint32_t isUsedRdmaRankPairNum; - uint64_t isUsedRdmaRankPair; - uint32_t pairLinkCounterNum; - uint64_t pairLinkCounter; - uint32_t nicNum; - uint64_t nicList; // niclist数组指针 - uint64_t complanRankLength; // complanRank占用的字节数 - uint64_t complanRank; // 指针 - uint64_t bridgeRankNum; // bridgeRank占用的个数 - uint64_t bridgeRank; // 指针 - uint64_t serverAndsuperPodRankLength; // serverAndsuperPodRank占用的字节数 - uint64_t serverAndsuperPodRank; // 指针 -}; +__aicore__ inline GM_ADDR GetURMACqInfoGMAddr(GM_ADDR hcclContext, uint32_t dstRankId) +{ + ascendc_assert((hcclContext != nullptr && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or dstRankId >= MAX_RANK_NUM"); -struct HcclOpConfig { - uint8_t deterministic; //确定性计算开关 - uint8_t retryEnable; // 是否重执行 - uint8_t highPerfEnable; - uint8_t padding[5]; // 大小需要64By对齐，未来添加参数时减小padding - uint8_t linkTimeOut[8]; // 发送超时时长 - uint64_t notifyWaitTime; // 超时时长，同HCCL_EXEC_TIMEOUT - uint32_t retryHoldTime; - uint32_t retryIntervalTime; - bool interHccsDisable = false; //使能rdma开关 - rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF; - uint32_t multiQpThreshold = 512; // 多QP每个QP分担数据量最小阈值 -}; + return (GM_ADDR)((__gm__ HcclAiRMACQ *)(((__gm__ HcclCombinOpParam *)hcclContext)->cqs) + dstRankId); +} -struct HcclMC2WorkSpace { - uint64_t workSpace; - uint64_t workSpaceSize; -}; +__aicore__ inline void GetURMASqInfoTensor(const AscendC::LocalTensor &sqInfoTensor, GM_ADDR hcclContext, + uint32_t dstRankId) +{ + GM_ADDR sqInfoGMAddr = GetURMASqInfoGMAddr(hcclContext, dstRankId); + ascendc_assert((sqInfoGMAddr != nullptr), "GetURMASqInfoGMAddr failed"); + + AscendC::GlobalTensor sqInfoGlobalTensor; + sqInfoGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoGMAddr)); + AscendC::DataCopyExtParams sqInfoParams = {1U, static_cast(sizeof(HcclAiRMAWQ)), 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams sqInfoPadParams{false, 0U, 0U, 0U}; + AscendC::DataCopyPad(sqInfoTensor, sqInfoGlobalTensor, sqInfoParams, sqInfoPadParams); +} -struct RemoteResPtr { - uint64_t nextHostPtr; - uint64_t nextDevicePtr; -}; +__aicore__ inline void GetURMACqInfoTensor(const AscendC::LocalTensor &cqInfoTensor, GM_ADDR hcclContext, + uint32_t dstRankId) +{ + GM_ADDR cqInfoGMAddr = GetURMACqInfoGMAddr(hcclContext, dstRankId); + ascendc_assert((cqInfoGMAddr != nullptr), "GetURMACqInfoGMAddr failed"); + + AscendC::GlobalTensor cqInfoGlobalTensor; + cqInfoGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoGMAddr)); + AscendC::DataCopyExtParams cqInfoParams = {1U, static_cast(sizeof(HcclAiRMACQ)), 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams cqInfoPadParams{false, 0U, 0U, 0U}; + AscendC::DataCopyPad(cqInfoTensor, cqInfoGlobalTensor, cqInfoParams, cqInfoPadParams); +} -struct HDCommunicateParams { - uint64_t hostAddr { 0 }; - uint64_t deviceAddr { 0 }; - uint64_t readCacheAddr { 0 }; - uint32_t devMemSize{ 0 }; - uint32_t buffLen{ 0 }; - uint32_t flag{ 0 }; -}; +__aicore__ inline void GenerateCommWriteSQE(const AscendC::LocalTensor &sqeTensor, uint32_t sqeCount = 1) +{ + AscendC::Duplicate(sqeTensor, 0, WRITE_SQE_SIZE); // 初始化为全0 + AscendC::SyncFunc(); + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) = 0b001; // odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_3) = 0b10110000; // owner=1, rmt_jetty_type=0b01, token_en=1, nf=0 + sqeTensor(SQE_COMMON_OPCODE_OFFSET) = 0x3; // opcode=0x3 + sqeTensor(SQE_COMMON_SGE_NUM_OFFSET) = 1; // sge_num=1 +} -struct HcclRankRelationResV2Custom { - uint32_t remoteUsrRankId; - uint32_t remoteWorldRank; - uint64_t windowsIn; - uint64_t windowsOut; - uint64_t windowsExp; - ListCommon nextTagRes; -}; +__aicore__ inline void GenerateCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, + uint32_t sqeCount = 1) +{ + AscendC::Duplicate(sqeTensor, 0, WRITE_WITH_NOTIFY_SQE_SIZE); // 初始化为全0 + AscendC::SyncFunc(); + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) = 0b001; // odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_3) = 0b10110000; // owner=1, rmt_jetty_type=0b01, token_en=1, nf=0 + sqeTensor(SQE_COMMON_OPCODE_OFFSET) = 0x5; // opcode=0x5，使用write with notify SQE + sqeTensor(SQE_COMMON_SGE_NUM_OFFSET) = 1; // sge_num=1 +} -struct HcclOpResParamCustom { - // 本地资源 - HcclMC2WorkSpace mc2WorkSpace; - uint32_t localUsrRankId; // usrrankid - uint32_t rankSize; // 通信域内total rank个数 - uint64_t winSize; // 每个win大小，静态图时，可能是0，如果通信域内也有动态图，则可能为非0 - uint64_t localWindowsIn; // 全F为无效值 - uint64_t localWindowsOut; // 全F为无效值 - char hcomId[128]; - // aicore识别remote window - uint64_t winExpSize; - uint64_t localWindowsExp; - uint32_t rWinStart; // 为HcclRankRelationRes起始位置 - uint32_t rWinOffset; // 为HcclRemoteRes的大小 - uint64_t version; - LocalResInfoV2 localRes; - AlgoTopoInfo topoInfo; - - // 外部配置参数 - HcclOpConfig config; - uint64_t hostStateInfo; - uint64_t aicpuStateInfo; - uint64_t lockAddr; - uint32_t rsv[16]; - uint32_t notifysize; // RDMA场景使用，910B/910_93为4B，其余芯片为8B - uint32_t remoteResNum; // 有效的remoteResNum - RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM]; //数组指针，指向HcclRankRelationResV2，下标为remoteUserRankId - - // communicate retry - HDCommunicateParams kfcControlTransferH2DParams; - HDCommunicateParams kfcStatusTransferD2HParams; - uint64_t tinyMem; // for all2all - uint64_t tinyMemSize; - // 零拷贝场景使用 - uint64_t zeroCopyHeadPtr; - uint64_t zeroCopyTailPtr; - uint64_t zeroCopyRingBuffer; - uint64_t zeroCopyIpcPtrs[16]; // 保存集合通信时每个对端的输入输出内存地址 - uint32_t zeroCopyDevicePhyId[16]; // 保存每个rank对应的物理卡Id - - bool utraceStatusFlag; -}; +__aicore__ inline void UpdateCommonSQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + // 更新tp_id(24b), jetty_id(20b), rmt_eid(128b), rmt_token_value(32b) + + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_COMMON_UINT32_OFFSET_2) &= 0xff000000; // 保留sge_num + templateSqeU32(SQE_COMMON_UINT32_OFFSET_2) |= (sqInfoU32(WQ_TP_ID_OFFSET) & 0x00ffffff); // tp_id + templateSqeU32(SQE_COMMON_RMT_JETTY_OR_SEG_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_jetty_or_seg_id + templateSqeU32(SQE_COMMON_RMT_EID_31_0_OFFSET) = sqInfoU32(WQ_RMTEID_0_3_OFFSET); // rmt_eid + templateSqeU32(SQE_COMMON_RMT_EID_63_32_OFFSET) = sqInfoU32(WQ_RMTEID_4_7_OFFSET); + templateSqeU32(SQE_COMMON_RMT_EID_95_64_OFFSET) = sqInfoU32(WQ_RMTEID_8_11_OFFSET); + templateSqeU32(SQE_COMMON_RMT_EID_127_96_OFFSET) = sqInfoU32(WQ_RMTEID_12_15_OFFSET); + templateSqeU32(SQE_COMMON_RMT_TOKEN_VALUE_OFFSET) = sqInfoU32(WQ_RMTTOKENVALUE_OFFSET); // rmt_token_value +} -// Transport 内存类型 -enum class HcclAiRMAMemType : uint32_t { - LOCAL_INPUT = 0, - REMOTE_INPUT, +__aicore__ inline void UpdateCommWriteSQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + UpdateCommonSQE(sqeTensor, sqInfoTensor); - LOCAL_OUTPUT, - REMOTE_OUTPUT, + // 更新rmt_token_id(20b) + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_token_id +} - // 可透传更多的内存，可在MAX_NUM之前追加，例如： - // LOCAL_EXP, - // REMOTE_EXP, - MAX_NUM -}; +__aicore__ inline void UpdateCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, + const AscendC::LocalTensor &sqInfoTensor) +{ + UpdateCommonSQE(sqeTensor, sqInfoTensor); + + // 更新rmt_token_id(20b), notify_token_value(32b), notify_token_id(20b) + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor templateSqeU32 = sqeTensor.ReinterpretCast(); + templateSqeU32(SQE_WITH_NOTIFY_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // rmt_token_id + templateSqeU32(SQE_WITH_NOTIFY_NOTIFY_TOKEN_VALUE_OFFSET) = + sqInfoU32(WQ_RMTTOKENVALUE_OFFSET); // notify_token_value + templateSqeU32(SQE_WITH_NOTIFY_NOTIFY_TOKEN_ID_OFFSET) = sqInfoU32(WQ_RMTOBJID_OFFSET); // notify_token_id +} -// Transport 内存信息 -struct HcclAiRMAMemInfo { - uint32_t memMaxNum{0}; // 最大内存数量，等于 HcclAiRMAMemType::MAX_NUM - uint32_t sizeOfMemDetails{0}; // sizeof(MemDetails)，用于内存校验和偏移计算 - uint64_t memDetailPtr{0}; // MemDetails数组首地址, 个数: HcclAiRMAMemType::MAX_NUM - // 可往后追加字段 -}; +__aicore__ inline void SetCommWriteSQE(const AscendC::LocalTensor &sqeTensor, uint64_t dataAddr, + uint64_t rmtAddr, uint32_t length, uint8_t cqe) +{ + // 设置cqe(1 bit), odr(3 bits), data_addr(64 bits), rmt_addr(64 bits), length(32 bits) + + AscendC::LocalTensor sqeLocalU32 = sqeTensor.ReinterpretCast(); + AscendC::LocalTensor sqeLocalU64 = sqeTensor.ReinterpretCast(); + + sqeLocalU64(SQE_RMT_ADDR_OFFSET) = rmtAddr; + sqeLocalU64(SQE_DATA_ADDR_OFFSET) = dataAddr; + sqeLocalU32(SQE_LENGTH_OFFSET) = length; + if (cqe > 0U) { + // cqe=1, odr=0b010 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00100010; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11111010; + } else { + // cqe=0, odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11011001; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00000001; + } +} -// 全部 Transport QP/Mem 信息 -struct HcclAiRMAInfo { - uint32_t curRankId{0}; // 当前rankId - uint32_t rankNum{0}; // rank数量 - uint32_t qpNum{0}; // 单个Transport的QP数量 - - uint32_t sizeOfAiRMAWQ{0}; // sizeof(HcclAiRMAWQ) - uint32_t sizeOfAiRMACQ{0}; // sizeof(HcclAiRMACQ) - uint32_t sizeOfAiRMAMem{0}; // sizeof(HcclAiRMAMemInfo) - - // HcclAiRMAWQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取SQ指针：sqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ - // 0 <= qpIndex < qpNum - uint64_t sqPtr{0}; - - // HcclAiRMACQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取SCQ指针：scqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ - // 0 <= qpIndex < qpNum - uint64_t scqPtr{0}; - - // HcclAiRMAWQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取RQ指针：rqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ - // 0 <= qpIndex < qpNum - uint64_t rqPtr{0}; - - // HcclAiRMACQ二维数组首地址 - // QP个数: rankNum * qpNum - // 计算偏移获取RCQ指针: rcqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ - // 0 <= qpIndex < qpNum - uint64_t rcqPtr{0}; - - // HcclAivMemInfo一维数组 - // 内存信息个数: rankNum - // 计算偏移获取内存信息指针: memPtr + rankId * sizeOfAiRMAMem - // srcRankId 获取自身内存信息，dstRankId 获取 Transport 内存信息 - uint64_t memPtr{0}; - // 可往后追加字段 -}; -struct CombinedCapability { - uint64_t dataplaneModeBitmap; -}; +__aicore__ inline void SetCommWriteWithNotifySQE(const AscendC::LocalTensor &sqeTensor, uint64_t dataAddr, + uint64_t rmtAddr, uint32_t length, uint64_t notifyAddr, + uint64_t notifyData, uint8_t cqe) +{ + // 设置cqe(1 bit), odr(3 bits), data_addr(64 bits), rmt_addr(64 bits), length(32 bits), notify_addr(64 bits), + // notify_data(64 bits) + + AscendC::LocalTensor sqeLocalU32 = sqeTensor.ReinterpretCast(); + AscendC::LocalTensor sqeLocalU64 = sqeTensor.ReinterpretCast(); + + sqeLocalU64(SQE_WITH_NOTIFY_RMT_ADDR_OFFSET) = rmtAddr; + sqeLocalU64(SQE_WITH_NOTIFY_DATA_ADDR_OFFSET) = dataAddr; + sqeLocalU32(SQE_WITH_NOTIFY_LENGTH_OFFSET) = length; + sqeLocalU64(SQE_WITH_NOTIFY_NOTIFY_ADDR_OFFSET) = notifyAddr; // notify_addr + sqeLocalU64(SQE_WITH_NOTIFY_NOTIFY_DATA_OFFSET) = notifyData; // notify_data + + if (cqe > 0U) { + // cqe=1, odr=0b010 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00100010; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11111010; + } else { + // cqe=0, odr=0b001 + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) &= 0b11011001; + sqeTensor(SQE_COMMON_UINT8_OFFSET_2) |= 0b00000001; + } +} -struct HcclA2CombineOpParam { - uint64_t workSpace; // Address for communication between client and server, - // hccl requests and clears - uint64_t workSpaceSize; // Space for communication between client and server - uint32_t rankId; // id of this rank - uint32_t rankNum; // num of ranks in this comm group - uint64_t winSize; // size of each windows memory - uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM]; // windows address for input, windowsIn[rankId] corresponds - // to the local card address, - // and others are cross-card mapping addresses. - uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM]; // windows address for output, windowsOut[rankId] corresponds - // to the local card address, - // and others are cross-card mapping addresses. - uint8_t res[8328]; - uint8_t multiFlag; - __gm__ AscendC::IbVerbsData *data; - uint64_t dataSize; - // 追加字段 - uint64_t sizeOfAiRMAInfo; // sizeof(HcclAiRMAInfo) - uint64_t aiRMAInfo; // HcclAiRMAInfo* 单个结构体指针 - - CombinedCapability* capability; // address of the communication capability information structure on the Device - uint64_t capabilitySize; // size of the communication capability information structure -}; -enum class DataplaneMode : uint32_t { - HOST = 0, - AICPU = 1, - AIV = 2, -}; +__aicore__ inline void SendJFSDoorBell(const AscendC::LocalTensor &jfsDoorBellTensor, + const AscendC::LocalTensor &sqInfoTensor, uint32_t sqPi) +{ + // 敲JFS DoorBell，更新硬件的sqPi + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + st_dev(sqPi, (__gm__ uint32_t *)(sqInfoU64(WQ_DBADDR_OFFSET)), 0); // Scalar操作 +} -enum class DBMode : int32_t { - INVALID_DB = -1, - HW_DB = 0, - SW_DB -}; +__aicore__ inline void SendJFCDoorBell(const AscendC::LocalTensor &jfcDoorBellTensor, + const AscendC::LocalTensor &cqInfoTensor, uint32_t cqCi) +{ + // 敲JFC DoorBell，更新硬件的cqCi + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU64 = cqInfoTensor.ReinterpretCast(); -struct HcclAiRMAWQ { - uint32_t wqn{0}; - uint64_t bufAddr{0}; - uint32_t wqeSize{0}; - uint32_t depth{0}; - uint64_t headAddr{0}; - uint64_t tailAddr{0}; - DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw - uint64_t dbAddr{0}; - uint32_t sl{0}; -}; + uint64_t jfcDbValue = (static_cast(cqInfoU32(CQ_JFCID_OFFSET)) << 32) + static_cast(cqCi); + st_dev(jfcDbValue, (__gm__ uint64_t *)(cqInfoU64(CQ_DBADDR_OFFSET)), 0); // Scalar操作 +} -struct HcclAiRMACQ { - uint32_t cqn{0}; - uint64_t bufAddr{0}; - uint32_t cqeSize{0}; - uint32_t depth{0}; - uint64_t headAddr{0}; - uint64_t tailAddr{0}; - DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw - uint64_t dbAddr{0}; -}; +__aicore__ inline bool CheckCqeStatus(const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &cqeTensor, uint32_t &outCqCi, + uint32_t &outCqCiLinear, uint32_t &newestCompletedPi) +{ + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU64 = cqInfoTensor.ReinterpretCast(); + + uint32_t cqeSize = cqInfoU32(CQ_CQESIZE_OFFSET); + AscendC::GlobalTensor cqGlobalTensor; + cqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoU64(CQ_CQVA_OFFSET))); + AscendC::DataCopyExtParams cqeParams = {1U, 8U, 0U, 0U, 0U}; + AscendC::DataCopyPadExtParams cqePadParams{false, 0U, 0U, 0U}; + + AscendC::DataCopyPad(cqeTensor, cqGlobalTensor[cqeSize * outCqCi], cqeParams, cqePadParams); + AscendC::SyncFunc(); + if (cqeTensor(CQE_STATUS_OFFSET) == 0xff) { + return true; + } else if (cqeTensor(CQE_STATUS_OFFSET) != 0) { + // 退出kernel并报错 + ascendc_assert(false, "CQE status is abnormal! status is %d, substatus is %d.\n", cqeTensor(CQE_STATUS_OFFSET), + cqeTensor(CQE_SUBSTATUS_OFFSET)); + } -struct hns_roce_rc_sq_wqe { - uint32_t byte_4; - uint32_t msg_len; - uint32_t immtdata; - uint32_t byte_16; - uint32_t byte_20; - uint32_t rkey; - uint64_t remoteVA; -}; + // status==0，处理当前CQE，从entry_idx获取对应WQE的sqPi + newestCompletedPi = (static_cast(cqeTensor(CQE_ENTRY_IDX_HIGH_OFFSET)) << UINT8_BITS_OFFSET) + + static_cast(cqeTensor(CQE_ENTRY_IDX_LOW_OFFSET)); + // 把CQE的status设置为无效值，并写回CQ + cqeTensor(CQE_STATUS_OFFSET) = 0xff; + AscendC::SyncFunc(); + AscendC::DataCopyPad(cqGlobalTensor[cqeSize * outCqCi], cqeTensor, cqeParams); + AscendC::SyncFunc(); // 防止cqeTensor被下一轮加载覆盖 + return false; +} -struct hns_roce_lite_wqe_data_seg { - uint32_t len; - uint32_t lkey; - uint64_t localVA; -}; +__aicore__ inline void PollCommCQUpdateSQCI(const AscendC::LocalTensor &sqInfoTensor, + const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t &outSqCi, + uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU32 = cqInfoTensor.ReinterpretCast(); + + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + uint32_t cqDepth = cqInfoU32(CQ_CQDEPTH_OFFSET); + + uint32_t pollTimes = 0; + uint32_t newestCompletedPi = 0; + while (pollTimes < cqDepth - 1) { + if (CheckCqeStatus(cqInfoTensor, cqeTensor, outCqCi, outCqCiLinear, newestCompletedPi)) { + break; + } + + // 递增本地的outCqCi + outCqCi = (outCqCi + 1) % cqDepth; + ++outCqCiLinear; + ++pollTimes; + } + + // 更新本地的outSqCi + if (pollTimes > 0) { + outSqCi = (newestCompletedPi + 1) % sqDepth; + } + AscendC::SyncFunc(); // 等cqe status写回无效值 + // 通过敲JFC DoorBell，更新硬件的cqCi + SendJFCDoorBell(jfcDoorBellTensor, cqInfoTensor, outCqCiLinear); +} + +__aicore__ inline void InvalidateCqeStatus(const AscendC::LocalTensor &cqeInfoTensor, + const AscendC::LocalTensor &cqeTensor) +{ + // 为cqeBuf_申请256*32B空间，初始化HGM上的CQ空间时，如果cqDepth>256，则循环多次DataCopy + AscendC::LocalTensor cqInfoU64 = cqeInfoTensor.ReinterpretCast(); + AscendC::LocalTensor cqInfoU32 = cqeInfoTensor.ReinterpretCast(); + + AscendC::GlobalTensor cqGlobalTensor; + cqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(cqInfoU64(CQ_CQVA_OFFSET))); + uint32_t cqDepth = cqInfoU32(CQ_CQDEPTH_OFFSET); + uint32_t cqeNum = (cqDepth > CQ_DEPTH_256) ? CQ_DEPTH_256 : cqDepth; + + AscendC::Duplicate(cqeTensor, 0xff, UB_ALIGN * cqeNum); // 初始化为全1 + AscendC::SyncFunc(); + + AscendC::DataCopyExtParams dataCopyParams = {static_cast(cqeNum), UB_ALIGN, 0U, + NORMAL_CQE_SIZE - UB_ALIGN, 0U}; + uint32_t loopTimes = AscendC::Ceil(cqDepth, CQ_DEPTH_256); + if (loopTimes > 1) { + for (uint32_t i = 0; i < loopTimes - 1; ++i) { + AscendC::DataCopyPad(cqGlobalTensor[NORMAL_CQE_SIZE * CQ_DEPTH_256 * i], cqeTensor, dataCopyParams); + } + uint32_t tailBlockCqeNum = cqDepth - CQ_DEPTH_256 * (loopTimes - 1); + if (tailBlockCqeNum != CQ_DEPTH_256) { + dataCopyParams = {static_cast(tailBlockCqeNum), UB_ALIGN, 0U, NORMAL_CQE_SIZE - UB_ALIGN, 0U}; + } + } + AscendC::DataCopyPad(cqGlobalTensor[NORMAL_CQE_SIZE * CQ_DEPTH_256 * (loopTimes - 1)], cqeTensor, dataCopyParams); +} + +__aicore__ inline uint32_t GetAvailableSpace(const AscendC::LocalTensor &sqInfoTensor, uint32_t &outSqPi, + uint32_t &outSqCi) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + + // 根据本地的outSqPi outSqCi 计算SQ可用空间大小（SQ不会用满，会预留一个WQE） + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + + if (outSqPi == outSqCi) { + return (sqDepth - 1); + } else if ((outSqPi + 1) % sqDepth == outSqCi) { + return 0; + } else if (outSqPi > outSqCi) { + return (sqDepth - (outSqPi - outSqCi) - 1); + } + return (outSqCi - outSqPi - 1); +} + +__aicore__ inline void +PutCommSQE(const AscendC::LocalTensor &sqInfoTensor, const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &sqeTensor, const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t sqeCount, uint32_t &outSqPi, + uint32_t &outSqPiLinear, uint32_t &outSqCi, uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + uint32_t sqeSize = sqInfoU32(WQ_WQESIZE_OFFSET); // 单个WQEBB占用大小（64B），由HCCL传递 + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + ascendc_assert(sqeCount < sqDepth, "too many SQE! SQE num[%d] should less than sqDepth[%d].", sqeCount, sqDepth); + + // 计算SQ可用空间大小 + uint32_t availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + while (availableSpace < sqeCount) { + // 可用空间不足时，轮询CQ，更新本地的cqCi_ sqCi_和硬件的cqCi + PollCommCQUpdateSQCI(sqInfoTensor, cqInfoTensor, cqeTensor, jfcDoorBellTensor, outSqCi, outCqCi, outCqCiLinear); + availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + } + + // 可用空间足够时，把WQE拷贝到HBM上的SQ中 + AscendC::GlobalTensor sqGlobalTensor; + sqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoU64(WQ_SQVA_OFFSET))); -__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) { - __gm__ uint8_t* start = - (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); - __gm__ uint8_t* end = - (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); - AscendC::GlobalTensor global; - global.SetGlobalBuffer(start); - for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) { - AscendC::DataCacheCleanAndInvalid(global[i]); + if (likely(sqDepth - outSqPi >= sqeCount)) { + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, WRITE_SQE_SIZE * sqeCount); + } else { + uint32_t firstPartSize = sqDepth - outSqPi; + uint32_t secondPartSize = sqeCount - firstPartSize; + + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, WRITE_SQE_SIZE * firstPartSize); + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_SQE_SIZE * firstPartSize], WRITE_SQE_SIZE * secondPartSize); } + + // 更新本地的outSqPi + outSqPiLinear += sqeCount; + outSqPi = outSqPiLinear % sqDepth; } -__aicore__ inline DataplaneMode GetDataplaneMode(GM_ADDR contextGM0) { - __gm__ HcclA2CombineOpParam *winContext_ = (__gm__ HcclA2CombineOpParam *)contextGM0; - CombinedCapability* capability = winContext_->capability; - uint64_t capabilitySize = winContext_->capabilitySize; - DataplaneMode dataplaneMode = DataplaneMode::AICPU; - if (capability == 0) { - return dataplaneMode; + +__aicore__ inline void +PutCommNotifySQE(const AscendC::LocalTensor &sqInfoTensor, const AscendC::LocalTensor &cqInfoTensor, + const AscendC::LocalTensor &sqeTensor, const AscendC::LocalTensor &cqeTensor, + const AscendC::LocalTensor &jfcDoorBellTensor, uint32_t sqeCount, uint32_t &outSqPi, + uint32_t &outSqPiLinear, uint32_t &outSqCi, uint32_t &outCqCi, uint32_t &outCqCiLinear) +{ + AscendC::LocalTensor sqInfoU32 = sqInfoTensor.ReinterpretCast(); + AscendC::LocalTensor sqInfoU64 = sqInfoTensor.ReinterpretCast(); + uint32_t sqeSize = sqInfoU32(WQ_WQESIZE_OFFSET); // 单个WQEBB占用大小（64B），由HCCL传递 + uint32_t sqDepth = sqInfoU32(WQ_SQDEPTH_OFFSET); + ascendc_assert((sqeCount << 1) < sqDepth, "too many SQE! SQE num[%d] should less than sqDepth[%d].", + (sqeCount << 1), sqDepth); + + // 计算SQ可用空间大小 + uint32_t availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); + while (availableSpace < (sqeCount << 1)) { + // 可用空间不足时，轮询CQ，更新本地的cqCi_ sqCi_和硬件的cqCi + PollCommCQUpdateSQCI(sqInfoTensor, cqInfoTensor, cqeTensor, jfcDoorBellTensor, outSqCi, outCqCi, outCqCiLinear); + availableSpace = GetAvailableSpace(sqInfoTensor, outSqPi, outSqCi); } - uint64_t dataplaneModeBitmap = capability->dataplaneModeBitmap; - if ((dataplaneModeBitmap & 0x04) == 0x04) { - dataplaneMode = DataplaneMode::AIV; + + // 可用空间足够时，把WQE拷贝到HBM上的SQ中 + AscendC::GlobalTensor sqGlobalTensor; + sqGlobalTensor.SetGlobalBuffer((__gm__ uint8_t *)(sqInfoU64(WQ_SQVA_OFFSET))); + + if (likely((sqDepth - outSqPi) >= (sqeCount << 1))) { + AscendC::DataCopyParams intriParams{static_cast(sqeCount), WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, intriParams); + } else { + uint32_t firstPartSize = (sqDepth - outSqPi) >> 1; + uint32_t secondPartSize = sqeCount - firstPartSize; + AscendC::DataCopyParams firstPartParams{static_cast(firstPartSize), WRITE_WITH_NOTIFY_SQE_SIZE >> 5, + 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize * outSqPi], sqeTensor, firstPartParams); + + if (((sqDepth - outSqPi) & 1) == 0) { + AscendC::DataCopyParams secondPartParams{static_cast(secondPartSize), + WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize], secondPartParams); + } else { + // 如果剩余奇数个WQEBB，其中一个WQE会被切块分在尾和头 + AscendC::DataCopy(sqGlobalTensor[sqeSize * (sqDepth - 1)], + sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize], sqeSize); + AscendC::DataCopy(sqGlobalTensor, sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * firstPartSize + sqeSize], + WRITE_WITH_NOTIFY_SQE_SIZE - sqeSize); + + AscendC::DataCopyParams secondPartParams{static_cast(secondPartSize - 1), + WRITE_WITH_NOTIFY_SQE_SIZE >> 5, 0, 1}; + AscendC::DataCopy(sqGlobalTensor[sqeSize], sqeTensor[WRITE_WITH_NOTIFY_SQE_SIZE * (firstPartSize + 1)], + secondPartParams); + } } - return dataplaneMode; + + // 更新本地的outSqPi + outSqPiLinear += (sqeCount << 1); + outSqPi = outSqPiLinear % sqDepth; } -__aicore__ inline int64_t GetCurrentTimestampUs() +__aicore__ inline void GetPICI(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, uint32_t &outSqPi, + uint32_t &outSqCi, uint32_t &outCqPi, uint32_t &outCqCi, uint32_t &outSqPiLinear, + uint32_t &outCqCiLinear) { - return AscendC::GetSystemCycle() / TIME_CYCLE; + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); + + outSqPi = piCiGlobalTensor(WIN_SQPI_OFFSET); + outSqCi = piCiGlobalTensor(WIN_SQCI_OFFSET); + outCqPi = piCiGlobalTensor(WIN_CQPI_OFFSET); + outCqCi = piCiGlobalTensor(WIN_CQCI_OFFSET); + outSqPiLinear = piCiGlobalTensor(WIN_SQPILINEAR_OFFSET); + outCqCiLinear = piCiGlobalTensor(WIN_CQCILINEAR_OFFSET); } -__aicore__ inline void RecordRankCommDuration(AscendC::LocalTensor performanceInfoU32Tensor, uint32_t rankId, int64_t startTime) +__aicore__ inline void GetIsFirstInComm(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, bool &outIsFirst) { - int64_t endTime = GetCurrentTimestampUs(); - int32_t duration = static_cast(endTime - startTime); // int32_t可以表示2^31(us)，约35min在实际场景下满足需要 - performanceInfoU32Tensor.SetValue(rankId * sizeof(int64_t) / sizeof(int32_t), duration); // 使用int32_t是因为atomicAdd不支持int64_t类型，这里只赋值到int64_t的低32位。 + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); + + outIsFirst = (piCiGlobalTensor(WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET) > 0); +} + +__aicore__ inline void UpdatePICI(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, uint32_t sqPi, + uint32_t sqCi, uint32_t cqPi, uint32_t cqCi, uint32_t sqPiLinear, uint32_t cqCiLinear) +{ + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + + piCiGlobalTensor(WIN_SQPI_OFFSET) = sqPi; + piCiGlobalTensor(WIN_SQCI_OFFSET) = sqCi; + piCiGlobalTensor(WIN_CQPI_OFFSET) = cqPi; + piCiGlobalTensor(WIN_CQCI_OFFSET) = cqCi; + piCiGlobalTensor(WIN_SQPILINEAR_OFFSET) = sqPiLinear; + piCiGlobalTensor(WIN_CQCILINEAR_OFFSET) = cqCiLinear; + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); } -#endif // MOE_DISTRIBUTE_BASE_H \ No newline at end of file + +__aicore__ inline void UpdateIsFirstInComm(GM_ADDR hcclContext, uint32_t curRankId, uint32_t dstRankId, bool isFirst) +{ + ascendc_assert((hcclContext != nullptr && curRankId < MAX_RANK_NUM && dstRankId < MAX_RANK_NUM), + "hcclContext is nullptr or curRankId >= MAX_RANK_NUM or dstRankId >= MAX_RANK_NUM"); + + GM_ADDR piCiSpaceGM = (GM_ADDR)(((__gm__ HcclCombinOpParam *)hcclContext)->windowsOut[curRankId] + WIN_PICI_OFFSET + + PICI_WIN_SIZE * dstRankId); + AscendC::GlobalTensor piCiGlobalTensor; + piCiGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)piCiSpaceGM); + + piCiGlobalTensor(WIN_FIRST_TIME_CREATE_WIN_FLAG_OFFSET) = isFirst ? 1 : 0; + AscendC::DataCacheCleanAndInvalid( + piCiGlobalTensor); +} + +#endif // MOE_DISTRIBUTE_BASE_H