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_v1.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 :
16 0 : BroadcastNHRV1::BroadcastNHRV1(const HcclDispatcher dispatcher) : NHRV1Base(dispatcher) {}
17 :
18 0 : BroadcastNHRV1::~BroadcastNHRV1() {}
19 :
20 0 : HcclResult BroadcastNHRV1::Prepare(PrepareData& param)
21 : {
22 0 : return AlgTemplateBase::Prepare(
23 0 : param.inputMem, param.outputMem, param.scratchMem, param.count, param.dataType, param.stream,
24 0 : HCCL_REDUCE_RESERVED, param.root, std::vector<Slice>(ZERO_SLICE), param.baseOffset);
25 : }
26 :
27 0 : HcclResult BroadcastNHRV1::RunScatterOnHorizontal(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
28 : {
29 : // 只有root节点所在的水平Ring做Scatter
30 0 : u32 rootVIndex = info.GetVIndex(root_);
31 0 : u32 vIndex = info.GetVIndex(rank);
32 0 : if (rootVIndex != vIndex) {
33 0 : return HCCL_SUCCESS;
34 : }
35 :
36 : // 收集link
37 0 : u32 hSize = info.GetHSizeByVIndex(vIndex);
38 0 : std::vector<LINK> subLinks(hSize);
39 0 : for (u32 hIdx = 0; hIdx < hSize; hIdx++) {
40 0 : u32 rankInRing = info.GetRank(vIndex, hIdx);
41 0 : CHK_PRT_RET(
42 : rankInRing >= links.size(),
43 : HCCL_ERROR(
44 : "[BroadcastNHRV1][Scatter-H] rank[%u] out of range, "
45 : "rankInRing=%u, links.size=%u",
46 : rank, rankInRing, links.size()),
47 : HCCL_E_INTERNAL);
48 0 : subLinks[hIdx] = links[rankInRing];
49 0 : HCCL_DEBUG("[BroadcastNHRV1][Scatter-H] rank[%u], ringRank[%u]=%u", rank, hIdx, rankInRing);
50 : }
51 :
52 : // 计算新的rank和root
53 0 : u32 subRank = info.GetHIndex(rank);
54 0 : u32 subRoot = info.GetHIndex(root_);
55 0 : HCCL_DEBUG("[BroadcastNHRV1][Scatter-H] rank[%u] subRank=%u, subRoot=%u", rank, subRank, subRoot);
56 :
57 : // 执行Ring - Scatter
58 : // 此处Prepare的baseOffset给0，因为偏移量已经在加在slices里面
59 : std::unique_ptr<AlgTemplateBase> tempAlg
60 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_SCATTER_RING, dispatcher_);
61 0 : CHK_SMART_PTR_NULL(tempAlg);
62 0 : if (!barrierSwitchOn_) {
63 0 : tempAlg->CloseBarrier();
64 : }
65 0 : CHK_RET(tempAlg->Prepare(scratch_, scratch_, scratch_, -1, dataType_, stream_, reductionOp_, subRoot, slices_));
66 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
67 0 : return tempAlg->RunAsync(subRank, subLinks.size(), subLinks);
68 0 : }
69 :
70 0 : HcclResult BroadcastNHRV1::RunBroadcastOnVertical(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
71 : {
72 : // 只有不出现在额外列的节点做Broadcast
73 0 : u32 hIndex = info.GetHIndex(rank);
74 0 : if (hIndex >= info.GetRowSize()) {
75 0 : return HCCL_SUCCESS;
76 : }
77 :
78 : // 收集link
79 0 : u32 vSize = info.GetVSizeByHIndex(hIndex);
80 0 : std::vector<LINK> subLinks(vSize);
81 0 : for (u32 vIdx = 0; vIdx < vSize; vIdx++) {
82 0 : u32 rankInRing = info.GetRank(vIdx, hIndex);
83 0 : CHK_PRT_RET(
84 : rankInRing >= links.size(),
85 : HCCL_ERROR(
86 : "[BroadcastNHRV1][Broadcast-V] rank[%u] out of range, "
87 : "rankInRing=%u, links.size=%u",
88 : rank, rankInRing, links.size()),
89 : HCCL_E_INTERNAL);
90 0 : subLinks[vIdx] = links[rankInRing];
91 0 : HCCL_DEBUG("[BroadcastNHRV1][Broadcast-V] rank[%u], ringRank[%u]=%u", rank, vIdx, rankInRing);
92 : }
93 :
94 : // 计算新的rank和root
95 0 : u32 subRank = info.GetVIndex(rank);
96 0 : u32 subRoot = info.GetVIndex(root_);
97 0 : HCCL_DEBUG("[BroadcastNHRV1][Broadcast-V] rank[%u] subRank=%u, subRoot=%u", rank, subRank, subRoot);
98 :
99 : // 计算新的内存块
100 0 : DeviceMem devMem = scratch_.range(slices_[hIndex].offset, slices_[hIndex].size);
101 0 : u64 memCount = 0;
102 0 : if (DataUnitSize(dataType_) != 0) {
103 0 : memCount = devMem.size() / DataUnitSize(dataType_);
104 : }
105 : // 执行Ring - Broadcast
106 : std::unique_ptr<AlgTemplateBase> tempAlg
107 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_BROADCAST_RING, dispatcher_);
108 0 : CHK_SMART_PTR_NULL(tempAlg);
109 0 : if (!barrierSwitchOn_) {
110 0 : tempAlg->CloseBarrier();
111 : }
112 0 : CHK_RET(tempAlg->Prepare(
113 : devMem, devMem, devMem, memCount, dataType_, stream_, reductionOp_, subRoot, std::vector<Slice>(0),
114 : slices_[hIndex].offset));
115 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
116 0 : return tempAlg->RunAsync(subRank, subLinks.size(), subLinks);
117 0 : }
118 :
119 : HcclResult
120 0 : BroadcastNHRV1::RunAllGatherOnHorizontal(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
121 : {
122 : // 收集link
123 0 : u32 vIndex = info.GetVIndex(rank);
124 0 : u32 hSize = info.GetHSizeByVIndex(vIndex);
125 0 : std::vector<LINK> subLinks(hSize);
126 0 : for (u32 hIdx = 0; hIdx < hSize; hIdx++) {
127 0 : u32 rankInRing = info.GetRank(vIndex, hIdx);
128 0 : CHK_PRT_RET(
129 : rankInRing >= links.size(),
130 : HCCL_ERROR(
131 : "[BroadcastNHRV1][AllGather-H] rank[%u] out of range, "
132 : "rankInRing=%u, links.size=%u",
133 : rank, rankInRing, links.size()),
134 : HCCL_E_INTERNAL);
135 0 : subLinks[hIdx] = links[rankInRing];
136 0 : HCCL_DEBUG("[BroadcastNHRV1][AllGather-H] rank[%u], ringRank[%u]=%u", rank, hIdx, rankInRing);
137 : }
138 :
139 : // 计算新的rank
140 0 : u32 subRank = info.GetHIndex(rank);
141 0 : HCCL_DEBUG("[BroadcastNHRV1][AllGather-H] rank[%u] subRank=%u", rank, subRank);
142 :
143 : // 执行Ring - AllGather
144 : // 此处Prepare的baseOffset给0，因为偏移量已经在加在slices里面
145 : std::unique_ptr<AlgTemplateBase> tempAlg
146 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_RING, dispatcher_);
147 0 : CHK_SMART_PTR_NULL(tempAlg);
148 0 : if (!barrierSwitchOn_) {
149 0 : tempAlg->CloseBarrier();
150 : }
151 0 : CHK_RET(tempAlg->Prepare(scratch_, scratch_, scratch_, -1, dataType_, stream_, reductionOp_, -1, slices_));
152 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
153 0 : return tempAlg->RunAsync(subRank, subLinks.size(), subLinks);
154 0 : }
155 :
156 0 : HcclResult BroadcastNHRV1::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
157 : {
158 : // 基本的检查
159 0 : CHK_RET(SimpleCheck(rank, rankSize, links));
160 0 : HCCL_DEBUG(
161 : "BroadcastNHRV1 run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
162 : inputMem_.ptr(), outputMem_.ptr(), count_);
163 :
164 : // 判断rank_size == 1
165 0 : if (rankSize == 1) {
166 0 : if (inputMem_ != outputMem_) {
167 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
168 : }
169 0 : return HCCL_SUCCESS;
170 : }
171 :
172 : // 创建scratch
173 0 : if (rank == root_) {
174 0 : scratch_ = DeviceMem::create(inputMem_.ptr(), inputMem_.size());
175 : } else {
176 0 : scratch_ = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
177 : }
178 :
179 0 : HCCL_DEBUG("[BroadcastNHRV1] root[%u] scratch[%p] memsize[%llu]", root_, scratch_.ptr(), scratch_.size());
180 :
181 : // 获取通信关系
182 0 : RingInfo info = GetRingInfo(rankSize);
183 :
184 : // 处理和检查Slices
185 0 : if (slices_.size() == 0) {
186 0 : CHK_RET(SetDefaultSlices(rank, info));
187 : }
188 0 : CHK_RET(CheckSlices(rank, info));
189 :
190 : // 水平方向做Ring Scatter（inputMem -> scratch_）
191 0 : CHK_RET(RunScatterOnHorizontal(rank, links, info));
192 :
193 : // 垂直方向做Ring Broadcast（scratch_ -> scratch_）
194 0 : CHK_RET(RunBroadcastOnVertical(rank, links, info));
195 :
196 : // 水平方向做Ring AllGather（scratch_ -> scratch_）
197 0 : CHK_RET(RunAllGatherOnHorizontal(rank, links, info));
198 :
199 0 : HCCL_INFO("BroadcastNHRV1 finished: rank[%u] end", rank);
200 0 : return HCCL_SUCCESS;
201 0 : }
202 :
203 0 : HcclResult BroadcastNHRV1::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
204 : {
205 : // 判断stream, dispatcher是否为空
206 0 : CHK_SMART_PTR_NULL(dispatcher_);
207 0 : CHK_PTR_NULL(stream_.ptr());
208 :
209 : // 判断Memory是否为空
210 0 : if (rank == root_) {
211 0 : CHK_PRT_RET(!inputMem_, HCCL_ERROR("[BroadcastNHRV1]rank[%u] inputmem is null", rank), HCCL_E_PTR);
212 : } else {
213 0 : CHK_PRT_RET(!outputMem_, HCCL_ERROR("[BroadcastNHRV1]rank[%u] outputmem is null", rank), HCCL_E_PTR);
214 : }
215 :
216 : // 判断links数量是否正确
217 0 : CHK_PRT_RET(
218 : links.size() < rankSize,
219 : HCCL_ERROR(
220 : "[BroadcastNHRV1]rank[%u] link size[%llu] is less than "
221 : "rank size[%u]",
222 : rank, links.size(), rankSize),
223 : HCCL_E_INTERNAL);
224 0 : return HCCL_SUCCESS;
225 : }
226 :
227 0 : HcclResult BroadcastNHRV1::SetDefaultSlices(const u32 rank, const RingInfo& info)
228 : {
229 0 : u32 unitSize = DataUnitSize(dataType_);
230 0 : if (unitSize == 0) {
231 0 : HCCL_ERROR("[BroadcastNHRV1] rank[%u] unit data size is zero", rank);
232 0 : return HCCL_E_INTERNAL;
233 : }
234 :
235 : // slices_只用于水平方向的Scatter和AllGather
236 0 : u32 rowSize = info.GetRowSize();
237 0 : u64 sliceCount = (count_ + rowSize - 1) / rowSize;
238 0 : u64 sliceSize = RoundUpWithDivisor(sliceCount * unitSize, HCCL_MIN_SLICE_ALIGN);
239 0 : u64 restSize = count_ * unitSize;
240 0 : slices_.resize(rowSize);
241 0 : for (u32 i = 0; i < rowSize; i++) {
242 : // broadcast逻辑与其他算子不太一样，impl传入的memory是总的大memory而不是预先切出server间的memory，需要在此处做处理
243 0 : slices_[i].offset = (i == 0) ? baseOffset_ : (slices_[i - 1].offset + slices_[i - 1].size);
244 0 : slices_[i].size = std::min(restSize, sliceSize);
245 0 : restSize -= slices_[i].size;
246 0 : HCCL_DEBUG(
247 : "[BroadcastNHRV1] rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu] ", rank, i, slices_[i].offset,
248 : i, slices_[i].size);
249 : }
250 :
251 : // 如果有多余rank，需要添加空白slice以实现BrokenRing
252 0 : if (info.GetHSizeByRank(rank) > info.GetRowSize()) {
253 0 : Slice slice;
254 0 : slice.offset = 0;
255 0 : slice.size = 0;
256 0 : slices_.push_back(slice);
257 0 : HCCL_DEBUG(
258 : "[BroadcastNHRV1] rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu] ", rank, slices_.size() - 1, 0,
259 : slices_.size() - 1, 0);
260 : }
261 0 : return HCCL_SUCCESS;
262 : }
263 :
264 0 : HcclResult BroadcastNHRV1::CheckSlices(const u32 rank, const RingInfo& info)
265 : {
266 0 : u32 expectedSlices = info.GetHSizeByRank(rank);
267 0 : CHK_PRT_RET(
268 : slices_.size() != expectedSlices,
269 : HCCL_ERROR(
270 : "[BroadcastNHRV1]slices.size[%u] should be equal to sqrt of rankSize[%u]", slices_.size(), expectedSlices),
271 : HCCL_E_INTERNAL);
272 :
273 0 : for (u32 idx = 1; idx < slices_.size(); idx++) {
274 0 : if (slices_[idx].size != 0) {
275 0 : CHK_PRT_RET(
276 : slices_[idx - 1].offset + slices_[idx - 1].size != slices_[idx].offset,
277 : HCCL_ERROR(
278 : "[BroadcastNHRV1]only support continuous slices, but get slices[%u].offset[%u]"
279 : ", slices[%u].size[%u], slices[%u].offset[%u]",
280 : idx - 1, slices_[idx - 1].offset, idx - 1, slices_[idx - 1].size, idx, slices_[idx].offset),
281 : HCCL_E_INTERNAL);
282 : }
283 : }
284 0 : return HCCL_SUCCESS;
285 : }
286 :
287 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_BROADCAST_NHR_V1, BroadcastNHRV1);
288 : } // namespace hccl
|
