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 "reduce_scatter_mesh_mix.h"
      12              : #include "alg_template_register.h"
      13              : 
      14              : namespace hccl {
      15              : using namespace std;
      16              : 
      17            1 : ReduceScatterMeshMix::ReduceScatterMeshMix(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
      18              : 
      19            2 : ReduceScatterMeshMix::~ReduceScatterMeshMix() {}
      20              : 
      21            1 : HcclResult ReduceScatterMeshMix::Prepare(
      22              :     DeviceMem& inputMem, DeviceMem& outputMem, DeviceMem& scratchMem, const u64 count, const HcclDataType dataType,
      23              :     const Stream& stream, const HcclReduceOp reductionOp, const u32 root, const std::vector<Slice>& slices,
      24              :     const u64 baseOffset, const u64 reduceAttrBitMap, std::vector<Stream>& meshStreams,
      25              :     const std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
      26              :     const std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 interRank, u32 interRankSize,
      27              :     HcomCollOpInfo* opInfo)
      28              : {
      29            1 :     reduceAttr_ = reduceAttrBitMap;
      30            1 :     meshStreams_ = meshStreams;
      31            1 :     meshSignalPtr_ = &meshSignal;
      32            1 :     meshSignalAuxPtr_ = &meshSignalAux;
      33            1 :     interRank_ = interRank;
      34            1 :     interRankSize_ = interRankSize;
      35            1 :     opInfo_ = opInfo;
      36            4 :     return AlgTemplateBase::Prepare(
      37            2 :         inputMem, outputMem, scratchMem, count, dataType, stream, reductionOp, root, slices, baseOffset);
      38              : }
      39              : 
      40            0 : HcclResult ReduceScatterMeshMix::MainRecordSub()
      41              : {
      42            0 :     for (u32 signalIndex = 0; signalIndex < meshSignalAuxPtr_->size(); signalIndex++) {
      43            0 :         CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
      44              :     }
      45            0 :     return HCCL_SUCCESS;
      46              : }
      47              : 
      48            0 : HcclResult ReduceScatterMeshMix::SubWaitMain()
      49              : {
      50            0 :     for (u32 streamIndex = 0; streamIndex < meshSignalAuxPtr_->size(); streamIndex++) {
      51            0 :         CHK_RET(LocalNotify::Wait(
      52              :             meshStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
      53              :     }
      54            0 :     return HCCL_SUCCESS;
      55              : }
      56              : 
      57            0 : HcclResult ReduceScatterMeshMix::MainWaitSub()
      58              : {
      59            0 :     for (u32 signalIndex = 0; signalIndex < meshSignalPtr_->size(); signalIndex++) {
      60            0 :         CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
      61              :     }
      62            0 :     return HCCL_SUCCESS;
      63              : }
      64              : 
      65            0 : HcclResult ReduceScatterMeshMix::SubRecordMain()
      66              : {
      67            0 :     for (u32 streamIndex = 0; streamIndex < meshSignalPtr_->size(); streamIndex++) {
      68            0 :         CHK_RET(LocalNotify::Post(
      69              :             meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
      70              :     }
      71            0 :     return HCCL_SUCCESS;
      72              : }
      73              : 
      74            0 : HcclResult ReduceScatterMeshMix::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
      75              : {
      76            0 :     HCCL_INFO(
      77              :         "ReduceScatterMeshMix run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
      78              :         inputMem_.ptr(), outputMem_.ptr(), count_);
      79              : 
      80              :     // 数据准备
      81            0 :     u32 unitSize = SIZE_TABLE[dataType_];
      82            0 :     u64 totalSize = opInfo_->count * unitSize;
      83            0 :     u64 sliceSize = count_ * unitSize;
      84              : 
      85            0 :     u8* curUerMemInPtr = static_cast<u8*>(opInfo_->inputAddr);
      86            0 :     u8* curCommMemOutPtr = static_cast<u8*>(outputMem_.ptr());
      87              : 
      88            0 :     DeviceMem userMemIn = DeviceMem::create(curUerMemInPtr + totalSize * rank, sliceSize);
      89            0 :     DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
      90              : 
      91            0 :     DeviceMem src;
      92            0 :     DeviceMem dst;
      93              : 
      94            0 :     for (u32 i = 0; i < interRankSize_; i++) {
      95            0 :         src = DeviceMem::create(curUerMemInPtr + (i * rankSize + rank) * totalSize, sliceSize);
      96            0 :         dst = DeviceMem::create(curCommMemOutPtr + (i * rankSize + rank) * sliceSize, sliceSize);
      97            0 :         CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
      98              :     }
      99              : 
     100            0 :     CHK_RET(MainRecordSub());
     101            0 :     CHK_RET(SubWaitMain());
     102              : 
     103              :     // 每个stream只负责一个对端的交互
     104            0 :     for (u32 round = 1; round < rankSize; round++) {
     105            0 :         u32 dstRank = (round + rank) % rankSize;
     106            0 :         const LINK& dstLink = links[dstRank];
     107            0 :         Stream& subStream = meshStreams_[round - 1];
     108            0 :         CHK_RET(dstLink->TxAck(subStream));
     109            0 :         CHK_RET(dstLink->RxAck(subStream));
     110              :     }
     111            0 :     CHK_RET(SubRecordMain());
     112            0 :     CHK_RET(MainWaitSub());
     113              :     // 为子图增加一个从stream到主stream的附着点
     114            0 :     CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
     115              : 
     116            0 :     CHK_RET(SubWaitMain());
     117            0 :     CHK_RET(MainRecordSub());
     118              : 
     119              :     // inline执行notice reduce
     120            0 :     for (u32 round = 1; round < rankSize; round++) {
     121            0 :         u32 dstRank = (round + rank) % rankSize;
     122            0 :         const LINK& dstLink = links[dstRank];
     123            0 :         Stream& subStream = meshStreams_[round - 1];
     124              :         // 本rank要发数据
     125            0 :         void* remMemPtr = nullptr;
     126              :         // 获取远端的commoutMem
     127            0 :         CHK_RET(dstLink->GetRemoteMem(UserMemType::INPUT_MEM, &remMemPtr));
     128              : 
     129            0 :         for (u32 i = 0; i < interRankSize_; i++) {
     130            0 :             src = DeviceMem::create(curUerMemInPtr + (i * rankSize + dstRank) * totalSize, sliceSize);
     131            0 :             dst = DeviceMem::create(static_cast<u8*>(remMemPtr) + (i * rankSize + dstRank) * sliceSize, sliceSize);
     132            0 :             CHK_RET(HcclReduceAsync(
     133              :                 dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, subStream,
     134              :                 static_cast<void*>(dst.ptr()), dstLink->GetRemoteRank(), dstLink->GetLinkType(), INLINE_REDUCE_BIT));
     135              :         }
     136              : 
     137            0 :         CHK_RET(dstLink->TxDataSignal(subStream));
     138            0 :         CHK_RET(dstLink->RxDataSignal(subStream));
     139              :     }
     140              : 
     141            0 :     CHK_RET(SubRecordMain());
     142            0 :     CHK_RET(MainWaitSub());
     143            0 :     CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
     144              : 
     145            0 :     HCCL_INFO("ReduceScatterMeshMix finished: rank[%u]", rank);
     146            0 :     return HCCL_SUCCESS;
     147            0 : }
     148              : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_MESH_MIX, ReduceScatterMeshMix);
     149              : } // namespace hccl