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 <cmath>
12 : #include "all_gather_hd_stage_pub.h"
13 :
14 : namespace hccl {
15 : namespace {
16 0 : static u32 GetStepNumInterServer(u32 rankSize)
17 : {
18 0 : u32 nSteps = 0;
19 0 : for (u32 tmp = rankSize - 1; tmp != 0; tmp >>= 1, nSteps++) {
20 : }
21 0 : return nSteps;
22 : }
23 :
24 0 : static void ReorderSequence(u32 start, u32 end, u32 len, std::vector<u32>& tree, std::vector<u32>& tmp)
25 : {
26 0 : const u32 divideTwo = 2;
27 :
28 0 : for (u32 i = start; i < end; i++) {
29 0 : u32 offset = i - start;
30 0 : if ((offset & 1) == 0) {
31 0 : tmp[start + offset / divideTwo] = tree[i];
32 : } else {
33 0 : tmp[start + (offset + len) / divideTwo] = tree[i];
34 : }
35 : }
36 0 : }
37 :
38 : // 参考NHRBase::GetRankMapping的实现
39 0 : static void GetRankMapping(const u32 rankSize, std::vector<u32>& sliceMap)
40 : {
41 0 : std::vector<u32> tree;
42 0 : for (u32 i = 0; i < rankSize; i++) {
43 0 : tree.push_back(i);
44 : }
45 :
46 : // 其他的再进行计算
47 0 : std::vector<u32> tmp(rankSize);
48 0 : u32 nSteps = GetStepNumInterServer(rankSize);
49 0 : u32 len = rankSize;
50 :
51 0 : for (u32 step = 0; step < nSteps; step++) {
52 0 : u32 nSlices = (rankSize - 1 + (1 << step)) / (1 << (step + 1));
53 0 : if (nSlices <= 1) {
54 0 : break;
55 : }
56 :
57 0 : bool endFlag = false;
58 0 : for (u32 part = 0; part * len < rankSize; part++) {
59 0 : u32 start = part * len;
60 0 : u32 end = std::min(start + len, rankSize);
61 0 : ReorderSequence(start, end, len, tree, tmp);
62 :
63 0 : if (((end - start) & 1) == 1) {
64 0 : endFlag = true;
65 : }
66 : }
67 :
68 0 : for (u32 i = 0; i < rankSize; i++) {
69 0 : tree[i] = tmp[i];
70 : }
71 :
72 0 : if (endFlag) {
73 0 : break;
74 : }
75 :
76 0 : len >>= 1;
77 : }
78 :
79 : // 因为取的是tree中rank的idx，所以直接返回反向的映射
80 0 : sliceMap.resize(rankSize);
81 0 : for (u32 i = 0; i < rankSize; i++) {
82 0 : sliceMap[tree[i]] = i;
83 : }
84 0 : return;
85 0 : }
86 : } // namespace
87 :
88 0 : AllGatherHDStage::AllGatherHDStage(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
89 :
90 0 : AllGatherHDStage::~AllGatherHDStage() {}
91 :
92 0 : HcclResult AllGatherHDStage::Prepare(PrepareData& param)
93 : {
94 0 : userRank_ = param.userRank;
95 0 : opInfo_ = param.opInfo;
96 :
97 0 : meshStreams_ = *param.subStreamsPtr;
98 0 : meshSignalPtr_ = param.signalPtr;
99 0 : meshSignalAuxPtr_ = param.signalAuxPtr;
100 :
101 0 : return AlgTemplateBase::Prepare(
102 0 : param.inputMem, param.outputMem, param.scratchMem, param.count, param.dataType, param.stream,
103 0 : HCCL_REDUCE_RESERVED, INVALID_VALUE_RANKID);
104 : }
105 :
106 0 : HcclResult AllGatherHDStage::MainRecordSub(u32 streamNum)
107 : {
108 0 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux = *meshSignalAuxPtr_;
109 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
110 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, meshSignalAux[signalIndex], profilerInput_.stage));
111 : }
112 0 : return HCCL_SUCCESS;
113 : }
114 :
115 0 : HcclResult AllGatherHDStage::SubWaitMain(u32 streamNum)
116 : {
117 0 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux = *meshSignalAuxPtr_;
118 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
119 0 : CHK_RET(LocalNotify::Wait(
120 : meshStreams_[streamIndex], dispatcher_, meshSignalAux[streamIndex], profilerInput_.stage));
121 : }
122 0 : return HCCL_SUCCESS;
123 : }
124 :
125 0 : HcclResult AllGatherHDStage::MainWaitSub(u32 streamNum)
126 : {
127 0 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignal = *meshSignalPtr_;
128 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
129 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, meshSignal[signalIndex], profilerInput_.stage));
130 : }
131 0 : return HCCL_SUCCESS;
132 : }
133 :
134 0 : HcclResult AllGatherHDStage::SubRecordMain(u32 streamNum)
135 : {
136 0 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignal = *meshSignalPtr_;
137 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
138 0 : CHK_RET(
139 : LocalNotify::Post(meshStreams_[streamIndex], dispatcher_, meshSignal[streamIndex], profilerInput_.stage));
140 : }
141 0 : return HCCL_SUCCESS;
142 : }
143 :
144 : // ringallreduce算法的函数入口
145 0 : HcclResult AllGatherHDStage::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
146 : {
147 0 : HcclResult ret = HCCL_SUCCESS;
148 0 : CHK_SMART_PTR_NULL(dispatcher_);
149 0 : CHK_PTR_NULL(stream_.ptr());
150 0 : HCCL_INFO(
151 : "AllGatherHDStage run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
152 : inputMem_.ptr(), outputMem_.ptr(), count_);
153 :
154 0 : if (links.size() < rankSize) {
155 0 : HCCL_ERROR(
156 : "[AllGatherHDStage][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(),
157 : rankSize);
158 0 : return HCCL_E_INTERNAL;
159 : }
160 :
161 0 : ret = RunAllGatherStage(rank, rankSize, links);
162 0 : CHK_PRT_RET(
163 : ret != HCCL_SUCCESS,
164 : HCCL_ERROR(
165 : "[AllGatherHDStage][RunAsync]rank[%u] count[%llu] failed"
166 : "step",
167 : rank, count_),
168 : ret);
169 :
170 0 : HCCL_INFO("AllGatherHDStage finished: rank[%u] ranksize[%u]", rank, rankSize);
171 0 : return HCCL_SUCCESS;
172 : }
173 :
174 0 : HcclResult AllGatherHDStage::ReverseId(u32 oriIdx, u32& revIdx)
175 : {
176 0 : revIdx = 0;
177 0 : u32 powerBase = 0;
178 0 : for (u32 i = 0; i < powerSteps_; i++) {
179 0 : powerBase = static_cast<u32>(pow(base, i));
180 0 : revIdx += (oriIdx / powerBase % base) * static_cast<u32>(pow(base, powerSteps_ - i - 1));
181 : }
182 0 : return HCCL_SUCCESS;
183 : }
184 :
185 : HcclResult
186 0 : AllGatherHDStage::PrepareSliceData(u32 subRank, u32 subRankSize, u32 size, u32 batchSize, std::vector<Slice>& slices)
187 : {
188 0 : Slice temp;
189 0 : u32 power = static_cast<u32>(log2(subRankSize));
190 0 : slices.clear();
191 0 : slices.reserve(power);
192 0 : for (u32 step = 0; step < power; step++) {
193 0 : u32 sliceNum = pow(base, step);
194 0 : u32 offset = static_cast<u32>(subRank ^ (1 << step)) / sliceNum * sliceNum;
195 0 : temp.offset = (offset * size) % batchSize;
196 0 : temp.size = sliceNum * size;
197 0 : slices.push_back(temp);
198 : }
199 0 : return HCCL_SUCCESS;
200 : }
201 :
202 0 : HcclResult AllGatherHDStage::RunAllGatherStage(u32 rank, u32 rankSize, const std::vector<LINK>& links)
203 : {
204 0 : HCCL_INFO(
205 : "RunAllGatherStage run: rank[%u] totalrank[%u] outputMem[%p] count[%llu]", rank, rankSize, outputMem_.ptr(),
206 : count_);
207 0 : u32 unitSize = SIZE_TABLE[dataType_];
208 0 : totalSize_ = unitSize * count_;
209 : // 对应因式分解中的2的幂次部分
210 0 : powerSteps_ = static_cast<u32>(log2(rankSize & (-rankSize)));
211 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
212 0 : if (powerSteps_ >= base) {
213 0 : finalSteps_ = base;
214 0 : } else if (powerSteps_ >= 1) {
215 0 : finalSteps_ = 1;
216 : }
217 : }
218 : // 对应因式分解中的奇数部分
219 0 : noPower_ = rankSize / (rankSize & (-rankSize));
220 0 : CHK_RET(RunPreCopy(rank, rankSize, links));
221 0 : if (noPower_ > 1) {
222 0 : CHK_RET(RunAllGatherNoPower(rank, rankSize, links));
223 : }
224 0 : if ((powerSteps_ - finalSteps_) >= 1) {
225 0 : CHK_RET(RunAllGatherPower(rank, rankSize, links));
226 : }
227 0 : if (finalSteps_ == base) {
228 0 : CHK_RET(RunAllGatherLastTwo(rank, rankSize, links));
229 0 : } else if (finalSteps_ == 1) {
230 0 : CHK_RET(RunAllGatherLastOne(rank, rankSize, links));
231 : } else {
232 0 : CHK_RET(RunAllGatherLast(rank, rankSize, links));
233 : }
234 0 : return HCCL_SUCCESS;
235 : }
236 :
237 0 : HcclResult AllGatherHDStage::RunPreCopy(u32 rank, u32 rankSize, const std::vector<LINK>& links)
238 : {
239 : // 交换数据
240 0 : HCCL_INFO(
241 : "RunPreCopy run: rank[%u] totalrank[%u] outputMem[%p] count[%llu]", rank, rankSize, outputMem_.ptr(), count_);
242 0 : std::vector<u32> noPowerMap;
243 0 : GetRankMapping(noPower_, noPowerMap);
244 0 : std::vector<u32> noPowerRevMap(noPower_);
245 0 : CHK_PRT_RET(
246 : noPowerMap.size() != noPower_,
247 : HCCL_ERROR("[AllGatherHDStage][RunPreCopy]rank[%u] count[%llu] failed", rank, count_), HCCL_E_RESERVED);
248 0 : for (u32 i = 0; i < noPowerMap.size(); i++) {
249 0 : noPowerRevMap[noPowerMap[i]] = i;
250 : }
251 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
252 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
253 0 : u32 revRank = 0;
254 0 : u32 revIdx = 0;
255 0 : CHK_RET(ReverseId(groupIdx, revIdx));
256 : // 将revRank的数据写到本卡
257 0 : revRank = revIdx * noPower_ + noPowerMap[group];
258 : // 将本卡数据写到revRankrev卡
259 0 : groupIdx = rank % noPower_;
260 0 : group = rank / noPower_;
261 0 : CHK_RET(ReverseId(group, revIdx));
262 0 : u32 revRankrev = noPowerRevMap[groupIdx] * static_cast<u32>(pow(base, powerSteps_)) + revIdx;
263 0 : DeviceMem UserMemIn = DeviceMem::create(opInfo_->inputAddr, totalSize_);
264 0 : if (revRank != rank && revRankrev != rank) {
265 0 : CHK_RET(links[revRank]->TxAck(stream_));
266 0 : CHK_RET(links[revRankrev]->RxAck(stream_));
267 0 : void* remMemPtr = nullptr;
268 0 : CHK_RET(links[revRankrev]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
269 : DeviceMem dst = DeviceMem::create(
270 0 : static_cast<u8*>(remMemPtr) + (rank % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_,
271 0 : totalSize_);
272 0 : DeviceMem src = UserMemIn;
273 0 : CHK_RET(HcclD2DMemcpyAsync(
274 : dispatcher_, dst, src, stream_, links[revRankrev]->GetRemoteRank(), links[revRankrev]->GetLinkType()));
275 0 : CHK_RET(links[revRankrev]->TxDataSignal(stream_));
276 0 : CHK_RET(links[revRank]->RxDataSignal(stream_));
277 0 : } else {
278 : DeviceMem dst
279 0 : = outputMem_.range((rank % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_, totalSize_);
280 0 : DeviceMem src = UserMemIn;
281 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
282 0 : }
283 0 : return HCCL_SUCCESS;
284 0 : }
285 :
286 0 : HcclResult AllGatherHDStage::RunAllGatherNoPower(u32 rank, u32 rankSize, const std::vector<LINK>& links)
287 : {
288 : std::unique_ptr<AlgTemplateBase> tempAlg
289 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_NHR, dispatcher_);
290 0 : CHK_SMART_PTR_NULL(tempAlg);
291 0 : CHK_RET(tempAlg->Prepare(true));
292 0 : HCCL_INFO(
293 : "[AllGatherHDStage][RunAllGather] rank[%u] tempAlg AllGatherNHR inputMem[%p] outputMem[%p] mem_size[%llu] "
294 : "count[%llu] planeID:[%d]",
295 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
296 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
297 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
298 0 : u32 revIdx = 0;
299 0 : CHK_RET(ReverseId(groupIdx, revIdx));
300 0 : u64 baseOffset = ((revIdx * noPower_) % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_;
301 0 : std::vector<Slice> slices;
302 0 : for (u32 i = 0; i < noPower_; i++) {
303 0 : Slice temp;
304 0 : temp.offset = i * totalSize_;
305 0 : temp.size = totalSize_;
306 0 : slices.push_back(temp);
307 : }
308 0 : DeviceMem nhrOutput = outputMem_.range(baseOffset, outputMem_.size() - baseOffset);
309 0 : CHK_RET(tempAlg->Prepare(
310 : nhrOutput, nhrOutput, nhrOutput, count_, dataType_, stream_, reductionOp_, 0, slices, baseOffset));
311 :
312 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
313 :
314 0 : std::vector<LINK> nhrLinks;
315 0 : for (u32 i = 0; i < noPower_; i++) {
316 0 : u32 remote = i * static_cast<u32>(pow(base, powerSteps_)) + groupIdx;
317 0 : nhrLinks.push_back(links[remote]);
318 : }
319 0 : return tempAlg->RunAsync(group, noPower_, nhrLinks);
320 0 : }
321 :
322 0 : HcclResult AllGatherHDStage::RunBetweenStep(u32 rank, u32 neighCur, u32 neighNext, const std::vector<LINK>& links)
323 : {
324 : (void)rank;
325 0 : CHK_RET(MainRecordSub(1));
326 0 : CHK_RET(SubWaitMain(1));
327 :
328 0 : CHK_RET(links[neighCur]->TxDataSignal(meshStreams_[0]));
329 0 : CHK_RET(links[neighCur]->RxDataSignal(meshStreams_[0]));
330 :
331 0 : CHK_RET(links[neighNext]->TxAck(stream_));
332 0 : CHK_RET(links[neighNext]->RxAck(stream_));
333 :
334 0 : CHK_RET(SubRecordMain(1));
335 0 : CHK_RET(MainWaitSub(1));
336 :
337 0 : return HCCL_SUCCESS;
338 : }
339 :
340 0 : HcclResult AllGatherHDStage::RunAllGatherPower(u32 rank, u32 rankSize, const std::vector<LINK>& links)
341 : {
342 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
343 :
344 0 : void* remMemPtr = nullptr;
345 0 : DeviceMem dst;
346 0 : DeviceMem src;
347 : u32 dstRank;
348 : u32 dstGroupIdx;
349 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
350 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
351 0 : u32 revGroup = 0;
352 0 : CHK_RET(ReverseId(groupIdx, revGroup));
353 0 : CHK_RET(PrepareSliceData(
354 : revGroup, pow(base, powerSteps_), noPower_ * totalSize_,
355 : totalSize_ * rankSize / static_cast<u32>(pow(base, finalSteps_)), slicePower_));
356 0 : for (u32 step = 0; step < powerSteps_ - finalSteps_; step++) {
357 0 : dstGroupIdx = groupIdx ^ (1 << (powerSteps_ - 1 - step));
358 0 : dstRank = group * pow(base, powerSteps_) + dstGroupIdx;
359 0 : if (step == 0) {
360 0 : CHK_RET(links[dstRank]->TxAck(stream_));
361 0 : CHK_RET(links[dstRank]->RxAck(stream_));
362 : }
363 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
364 0 : dst = outputMem_.range(slicePower_[step].offset, slicePower_[step].size);
365 0 : src = DeviceMem::create(static_cast<u8*>(remMemPtr) + slicePower_[step].offset, slicePower_[step].size);
366 0 : CHK_RET(HcclD2DMemcpyAsync(
367 : dispatcher_, dst, src, stream_, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
368 0 : if (step != (powerSteps_ - finalSteps_ - 1)) {
369 0 : CHK_RET(RunBetweenStep(
370 : rank, dstRank, group * pow(base, powerSteps_) + (groupIdx ^ (1 << (powerSteps_ - 1 - step - 1))),
371 : links));
372 : } else {
373 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
374 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
375 : }
376 : }
377 0 : return HCCL_SUCCESS;
378 0 : }
379 0 : HcclResult AllGatherHDStage::RunAllGatherLastTwo(u32 rank, u32 rankSize, const std::vector<LINK>& links)
380 : {
381 : // mesh
382 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
383 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
384 0 : CHK_RET(MainRecordSub(base));
385 0 : CHK_RET(SubWaitMain(base));
386 :
387 0 : u32 subGroupSize = static_cast<u32>(pow(base, base));
388 0 : CHK_PRT_RET(
389 : meshStreams_.size() < (subGroupSize - 1),
390 : HCCL_ERROR("[AllGatherHDStage][RunAllGatherLastTwo]rank[%u] count[%llu] failed", rank, count_),
391 : HCCL_E_RESERVED);
392 0 : for (u32 round = 1; round < subGroupSize; round++) {
393 0 : u32 dstRank = rank / subGroupSize * subGroupSize + BackwardRank(rank % subGroupSize, subGroupSize, round);
394 0 : Stream& subStream = round == (subGroupSize - 1) ? stream_ : meshStreams_[round - 1];
395 0 : CHK_RET(links[dstRank]->TxAck(subStream));
396 0 : CHK_RET(links[dstRank]->RxAck(subStream));
397 : }
398 :
399 0 : CHK_RET(SubRecordMain(base));
400 0 : CHK_RET(MainWaitSub(base));
401 :
402 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
403 :
404 0 : CHK_RET(SubWaitMain(meshStreams_.size()));
405 0 : CHK_RET(MainRecordSub(meshStreams_.size()));
406 :
407 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
408 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize / subGroupSize);
409 : DeviceMem dst = userMemOut.range(
410 0 : totalSize_ * rankSize / subGroupSize * (resMap[rank % subGroupSize]), totalSize_ * rankSize / subGroupSize);
411 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
412 0 : }
413 :
414 0 : for (u32 round = 1; round < subGroupSize; round++) {
415 0 : u32 dstRank = rank / subGroupSize * subGroupSize + BackwardRank(rank % subGroupSize, subGroupSize, round);
416 0 : Stream& subStream = meshStreams_[round - 1];
417 0 : void* remMemPtr = nullptr;
418 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
419 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remMemPtr), totalSize_ * rankSize / subGroupSize);
420 : DeviceMem dst = userMemOut.range(
421 0 : totalSize_ * rankSize / subGroupSize * (resMap[dstRank % subGroupSize]),
422 0 : totalSize_ * rankSize / subGroupSize);
423 0 : CHK_RET(HcclD2DMemcpyAsync(
424 : dispatcher_, dst, src, subStream, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
425 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
426 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
427 0 : }
428 0 : CHK_RET(SubRecordMain(meshStreams_.size()));
429 0 : CHK_RET(MainWaitSub(meshStreams_.size()));
430 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
431 0 : return HCCL_SUCCESS;
432 0 : }
433 :
434 0 : HcclResult AllGatherHDStage::RunAllGatherLastOne(u32 rank, u32 rankSize, const std::vector<LINK>& links)
435 : {
436 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
437 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
438 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
439 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
440 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
441 :
442 0 : u32 dstRank = group * pow(base, powerSteps_) + (groupIdx ^ (1 << 0));
443 0 : CHK_RET(links[dstRank]->TxAck(stream_));
444 0 : CHK_RET(links[dstRank]->RxAck(stream_));
445 :
446 0 : CHK_RET(MainRecordSub(1));
447 0 : CHK_RET(SubWaitMain(1));
448 : // 本地拷贝
449 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
450 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize / base);
451 0 : DeviceMem dst = userMemOut.range(totalSize_ * rankSize / base * (rank % base), totalSize_ * rankSize / base);
452 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
453 0 : }
454 : // 对端拷到usrout上
455 0 : void* remMemPtr = nullptr;
456 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
457 0 : DeviceMem dst = userMemOut.range(totalSize_ * rankSize / base * (1 - (rank % base)), totalSize_ * rankSize / base);
458 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remMemPtr) + 0, totalSize_ * rankSize / base);
459 0 : CHK_RET(HcclD2DMemcpyAsync(
460 : dispatcher_, dst, src, meshStreams_[0], links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
461 :
462 0 : CHK_RET(SubRecordMain(1));
463 0 : CHK_RET(MainWaitSub(1));
464 :
465 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
466 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
467 :
468 0 : return HCCL_SUCCESS;
469 0 : }
470 :
471 0 : HcclResult AllGatherHDStage::RunAllGatherLast(
472 : [[maybe_unused]] u32 rank, u32 rankSize, [[maybe_unused]] const std::vector<LINK>& links)
473 : {
474 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
475 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
476 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize);
477 0 : DeviceMem dst = userMemOut.range(0, totalSize_ * rankSize);
478 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
479 0 : }
480 0 : return HCCL_SUCCESS;
481 : }
482 :
483 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_HD_STAGE, AllGatherHDStage);
484 : } // namespace hccl
|
