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 "alg_template_register.h"
      12              : #include "all_reduce_mesh_oneshot.h"
      13              : 
      14              : namespace hccl {
      15            0 : AllReduceMeshDirectOneshot::AllReduceMeshDirectOneshot(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
      16              : 
      17            0 : AllReduceMeshDirectOneshot::~AllReduceMeshDirectOneshot() {}
      18              : 
      19            0 : HcclResult AllReduceMeshDirectOneshot::Prepare(
      20              :     u64 reduceAttrBitMap, std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
      21              :     std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 interRank, u32 interRankSize, u32 userRank,
      22              :     HcomCollOpInfo* opInfo)
      23              : {
      24            0 :     reduceAttr_ = reduceAttrBitMap;
      25            0 :     localRank_ = interRank;
      26            0 :     localRankSize_ = interRankSize;
      27            0 :     userRank_ = userRank;
      28            0 :     meshStreams_ = meshStreams;
      29            0 :     meshSignal_ = &meshSignal;
      30            0 :     meshSignalAux_ = &meshSignalAux;
      31            0 :     opInfo_ = opInfo;
      32            0 :     return HCCL_SUCCESS;
      33              : }
      34              : 
      35              : // ringallreduce算法的函数入口
      36            0 : HcclResult AllReduceMeshDirectOneshot::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
      37              : {
      38            0 :     HcclResult ret = HCCL_SUCCESS;
      39            0 :     CHK_SMART_PTR_NULL(dispatcher_);
      40            0 :     CHK_PTR_NULL(stream_.ptr());
      41            0 :     HCCL_INFO(
      42              :         "AllReduceMeshDirectOneshot run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
      43              :         inputMem_.ptr(), outputMem_.ptr(), count_);
      44              : 
      45            0 :     if (links.size() < rankSize) {
      46            0 :         HCCL_ERROR(
      47              :             "[AllReduceMeshDirectOneshot][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank,
      48              :             links.size(), rankSize);
      49            0 :         return HCCL_E_INTERNAL;
      50              :     }
      51              : 
      52              :     // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
      53            0 :     if (rankSize == 1) {
      54            0 :         if (opInfo_->inputAddr != opInfo_->outputAddr) {
      55            0 :             DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * SIZE_TABLE[dataType_]);
      56            0 :             DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, count_ * SIZE_TABLE[dataType_]);
      57            0 :             ret = HcclD2DMemcpyAsync(dispatcher_, userMemOut, userMemIn, stream_);
      58            0 :             CHK_PRT_RET(
      59              :                 ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceMeshOneshot][RunAsync]rank[%u] memcpy async failed", rank),
      60              :                 ret);
      61            0 :         }
      62            0 :         return ret;
      63              :     }
      64              : 
      65            0 :     ret = RunAllReduceOne(rank, rankSize, links);
      66            0 :     CHK_PRT_RET(
      67              :         ret != HCCL_SUCCESS,
      68              :         HCCL_ERROR(
      69              :             "[AllReduceMeshOneshot][RunAsync]rank[%u] count[%llu] failed"
      70              :             "step",
      71              :             rank, count_),
      72              :         ret);
      73              : 
      74            0 :     HCCL_INFO("AllReduceMeshDirectOneshot finished: rank[%u] ranksize[%u]", rank, rankSize);
      75            0 :     return HCCL_SUCCESS;
      76              : }
      77              : 
      78            0 : HcclResult AllReduceMeshDirectOneshot::RunAllReduceOne(u32 rank, u32 rankSize, const std::vector<LINK>& links)
      79              : {
      80            0 :     HCCL_INFO(
      81              :         "RunAllReduceOne run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
      82              :         inputMem_.ptr(), outputMem_.ptr(), count_);
      83              : 
      84              :     // 数据准备
      85            0 :     u32 unitSize = SIZE_TABLE[dataType_];
      86            0 :     u32 totalSize = unitSize * count_;
      87              : 
      88            0 :     DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, totalSize);
      89            0 :     DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
      90            0 :     DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize);
      91              : 
      92            0 :     DeviceMem src = DeviceMem::create(static_cast<char*>(opInfo_->inputAddr), totalSize);
      93            0 :     DeviceMem dst = DeviceMem::create(static_cast<char*>(opInfo_->outputAddr), totalSize);
      94            0 :     CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
      95              : 
      96            0 :     if (opInfo_->outputAddr != outputMem_.ptr()) {
      97            0 :         src = DeviceMem::create(static_cast<char*>(opInfo_->inputAddr), totalSize);
      98            0 :         dst = commMemOut.range(0, totalSize);
      99            0 :         CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
     100              :     }
     101              : 
     102            0 :     for (u32 round = 1; round < rankSize; round++) {
     103            0 :         u32 dstRank = (round + rank) % rankSize;
     104            0 :         CHK_RET(links[dstRank]->TxAck(stream_));
     105            0 :         CHK_RET(links[dstRank]->RxAck(stream_));
     106              : 
     107            0 :         void* remMemPtr = nullptr;
     108            0 :         CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
     109              : 
     110            0 :         src = DeviceMem::create(static_cast<char*>(remMemPtr), totalSize);
     111            0 :         dst = userMemOut.range(0, totalSize);
     112            0 :         CHK_RET(HcclReduceAsync(
     113              :             dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, stream_,
     114              :             static_cast<void*>(dst.ptr()), links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType(),
     115              :             INLINE_REDUCE_BIT));
     116              : 
     117            0 :         CHK_RET(links[dstRank]->TxDataSignal(stream_));
     118            0 :         CHK_RET(links[dstRank]->RxDataSignal(stream_));
     119              :     }
     120            0 :     return HCCL_SUCCESS;
     121            0 : }
     122              : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_MESH_DIRECT_ONESHOT, AllReduceMeshDirectOneshot);
     123              : } // namespace hccl