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