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("[AllGatherHDStage][RunAsync]rank[%u] count[%llu] failed in step", rank, count_), ret);
165 :
166 0 : HCCL_INFO("AllGatherHDStage finished: rank[%u] ranksize[%u]", rank, rankSize);
167 0 : return HCCL_SUCCESS;
168 : }
169 :
170 0 : HcclResult AllGatherHDStage::ReverseId(u32 oriIdx, u32& revIdx)
171 : {
172 0 : revIdx = 0;
173 0 : u32 powerBase = 0;
174 0 : for (u32 i = 0; i < powerSteps_; i++) {
175 0 : powerBase = static_cast<u32>(pow(base, i));
176 0 : revIdx += (oriIdx / powerBase % base) * static_cast<u32>(pow(base, powerSteps_ - i - 1));
177 : }
178 0 : return HCCL_SUCCESS;
179 : }
180 :
181 0 : HcclResult AllGatherHDStage::PrepareSliceData(
182 : u32 subRank, u32 subRankSize, u32 size, u32 batchSize, std::vector<Slice>& slices) const
183 : {
184 0 : Slice temp;
185 0 : u32 power = static_cast<u32>(log2(subRankSize));
186 0 : slices.clear();
187 0 : slices.reserve(power);
188 0 : for (u32 step = 0; step < power; step++) {
189 0 : u32 sliceNum = pow(base, step);
190 0 : u32 offset = static_cast<u32>(subRank ^ (1 << step)) / sliceNum * sliceNum;
191 0 : temp.offset = (offset * size) % batchSize;
192 0 : temp.size = sliceNum * size;
193 0 : slices.push_back(temp);
194 : }
195 0 : return HCCL_SUCCESS;
196 : }
197 :
198 0 : HcclResult AllGatherHDStage::RunAllGatherStage(u32 rank, u32 rankSize, const std::vector<LINK>& links)
199 : {
200 0 : HCCL_INFO(
201 : "RunAllGatherStage run: rank[%u] totalrank[%u] outputMem[%p] count[%llu]", rank, rankSize, outputMem_.ptr(),
202 : count_);
203 0 : u32 unitSize = SIZE_TABLE[dataType_];
204 0 : totalSize_ = unitSize * count_;
205 : // 对应因式分解中的2的幂次部分
206 0 : powerSteps_ = static_cast<u32>(log2(rankSize & (-rankSize)));
207 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
208 0 : if (powerSteps_ >= base) {
209 0 : finalSteps_ = base;
210 0 : } else if (powerSteps_ >= 1) {
211 0 : finalSteps_ = 1;
212 : }
213 : }
214 : // 对应因式分解中的奇数部分
215 0 : noPower_ = rankSize / (rankSize & (-rankSize));
216 0 : CHK_RET(RunPreCopy(rank, rankSize, links));
217 0 : if (noPower_ > 1) {
218 0 : CHK_RET(RunAllGatherNoPower(rank, rankSize, links));
219 : }
220 0 : if ((powerSteps_ - finalSteps_) >= 1) {
221 0 : CHK_RET(RunAllGatherPower(rank, rankSize, links));
222 : }
223 0 : if (finalSteps_ == base) {
224 0 : CHK_RET(RunAllGatherLastTwo(rank, rankSize, links));
225 0 : } else if (finalSteps_ == 1) {
226 0 : CHK_RET(RunAllGatherLastOne(rank, rankSize, links));
227 : } else {
228 0 : CHK_RET(RunAllGatherLast(rank, rankSize, links));
229 : }
230 0 : return HCCL_SUCCESS;
231 : }
232 :
233 0 : HcclResult AllGatherHDStage::RunPreCopy(u32 rank, u32 rankSize, const std::vector<LINK>& links)
234 : {
235 : // 交换数据
236 0 : HCCL_INFO(
237 : "RunPreCopy run: rank[%u] totalrank[%u] outputMem[%p] count[%llu]", rank, rankSize, outputMem_.ptr(), count_);
238 0 : std::vector<u32> noPowerMap;
239 0 : GetRankMapping(noPower_, noPowerMap);
240 0 : std::vector<u32> noPowerRevMap(noPower_);
241 0 : CHK_PRT_RET(
242 : noPowerMap.size() != noPower_,
243 : HCCL_ERROR("[AllGatherHDStage][RunPreCopy]rank[%u] count[%llu] failed", rank, count_), HCCL_E_RESERVED);
244 0 : for (u32 i = 0; i < noPowerMap.size(); i++) {
245 0 : noPowerRevMap[noPowerMap[i]] = i;
246 : }
247 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
248 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
249 0 : u32 revRank = 0;
250 0 : u32 revIdx = 0;
251 0 : CHK_RET(ReverseId(groupIdx, revIdx));
252 : // 将revRank的数据写到本卡
253 0 : revRank = revIdx * noPower_ + noPowerMap[group];
254 : // 将本卡数据写到revRankrev卡
255 0 : groupIdx = rank % noPower_;
256 0 : group = rank / noPower_;
257 0 : CHK_RET(ReverseId(group, revIdx));
258 0 : u32 revRankrev = noPowerRevMap[groupIdx] * static_cast<u32>(pow(base, powerSteps_)) + revIdx;
259 0 : DeviceMem UserMemIn = DeviceMem::create(opInfo_->inputAddr, totalSize_);
260 0 : if (revRank != rank && revRankrev != rank) {
261 0 : CHK_RET(links[revRank]->TxAck(stream_));
262 0 : CHK_RET(links[revRankrev]->RxAck(stream_));
263 0 : void* remMemPtr = nullptr;
264 0 : CHK_RET(links[revRankrev]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
265 : DeviceMem dst = DeviceMem::create(
266 0 : static_cast<u8*>(remMemPtr) + (rank % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_,
267 0 : totalSize_);
268 0 : DeviceMem src = UserMemIn;
269 0 : CHK_RET(HcclD2DMemcpyAsync(
270 : dispatcher_, dst, src, stream_, links[revRankrev]->GetRemoteRank(), links[revRankrev]->GetLinkType()));
271 0 : CHK_RET(links[revRankrev]->TxDataSignal(stream_));
272 0 : CHK_RET(links[revRank]->RxDataSignal(stream_));
273 0 : } else {
274 : DeviceMem dst
275 0 : = outputMem_.range((rank % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_, totalSize_);
276 0 : DeviceMem src = UserMemIn;
277 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
278 0 : }
279 0 : return HCCL_SUCCESS;
280 0 : }
281 :
282 0 : HcclResult AllGatherHDStage::RunAllGatherNoPower(u32 rank, u32 rankSize, const std::vector<LINK>& links)
283 : {
284 : std::unique_ptr<AlgTemplateBase> tempAlg
285 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_NHR, dispatcher_);
286 0 : CHK_SMART_PTR_NULL(tempAlg);
287 0 : CHK_RET(tempAlg->Prepare(true));
288 0 : HCCL_INFO(
289 : "[AllGatherHDStage][RunAllGather] rank[%u] tempAlg AllGatherNHR inputMem[%p] outputMem[%p] mem_size[%llu] "
290 : "count[%llu] planeID:[%d]",
291 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
292 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
293 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
294 0 : u32 revIdx = 0;
295 0 : CHK_RET(ReverseId(groupIdx, revIdx));
296 0 : u64 baseOffset = ((revIdx * noPower_) % (rankSize / static_cast<u32>(pow(base, finalSteps_)))) * totalSize_;
297 0 : std::vector<Slice> slices;
298 0 : for (u32 i = 0; i < noPower_; i++) {
299 0 : Slice temp;
300 0 : temp.offset = i * totalSize_;
301 0 : temp.size = totalSize_;
302 0 : slices.push_back(temp);
303 : }
304 0 : DeviceMem nhrOutput = outputMem_.range(baseOffset, outputMem_.size() - baseOffset);
305 0 : CHK_RET(tempAlg->Prepare(
306 : nhrOutput, nhrOutput, nhrOutput, count_, dataType_, stream_, reductionOp_, 0, slices, baseOffset));
307 :
308 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
309 :
310 0 : std::vector<LINK> nhrLinks;
311 0 : for (u32 i = 0; i < noPower_; i++) {
312 0 : u32 remote = i * static_cast<u32>(pow(base, powerSteps_)) + groupIdx;
313 0 : nhrLinks.push_back(links[remote]);
314 : }
315 0 : return tempAlg->RunAsync(group, noPower_, nhrLinks);
316 0 : }
317 :
318 0 : HcclResult AllGatherHDStage::RunBetweenStep(u32 rank, u32 neighCur, u32 neighNext, const std::vector<LINK>& links)
319 : {
320 : (void)rank;
321 0 : CHK_RET(MainRecordSub(1));
322 0 : CHK_RET(SubWaitMain(1));
323 :
324 0 : CHK_RET(links[neighCur]->TxDataSignal(meshStreams_[0]));
325 0 : CHK_RET(links[neighCur]->RxDataSignal(meshStreams_[0]));
326 :
327 0 : CHK_RET(links[neighNext]->TxAck(stream_));
328 0 : CHK_RET(links[neighNext]->RxAck(stream_));
329 :
330 0 : CHK_RET(SubRecordMain(1));
331 0 : CHK_RET(MainWaitSub(1));
332 :
333 0 : return HCCL_SUCCESS;
334 : }
335 :
336 0 : HcclResult AllGatherHDStage::RunAllGatherPower(u32 rank, u32 rankSize, const std::vector<LINK>& links)
337 : {
338 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
339 :
340 0 : void* remMemPtr = nullptr;
341 0 : DeviceMem dst;
342 0 : DeviceMem src;
343 : u32 dstRank;
344 : u32 dstGroupIdx;
345 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
346 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
347 0 : u32 revGroup = 0;
348 0 : CHK_RET(ReverseId(groupIdx, revGroup));
349 0 : CHK_RET(PrepareSliceData(
350 : revGroup, pow(base, powerSteps_), noPower_ * totalSize_,
351 : totalSize_ * rankSize / static_cast<u32>(pow(base, finalSteps_)), slicePower_));
352 0 : for (u32 step = 0; step < powerSteps_ - finalSteps_; step++) {
353 0 : dstGroupIdx = groupIdx ^ (1 << (powerSteps_ - 1 - step));
354 0 : dstRank = group * pow(base, powerSteps_) + dstGroupIdx;
355 0 : if (step == 0) {
356 0 : CHK_RET(links[dstRank]->TxAck(stream_));
357 0 : CHK_RET(links[dstRank]->RxAck(stream_));
358 : }
359 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
360 0 : dst = outputMem_.range(slicePower_[step].offset, slicePower_[step].size);
361 0 : src = DeviceMem::create(static_cast<u8*>(remMemPtr) + slicePower_[step].offset, slicePower_[step].size);
362 0 : CHK_RET(HcclD2DMemcpyAsync(
363 : dispatcher_, dst, src, stream_, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
364 0 : if (step != (powerSteps_ - finalSteps_ - 1)) {
365 0 : CHK_RET(RunBetweenStep(
366 : rank, dstRank, group * pow(base, powerSteps_) + (groupIdx ^ (1 << (powerSteps_ - 1 - step - 1))),
367 : links));
368 : } else {
369 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
370 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
371 : }
372 : }
373 0 : return HCCL_SUCCESS;
374 0 : }
375 0 : HcclResult AllGatherHDStage::RunAllGatherLastTwo(u32 rank, u32 rankSize, const std::vector<LINK>& links)
376 : {
377 : // mesh
378 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
379 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
380 0 : CHK_RET(MainRecordSub(base));
381 0 : CHK_RET(SubWaitMain(base));
382 :
383 0 : u32 subGroupSize = static_cast<u32>(pow(base, base));
384 0 : CHK_PRT_RET(
385 : meshStreams_.size() < (subGroupSize - 1),
386 : HCCL_ERROR("[AllGatherHDStage][RunAllGatherLastTwo]rank[%u] count[%llu] failed", rank, count_),
387 : HCCL_E_RESERVED);
388 0 : for (u32 round = 1; round < subGroupSize; round++) {
389 0 : u32 dstRank = rank / subGroupSize * subGroupSize + BackwardRank(rank % subGroupSize, subGroupSize, round);
390 0 : Stream& subStream = round == (subGroupSize - 1) ? stream_ : meshStreams_[round - 1];
391 0 : CHK_RET(links[dstRank]->TxAck(subStream));
392 0 : CHK_RET(links[dstRank]->RxAck(subStream));
393 : }
394 :
395 0 : CHK_RET(SubRecordMain(base));
396 0 : CHK_RET(MainWaitSub(base));
397 :
398 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
399 :
400 0 : CHK_RET(SubWaitMain(meshStreams_.size()));
401 0 : CHK_RET(MainRecordSub(meshStreams_.size()));
402 :
403 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
404 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize / subGroupSize);
405 : DeviceMem dst = userMemOut.range(
406 0 : totalSize_ * rankSize / subGroupSize * (resMap[rank % subGroupSize]), totalSize_ * rankSize / subGroupSize);
407 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
408 0 : }
409 :
410 0 : for (u32 round = 1; round < subGroupSize; round++) {
411 0 : u32 dstRank = rank / subGroupSize * subGroupSize + BackwardRank(rank % subGroupSize, subGroupSize, round);
412 0 : Stream& subStream = meshStreams_[round - 1];
413 0 : void* remMemPtr = nullptr;
414 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
415 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remMemPtr), totalSize_ * rankSize / subGroupSize);
416 : DeviceMem dst = userMemOut.range(
417 0 : totalSize_ * rankSize / subGroupSize * (resMap[dstRank % subGroupSize]),
418 0 : totalSize_ * rankSize / subGroupSize);
419 0 : CHK_RET(HcclD2DMemcpyAsync(
420 : dispatcher_, dst, src, subStream, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
421 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
422 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
423 0 : }
424 0 : CHK_RET(SubRecordMain(meshStreams_.size()));
425 0 : CHK_RET(MainWaitSub(meshStreams_.size()));
426 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
427 0 : return HCCL_SUCCESS;
428 0 : }
429 :
430 0 : HcclResult AllGatherHDStage::RunAllGatherLastOne(u32 rank, u32 rankSize, const std::vector<LINK>& links)
431 : {
432 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
433 0 : u32 group = rank / static_cast<u32>(pow(base, powerSteps_));
434 0 : u32 groupIdx = rank % static_cast<u32>(pow(base, powerSteps_));
435 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
436 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
437 :
438 0 : u32 dstRank = group * pow(base, powerSteps_) + (groupIdx ^ (1 << 0));
439 0 : CHK_RET(links[dstRank]->TxAck(stream_));
440 0 : CHK_RET(links[dstRank]->RxAck(stream_));
441 :
442 0 : CHK_RET(MainRecordSub(1));
443 0 : CHK_RET(SubWaitMain(1));
444 : // 本地拷贝
445 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
446 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize / base);
447 0 : DeviceMem dst = userMemOut.range(totalSize_ * rankSize / base * (rank % base), totalSize_ * rankSize / base);
448 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
449 0 : }
450 : // 对端拷到usrout上
451 0 : void* remMemPtr = nullptr;
452 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
453 0 : DeviceMem dst = userMemOut.range(totalSize_ * rankSize / base * (1 - (rank % base)), totalSize_ * rankSize / base);
454 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remMemPtr) + 0, totalSize_ * rankSize / base);
455 0 : CHK_RET(HcclD2DMemcpyAsync(
456 : dispatcher_, dst, src, meshStreams_[0], links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
457 :
458 0 : CHK_RET(SubRecordMain(1));
459 0 : CHK_RET(MainWaitSub(1));
460 :
461 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
462 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
463 :
464 0 : return HCCL_SUCCESS;
465 0 : }
466 :
467 0 : HcclResult AllGatherHDStage::RunAllGatherLast(
468 : [[maybe_unused]] u32 rank, u32 rankSize, [[maybe_unused]] const std::vector<LINK>& links)
469 : {
470 0 : if (outputMem_.ptr() != opInfo_->outputAddr) {
471 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize_ * rankSize);
472 0 : DeviceMem src = outputMem_.range(0, totalSize_ * rankSize);
473 0 : DeviceMem dst = userMemOut.range(0, totalSize_ * rankSize);
474 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
475 0 : }
476 0 : return HCCL_SUCCESS;
477 : }
478 :
479 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_HD_STAGE, AllGatherHDStage);
480 : } // namespace hccl
|
