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