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 "coll_scatter_mesh_executor.h"
      12              : 
      13              : namespace hccl {
      14            1 : CollScatterMeshExecutor::CollScatterMeshExecutor(const HcclDispatcher dispatcher,
      15            1 :                                 std::unique_ptr<TopoMatcher> &topoMatcher)
      16            1 :     : CollScatterExecutor(dispatcher, topoMatcher)
      17              : {
      18            4 : }
      19              : 
      20            4 : HcclResult CollScatterMeshExecutor::CalcLevel0CommInfo(TransportMemType inputType,
      21              :     TransportMemType outputType,
      22              :     std::vector<LevelNSubCommTransport>& opTransport)
      23              : {
      24            4 :     HCCL_INFO("[CollScatterMeshExecutor][CalcLevel0CommInfo]tag[%s] start", tag_.c_str());
      25            4 :     CommParaInfo commParaLevel0(COMM_LEVEL0, CommType::COMM_TAG_MESH);
      26              : 
      27            4 :     CHK_RET(CalcCommPlaneInfo(tag_, commParaLevel0, opTransport[COMM_LEVEL0], inputType, outputType));
      28            4 :     HCCL_INFO("[CollScatterMeshExecutor][CalcLevel0CommInfo]tag[%s] Calc meshComm finish", tag_.c_str());
      29            4 :     return HCCL_SUCCESS;
      30            4 : }
      31              : 
      32            4 : HcclResult CollScatterMeshExecutor::CalcStreamNum(u32& streamNum)
      33              : {
      34            4 :     u32 totalStreamNum = topoAttr_.deviceNumPerAggregation > 1U ? topoAttr_.deviceNumPerAggregation - 1U : 1U;
      35            4 :     streamNum = totalStreamNum - 1U;
      36            4 :     return HCCL_SUCCESS;
      37              : }
      38              : 
      39            0 : HcclResult CollScatterMeshExecutor::KernelRun(const OpParam &param, ExecMem &execMem)
      40              : {
      41            0 :     HCCL_CONFIG_INFO(HCCL_ALG, "[CollScatterMeshExecutor] scatter starts.");
      42            0 :     Stream& stream = const_cast<Stream&>(param.stream);
      43              : 
      44            0 :     u32 perDataSize = SIZE_TABLE[param.DataDes.dataType];
      45              : 
      46            0 :     SubCommInfo level0CommInfo = GetSubCommInfo(COMM_LEVEL0, COMM_INDEX_0);
      47            0 :     u32 level0LocalRank = level0CommInfo.localRank;
      48            0 :     u32 level0LocalRankSize = level0CommInfo.localRankSize;
      49              : 
      50            0 :     u32 commIndex = level0LocalRank;
      51              : 
      52            0 :     CHK_RET(CheckCommSize(COMM_LEVEL1, commIndex + 1));
      53            0 :     SubCommInfo level1CommInfo = GetSubCommInfo(COMM_LEVEL1, commIndex);
      54            0 :     u32 level1LocalRank = level1CommInfo.localRank;
      55            0 :     u32 level1LocalRankSize = level1CommInfo.localRankSize;
      56              : 
      57            0 :     bool bRet = level0LocalRankSize == 0;
      58            0 :     CHK_PRT_RET(bRet, HCCL_ERROR("[CollScatterMeshExecutor][KernelRun]tag[%s],comm level0 is empty", tag_.c_str()),
      59              :         HCCL_E_INTERNAL);
      60              : 
      61              :     /* ***********第一步: 节点间scatter ****************************/
      62            0 :     u32 subRoot = INVALID_VALUE_RANKID;
      63            0 :     CHK_RET(topoMatcher_->GetSubRootForScatter(param.root, subRoot));
      64            0 :     CHK_PRT_RET(subRoot == INVALID_VALUE_RANKID,
      65              :         HCCL_ERROR("[CollScatterMeshExecutor][KernelRun]GetSubRootForScatter failed, "\
      66              :         "userRank[%u], root[%u], subRoot[%u]", topoAttr_.userRank, param.root, subRoot), HCCL_E_INTERNAL);
      67            0 :     HCCL_DEBUG("[CollScatterMeshExecutor][KernelRun]GetSubRootForScatter, userRank[%u], root[%u], subRoot[%u]",
      68              :         topoAttr_.userRank, param.root, subRoot);
      69            0 :     CHK_RET(KernelRunLevel1(execMem.inputMem, execMem.count, param.DataDes.dataType, commIndex,
      70              :         param.root, subRoot, COMM_LEVEL1, stream));
      71              : 
      72              :     /* ***********第二步: 节点内scatter*****************************/
      73              :     // 根据数据量算每个环上数据的偏移和大小
      74            0 :     u32 sliceNum = level0LocalRankSize;
      75            0 :     std::vector<Slice> dataSegsSlice;
      76            0 :     CHK_RET(PrepareDataSlice(execMem.count, perDataSize, sliceNum, dataSegsSlice));
      77              : 
      78              :     // 每个server分配的slice大小
      79            0 :     u64 serverSliceSize = execMem.inputMem.size() / level1LocalRankSize;
      80              :     // 每个服务器对应的偏移
      81            0 :     u64 serverSliceOffset = serverSliceSize * level1LocalRank;
      82            0 :     DeviceMem scatterMeshInput = execMem.inputMem.range(serverSliceOffset, serverSliceSize);
      83            0 :     CHK_SMART_PTR_NULL(scatterMeshInput);
      84            0 :     DeviceMem scatterMeshOutput = execMem.inputMem.range(serverSliceOffset, serverSliceSize);
      85            0 :     CHK_SMART_PTR_NULL(scatterMeshOutput);
      86              : 
      87            0 :     std::unique_ptr<AlgTemplateBase> level0TempAlg = AlgTemplateRegistry::Instance().GetAlgTemplate(
      88            0 :         TemplateType::TEMPLATE_SCATTER_MESH, dispatcher_);
      89            0 :     CHK_SMART_PTR_NULL(level0TempAlg);
      90            0 :     CHK_RET(level0TempAlg->Prepare(level0LocalRank, level0LocalRankSize));
      91              :     // 偏移需要带入prepare
      92            0 :     u32 rootRankLevel0 = 0;
      93            0 :     CHK_RET(GetRankByUserRank(COMM_LEVEL0, COMM_INDEX_0, subRoot, rootRankLevel0));
      94            0 :     CHK_PRT_RET(rootRankLevel0 == INVALID_VALUE_RANKID,
      95              :         HCCL_ERROR("[CollScatterMeshExecutor][KernelRun]rootRankLevel0[%u] is invalid, userRank[%u], subRoot[%u]",
      96              :         rootRankLevel0, topoAttr_.userRank, subRoot), HCCL_E_INTERNAL);
      97              : 
      98            0 :     CHK_RET(level0TempAlg->Prepare(scatterMeshInput, scatterMeshOutput, execMem.inputMem, execMem.count,
      99              :         param.DataDes.dataType, stream, HCCL_REDUCE_RESERVED, rootRankLevel0, dataSegsSlice, serverSliceOffset));
     100              : 
     101            0 :     HcclResult ret = RunTemplate(level0TempAlg, level0CommInfo);
     102            0 :     CHK_PRT_RET(ret != HCCL_SUCCESS,
     103              :         HCCL_ERROR("[CollScatterMeshExecutor][KernelRun]scatter(mesh) RunTemplate failed,return[%d]", ret), ret);
     104              : 
     105              :     // 将scratchMem赋值给outputMem
     106            0 :     u8 *scatterMeshOutputPtr = static_cast<u8 *>(scatterMeshOutput.ptr());
     107            0 :     DeviceMem resultMem(scatterMeshOutputPtr + execMem.outputMem.size() * level0LocalRank, execMem.outputMem.size());
     108            0 :     CHK_RET(HcclD2DMemcpyAsync(dispatcher_, execMem.outputMem, resultMem, stream));
     109            0 :     return HCCL_SUCCESS;
     110            0 : }
     111              : 
     112              : REGISTER_EXEC("ScatterMeshExecutor", ScatterMesh, CollScatterMeshExecutor);
     113              : 
     114              : }