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_reduce_nhr_oneshot.h"
      12              : #include "alg_template_register.h"
      13              : #include "reduce_nhr_oneshot.h"
      14              : #include "broadcast_nhr_oneshot.h"
      15              : 
      16              : namespace hccl {
      17           10 : AllReduceNHROneshot::AllReduceNHROneshot(const HcclDispatcher dispatcher) : NHRBase(dispatcher)
      18              : {
      19           10 : }
      20              : 
      21           20 : AllReduceNHROneshot::~AllReduceNHROneshot()
      22              : {
      23           20 : }
      24              : 
      25           10 : HcclResult AllReduceNHROneshot::Prepare(u64 reduceAttrBitMap, HcomCollOpInfo *opInfo)
      26              : {
      27           10 :     reduceAttr_ = reduceAttrBitMap;
      28           10 :     return HCCL_SUCCESS;
      29              : }
      30              : 
      31            3 : HcclResult AllReduceNHROneshot::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
      32              : {
      33              :     // 基本的检查
      34            3 :     CHK_RET(SimpleCheck(rank, rankSize, links));
      35            3 :     HCCL_INFO("[AllReduceNHROneshot][RunAsync] run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]",
      36              :         rank, rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
      37              : 
      38            3 :     HcclResult ret = HCCL_SUCCESS;
      39              :     // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
      40            3 :     if (rankSize == 1) {
      41            0 :         if (inputMem_ != outputMem_) {
      42            0 :             ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
      43            0 :             CHK_PRT_RET(ret != HCCL_SUCCESS,
      44              :                 HCCL_ERROR("[AllReduceNHROneshot][RunAsync] rank[%u] memcpy async failed", rank), ret);
      45              :         }
      46              : 
      47            0 :         return ret;
      48              :     }
      49              : 
      50              :     // 先执行1-reduce
      51            3 :     ret = RunReduceOneshot(rank, rankSize, links);
      52            3 :     CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceNHROneshot][RunAsync] rank[%u] count[%llu] failed in "
      53              :         "1-reduce step", rank, count_), ret);
      54              : 
      55              :     // 再执行1-bcast
      56            3 :     ret = RunBroadcastOneshot(rank, rankSize, links);
      57            3 :     CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceNHROneshot][RunAsync] rank[%u] count[%llu] failed in "
      58              :         "1-bcast step", rank, count_), ret);
      59              : 
      60            3 :     HCCL_INFO("[AllReduceNHROneshot][RunAsync] finished: rank[%u] ranksize[%u]", rank, rankSize);
      61            3 :     return HCCL_SUCCESS;
      62              : }
      63              : 
      64            3 : HcclResult AllReduceNHROneshot::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
      65              : {
      66              :     // 判断stream, dispatcher是否为空
      67            3 :     CHK_SMART_PTR_NULL(dispatcher_);
      68            3 :     CHK_PTR_NULL(stream_.ptr());
      69              : 
      70              :     // 检查memory
      71            3 :     CHK_PRT_RET(!outputMem_ || !inputMem_,
      72              :         HCCL_ERROR("[AllReduceNHROneshot][SimpleCheck] rank[%u] inputmem or outputmem is null", rank), HCCL_E_PTR);
      73              : 
      74              :     // 判断links数量是否正确
      75            3 :     CHK_PRT_RET(links.size() < rankSize, HCCL_ERROR("[AllReduceNHROneshot][SimpleCheck] rank[%u] link size[%llu] is "
      76              :         "less than rank size[%u]", rank, links.size(), rankSize), HCCL_E_INTERNAL);
      77              : 
      78            3 :     return HCCL_SUCCESS;
      79              : }
      80              : 
      81            3 : HcclResult AllReduceNHROneshot::RunReduceOneshot(u32 rank, u32 rankSize, const std::vector<LINK> &links)
      82              : {
      83            3 :     std::unique_ptr<AlgTemplateBase> tempAlg;
      84            3 :     tempAlg = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCE_NHR_ONE_SHOT, dispatcher_);
      85            3 :     CHK_SMART_PTR_NULL(tempAlg);
      86            3 :     CHK_RET(tempAlg->Prepare(reduceAttr_));
      87            3 :     if (!barrierSwitchOn_) {
      88            0 :         tempAlg->CloseBarrier();
      89              :     }
      90            3 :     HCCL_INFO("[AllReduceNHROneshot][RunReduceOneshot] 1-reduce tempAlg rank[%u] inputMem[%p] outputMem[%p] "
      91              :         "mem_size[%llu] count[%llu] planeID:[%d]",
      92              :         rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
      93            9 :     CHK_RET(tempAlg->Prepare(inputMem_, inputMem_, outputMem_, count_, dataType_, stream_,
      94              :         reductionOp_, root_, slices_, baseOffset_));
      95              : 
      96            3 :     CHK_RET(tempAlg->RegisterProfiler(
      97              :         profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
      98              : 
      99            3 :     return tempAlg->RunAsync(rank, rankSize, links);
     100            3 : }
     101              : 
     102            3 : HcclResult AllReduceNHROneshot::RunBroadcastOneshot(u32 rank, u32 rankSize, const std::vector<LINK> &links)
     103              : {
     104            3 :     BroadcastNHROneshot tempAlg(dispatcher_);
     105            3 :     HCCL_INFO("[AllReduceNHROneshot][RunBroadcastOneshot] 1-broadcast tempAlg rank[%u] inputMem[%p] outputMem[%p] "
     106              :         "mem_size[%llu] count[%llu] planeID:[%d]", rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(),
     107              :         count_, profilerInput_.planeID);
     108            9 :     CHK_RET(tempAlg.Prepare(inputMem_, outputMem_, outputMem_, count_, dataType_, stream_,
     109              :         reductionOp_, root_, slices_, baseOffset_));
     110              : 
     111            3 :     CHK_RET(tempAlg.RegisterProfiler(
     112              :         profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
     113              : 
     114            3 :     return tempAlg.RunAsyncForAllReduce(rank, rankSize, links);
     115            3 : }
     116              : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_NHR_ONESHOT, AllReduceNHROneshot);
     117              : }  // namespace hccl