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