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 "broadcast_nhr.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : BroadcastNHR::BroadcastNHR(const HcclDispatcher dispatcher)
16 0 : : NHRBase(dispatcher)
17 : {
18 0 : }
19 :
20 0 : BroadcastNHR::~BroadcastNHR()
21 : {
22 0 : }
23 :
24 0 : HcclResult BroadcastNHR::RunAsync(const u32 rank, const u32 rankSize,
25 : const std::vector<std::shared_ptr<Transport> > &links)
26 : {
27 0 : CHK_SMART_PTR_NULL(dispatcher_);
28 0 : CHK_PTR_NULL(stream_.ptr());
29 0 : HCCL_INFO("[BroadcastNHR][RunAsync] run: rank[%u] totalrank[%u] count[%llu]", rank, rankSize, count_);
30 :
31 0 : if (rankSize == 1) {
32 0 : return HCCL_SUCCESS;
33 : }
34 :
35 0 : CHK_PRT_RET(links.size() < rankSize,
36 : HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] linksize[%llu] is less than rank size", rank, links.size()),
37 : HCCL_E_INTERNAL);
38 :
39 0 : u32 unitSize = DataUnitSize(dataType_);
40 0 : CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] unit data size is zero", rank),
41 : HCCL_E_INTERNAL);
42 :
43 0 : DeviceMem srcMem = inputMem_.range(baseOffset_, count_ * unitSize);
44 0 : DeviceMem dstMem = outputMem_.range(baseOffset_, count_ * unitSize);
45 :
46 0 : if (inputMem_ != outputMem_ && rank == root_) {
47 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, stream_));
48 : }
49 :
50 : // 准备slice
51 0 : PrepareSlice(rank, rankSize);
52 :
53 : // scatter
54 0 : std::unique_ptr<AlgTemplateBase> tempAlgScatter = AlgTemplateRegistry::Instance().GetAlgTemplate(
55 0 : TemplateType::TEMPLATE_SCATTER_NHR, dispatcher_);
56 0 : CHK_SMART_PTR_NULL(tempAlgScatter);
57 0 : CHK_RET(tempAlgScatter->Prepare(true));
58 0 : tempAlgScatter->CloseBarrier();
59 0 : CHK_RET(tempAlgScatter->Prepare(srcMem, srcMem, srcMem, count_, dataType_, stream_,
60 : reductionOp_, root_, slices_, baseOffset_));
61 :
62 0 : CHK_RET(tempAlgScatter->RegisterProfiler(
63 : profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
64 0 : CHK_RET(tempAlgScatter->RunAsync(rank, rankSize, links));
65 :
66 : // allgather
67 0 : std::unique_ptr<AlgTemplateBase> tempAlgAllgather = AlgTemplateRegistry::Instance().GetAlgTemplate(
68 0 : TemplateType::TEMPLATE_ALL_GATHER_NHR, dispatcher_);
69 0 : CHK_SMART_PTR_NULL(tempAlgAllgather);
70 0 : CHK_RET(tempAlgAllgather->Prepare(true));
71 0 : CHK_RET(tempAlgAllgather->Prepare(srcMem, srcMem, srcMem, count_, dataType_, stream_,
72 : reductionOp_, root_, slices_, baseOffset_));
73 :
74 0 : CHK_RET(tempAlgAllgather->RegisterProfiler(
75 : profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
76 0 : CHK_RET(tempAlgAllgather->RunAsync(rank, rankSize, links));
77 :
78 0 : HCCL_INFO("[BroadcastNHR][RunAsync] finished: rank[%u] end count[%llu]", rank, count_);
79 0 : return HCCL_SUCCESS;
80 0 : }
81 :
82 0 : HcclResult BroadcastNHR::PrepareSlice(const u32 rank, const u32 rankSize)
83 : {
84 0 : u32 unitSize = DataUnitSize(dataType_);
85 :
86 : // 所有的数据平均到每个rank上
87 0 : u64 sizeAvg = ((count_ + rankSize - 1) / rankSize) * unitSize;
88 0 : u64 sizePerSlice = AlgTemplateBase::RoundUpWithDivisor(sizeAvg, HCCL_MIN_SLICE_ALIGN);
89 0 : HCCL_DEBUG("[BroadcastNHR][PrepareSlice] bcast total count[%llu] sizeAverage[%llu] sizePerSlice after aligns[%llu]",
90 : count_, sizeAvg, sizePerSlice);
91 :
92 : // 准备slice
93 0 : slices_.resize(rankSize);
94 0 : u64 sizeResidue = count_ * unitSize;
95 0 : u64 sizePerRound = 0;
96 :
97 0 : for (u32 i = 0; i < rankSize; i++) {
98 0 : sizePerRound = (sizeResidue > sizePerSlice) ? sizePerSlice : sizeResidue;
99 0 : slices_[i].offset = count_ * unitSize - sizeResidue;
100 0 : slices_[i].size = sizePerRound;
101 :
102 0 : sizeResidue -= sizePerRound;
103 0 : HCCL_DEBUG("[BroadcastNHR][PrepareSlice] rank[%u] default slice[%u]: offset: [%llu] size[%llu]", rank, i,
104 : slices_[i].offset, slices_[i].size);
105 : }
106 0 : return HCCL_SUCCESS;
107 : }
108 :
109 0 : HcclResult BroadcastNHR::GetNslbAdjInfo(const u32 rank, const u32 rankSize,
110 : const std::vector<LINK> &links, AdjInfo& nslbAdjInfo)
111 : {
112 0 : if (rankSize == 1) {
113 0 : return HCCL_SUCCESS;
114 : }
115 0 : if (links.size() < rankSize) {
116 0 : return HCCL_SUCCESS;
117 : }
118 0 : u32 nSteps = 0;
119 0 : for(u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps){}
120 :
121 0 : u32 deltaRoot = (rankSize - rank) % rankSize;
122 : // 先执行 scatter 流程
123 0 : for (u32 step = 0; step < nSteps; step++) {
124 0 : u32 deltaRankPair = 1 << step;
125 0 : u32 nRanks = 0;
126 0 : bool isPerfect = (rankSize & (rankSize - 1)) == 0;
127 0 : if (!isPerfect && step == nSteps - 1) {
128 0 : nRanks = rankSize - deltaRankPair;
129 : } else {
130 0 : nRanks = deltaRankPair;
131 : }
132 0 : if (deltaRoot >= nRanks) {
133 0 : continue;
134 : }
135 0 : u32 sendTo =(rank + rankSize - deltaRankPair) % rankSize;
136 0 : LINK linkRight = links[sendTo];
137 0 : CHK_SMART_PTR_NULL(linkRight);
138 :
139 0 : NslbDpAdjInfo adjInfoStep = {0};
140 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
141 0 : adjInfoStep.phaseId = step + 1;
142 0 : adjInfoStep.rev = 0;
143 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
144 0 : }
145 0 : u32 begin = nSteps;
146 : //后续执行AllGather的NHR流程
147 0 : for (u32 step = 0; step < nSteps; step++) {
148 0 : u32 deltaRank = 1 << (nSteps - 1 - step);
149 0 : u32 sendTo =(rank + deltaRank) % rankSize;
150 0 : LINK linkRight = links[sendTo];
151 0 : CHK_SMART_PTR_NULL(linkRight);
152 :
153 0 : NslbDpAdjInfo allGatherInfoStep = {0};
154 0 : allGatherInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
155 0 : allGatherInfoStep.phaseId = step + begin + 1;
156 0 : allGatherInfoStep.rev = 0;
157 0 : nslbAdjInfo.nsAdjInfo.push_back(allGatherInfoStep);
158 0 : }
159 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
160 0 : return HCCL_SUCCESS;
161 : }
162 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_BROADCAST_NHR, BroadcastNHR);
163 : } // namespace hccl
|
