Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #include "all_gather_mesh_mix.h"
      12              : #include "alg_template_register.h"
      13              : 
      14              : namespace hccl {
      15            0 : AllgatherMeshMix::AllgatherMeshMix(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
      16              : 
      17            0 : AllgatherMeshMix::~AllgatherMeshMix() {}
      18              : 
      19            0 : HcclResult AllgatherMeshMix::Prepare(
      20              :     std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
      21              :     std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, [[maybe_unused]] u32 userRank, HcomCollOpInfo* opInfo,
      22              :     u32 interRank, u32 interRankSize)
      23              : {
      24            0 :     meshStreams_ = meshStreams;
      25            0 :     meshSignal_ = &meshSignal;
      26            0 :     meshSignalAux_ = &meshSignalAux;
      27            0 :     interRank_ = interRank;
      28            0 :     interRankSize_ = interRankSize;
      29            0 :     opInfo_ = opInfo;
      30            0 :     return HCCL_SUCCESS;
      31              : }
      32              : 
      33            0 : HcclResult AllgatherMeshMix::MainRecordSub()
      34              : {
      35            0 :     for (u32 signalIndex = 0; signalIndex < (*meshSignalAux_).size(); signalIndex++) {
      36            0 :         CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex], profilerInput_.stage));
      37              :     }
      38            0 :     return HCCL_SUCCESS;
      39              : }
      40              : 
      41            0 : HcclResult AllgatherMeshMix::SubWaitMain()
      42              : {
      43            0 :     for (u32 streamIndex = 0; streamIndex < (*meshSignalAux_).size(); streamIndex++) {
      44            0 :         CHK_RET(LocalNotify::Wait(
      45              :             meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
      46              :     }
      47            0 :     return HCCL_SUCCESS;
      48              : }
      49              : 
      50            0 : HcclResult AllgatherMeshMix::MainWaitSub()
      51              : {
      52            0 :     for (u32 signalIndex = 0; signalIndex < (*meshSignal_).size(); signalIndex++) {
      53            0 :         CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
      54              :     }
      55            0 :     return HCCL_SUCCESS;
      56              : }
      57              : 
      58            0 : HcclResult AllgatherMeshMix::SubRecordMain()
      59              : {
      60            0 :     for (u32 streamIndex = 0; streamIndex < (*meshSignal_).size(); streamIndex++) {
      61            0 :         CHK_RET(LocalNotify::Post(
      62              :             meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
      63              :     }
      64            0 :     return HCCL_SUCCESS;
      65              : }
      66              : 
      67              : // allgather的入口函数
      68            0 : HcclResult AllgatherMeshMix::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
      69              : {
      70            0 :     HCCL_INFO(
      71              :         "AllGatherMesh run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
      72              :         inputMem_.ptr(), outputMem_.ptr(), count_);
      73            0 :     u32 unitSize = DataUnitSize(dataType_);
      74            0 :     u64 sliceSize = count_ * unitSize;         // 当前count
      75            0 :     u64 totalSize = opInfo_->count * unitSize; // 总输入count
      76              : 
      77            0 :     u8* curUerMemOutPtr = static_cast<u8*>(opInfo_->outputAddr);
      78            0 :     u8* curCommMemOutPtr = static_cast<u8*>(outputMem_.ptr());
      79              : 
      80            0 :     CHK_RET(MainRecordSub());
      81            0 :     CHK_RET(SubWaitMain());
      82              : 
      83            0 :     for (u32 round = 1; round < rankSize; round++) {
      84            0 :         u32 dstRank = BackwardRank(rank, rankSize, round);
      85            0 :         Stream& subStream = meshStreams_[round - 1];
      86            0 :         CHK_RET(links[dstRank]->TxAck(subStream));
      87            0 :         CHK_RET(links[dstRank]->RxAck(subStream));
      88              :     }
      89              : 
      90            0 :     CHK_RET(SubRecordMain());
      91            0 :     CHK_RET(MainWaitSub());
      92              : 
      93            0 :     CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
      94              : 
      95            0 :     CHK_RET(SubWaitMain());
      96            0 :     CHK_RET(MainRecordSub());
      97              : 
      98            0 :     DeviceMem src;
      99            0 :     DeviceMem dst;
     100            0 :     for (u32 i = 0; i < interRankSize_; i++) {
     101            0 :         src = DeviceMem::create(curCommMemOutPtr + (i * rankSize + rank) * sliceSize, sliceSize);
     102            0 :         dst = DeviceMem::create(curUerMemOutPtr + (i * rankSize + rank) * totalSize, sliceSize);
     103            0 :         CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
     104              :     }
     105              : 
     106            0 :     for (u32 round = 1; round < rankSize; round++) {
     107            0 :         u32 dstRank = BackwardRank(rank, rankSize, round);
     108            0 :         Stream& subStream = meshStreams_[round - 1];
     109              :         // 本rank要收数据
     110            0 :         void* remMemPtr = nullptr;
     111              :         // 从对端的input内存拿数据，input==output也没有关系
     112            0 :         CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
     113              : 
     114            0 :         for (u32 i = 0; i < interRankSize_; i++) {
     115            0 :             src = DeviceMem::create(static_cast<u8*>(remMemPtr) + (i * rankSize + dstRank) * sliceSize, sliceSize);
     116            0 :             dst = DeviceMem::create(curUerMemOutPtr + (i * rankSize + dstRank) * totalSize, sliceSize);
     117            0 :             CHK_RET(HcclD2DMemcpyAsync(
     118              :                 dispatcher_, dst, src, subStream, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
     119              :         }
     120              : 
     121            0 :         CHK_RET(links[dstRank]->TxDataSignal(subStream));
     122            0 :         CHK_RET(links[dstRank]->RxDataSignal(subStream));
     123              :     }
     124            0 :     CHK_RET(SubRecordMain());
     125            0 :     CHK_RET(MainWaitSub());
     126            0 :     CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
     127              : 
     128            0 :     HCCL_INFO("AllGatherMesh finished: rank[%u]", rank);
     129            0 :     return HCCL_SUCCESS;
     130            0 : }
     131              : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_MESH_MIX, AllgatherMeshMix);
     132              : } // namespace hccl