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 "ccu_context_all_to_all_mesh2d.h"
12 : #include "ccu_instruction_all_to_all_mesh2d.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr uint16_t CKE_ID_0 = 0;
17 : constexpr uint16_t CKE_ID_1 = 1;
18 : constexpr uint16_t CKE_ID_2 = 2;
19 : constexpr uint16_t CKE_ID_3 = 3;
20 : constexpr uint16_t FST_AXIS_ID = 0;
21 : constexpr uint16_t SEC_AXIS_ID = 1;
22 :
23 3 : CcuContextAlltoAllMesh2D::CcuContextAlltoAllMesh2D(
24 3 : const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
25 3 : : CcuContextAlgBase(arg, transports, group)
26 : {
27 3 : goSize_ = CreateGroupOpSize();
28 3 : input = CreateVariable();
29 3 : bufferB = CreateVariable();
30 3 : sliceSize_ = CreateVariable();
31 3 : baseOffset = CreateVariable();
32 3 : firstTransportSize = CreateVariable();
33 3 : firstChunkOffset = CreateVariable();
34 3 : firstInputStrideLocal = CreateVariable();
35 3 : firstInputStrideAnother = CreateVariable();
36 3 : firstBufferOffset = CreateVariable();
37 3 : firstBufferStride = CreateVariable();
38 3 : firstOutputOffset = CreateVariable();
39 3 : secondTransportSize = CreateVariable();
40 3 : secondChunkOffset = CreateVariable();
41 3 : secondInputOffset = CreateVariable();
42 3 : secondInputStride = CreateVariable();
43 3 : secondBufferStrideLocal = CreateVariable();
44 3 : secondBufferStrideAnother = CreateVariable();
45 3 : secondOutputOffset = CreateVariable();
46 3 : secondOutputStride = CreateVariable();
47 3 : localAxisSignal = CreateMaskSignal();
48 :
49 3 : const CcuCtxArgAlltoAllMesh2D* ctxArg = dynamic_cast<const CcuCtxArgAlltoAllMesh2D*>(&arg);
50 3 : if (ctxArg == nullptr) {
51 0 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D::ctxArg ptr is null"));
52 : }
53 3 : if (transports.size() == 0) {
54 0 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D transports is empty"));
55 : }
56 3 : rankId_ = ctxArg->rankId;
57 3 : dimSize = ctxArg->dimSize;
58 3 : axisId = ctxArg->axisId;
59 3 : uint32_t max_dimSize = 2;
60 3 : if (dimSize.size() != max_dimSize or axisId > 1) { // dimSize不为2，或axisId超过1，则不为2D场景
61 0 : THROW<NullPtrException>(
62 0 : StringFormat("[CcuContextAlltoAllMesh2D] dimSize[%u] or axisId[%u] is invalid", dimSize.size(), axisId));
63 : }
64 3 : CHK_PRT_THROW(
65 : dimSize[0] == 0 || dimSize[1] == 0,
66 : HCCL_ERROR("[CcuContextAlltoAllMesh2D] dimSize0[%llu] or dimSize1[%llu] is zero", dimSize[0], dimSize[1]),
67 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
68 3 : dimId.emplace_back(rankId_ % dimSize[0]);
69 3 : dimId.emplace_back(rankId_ / dimSize[0]);
70 3 : localId = dimId[axisId];
71 3 : localSize = dimSize[axisId];
72 3 : anotherId = dimId[1 - axisId]; // 本rank在另一个轴上的Id
73 3 : anotherSize = dimSize[1 - axisId];
74 9 : HCCL_INFO(
75 : "[CcuContextAlltoAllMesh2D] RankId[%u], DimSize: D0[%u]--D1[%u], localId[%u], lcoalSize[%u]", rankId_,
76 : dimSize[0], dimSize[1], localId, localSize);
77 :
78 3 : AllocGoResource(LOC_CPY_LOOP_NUM); // 只用8个loop做本地搬运，每个loop搬4K
79 :
80 3 : localAxisSignalName = "CcuContextAlltoAllMesh2DAxisSync_" + std::to_string(axisId);
81 3 : anotherAxisSignalName = "CcuContextAlltoAllMesh2DAxisSync_" + std::to_string(1 - axisId);
82 3 : }
83 :
84 2 : void CcuContextAlltoAllMesh2D::InitResources()
85 : {
86 : // 用write语义，input只有本地的1个，scratch和output需要交换
87 2 : ExportMaskSignal(localAxisSignal, localAxisSignalName);
88 2 : anotherAxisSignal = ImportMaskSignal(anotherAxisSignalName);
89 :
90 2 : uint32_t transportIdx = 0;
91 6 : for (uint32_t peerId = 0; peerId < localSize; peerId++) {
92 4 : if (peerId == localId) {
93 2 : bufferA.emplace_back(CreateVariable());
94 2 : output.emplace_back(CreateVariable());
95 2 : token.emplace_back(CreateVariable());
96 : } else {
97 6 : HCCL_INFO("[CcuContextAlltoAllMesh2D]Rank[%u], PeerId[%u], TransportId[%u]", rankId_, peerId, transportIdx);
98 2 : bufferA.emplace_back(
99 4 : CreateVariable(*(transports[transportIdx]), 0)); // 获取transport中id=1的Var来传递bufferA
100 2 : output.emplace_back(CreateVariable(*(transports[transportIdx]), 1)); // 1 for output
101 2 : token.emplace_back(CreateVariable(*(transports[transportIdx]), 2)); // 2 for token
102 2 : transportIdx++;
103 : }
104 : }
105 :
106 6 : for (uint16_t i = 0; i < localSize; i++) {
107 4 : inputAddrs.emplace_back(CreateMemory());
108 4 : bufferAddrs.emplace_back(CreateMemory());
109 4 : outputAddrs.emplace_back(CreateMemory());
110 : }
111 :
112 6 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) {
113 4 : firstSignal.emplace_back(
114 8 : CreateMaskSignal()); // 每个对端发anotherSize个分片，localSize个分片共用一个信号，共anotherSize个
115 4 : secondSignal.emplace_back(CreateMaskSignal());
116 : }
117 :
118 2 : return;
119 : }
120 :
121 2 : void CcuContextAlltoAllMesh2D::LoadArgs()
122 : {
123 2 : Load(input);
124 2 : Load(output[localId]);
125 2 : Load(token[localId]);
126 2 : Load(bufferA[localId]);
127 2 : Load(bufferB);
128 2 : Load(sliceSize_);
129 2 : Load(goSize_);
130 :
131 2 : Load(baseOffset); // 10号
132 2 : Load(firstTransportSize);
133 2 : Load(firstChunkOffset);
134 2 : Load(firstInputStrideLocal);
135 2 : Load(firstInputStrideAnother);
136 2 : Load(firstBufferOffset); // 15号
137 2 : Load(firstBufferStride); // 16号
138 2 : Load(firstOutputOffset);
139 :
140 2 : Load(secondTransportSize);
141 2 : Load(secondChunkOffset);
142 2 : Load(secondInputOffset);
143 2 : Load(secondInputStride);
144 2 : Load(secondBufferStrideLocal);
145 2 : Load(secondBufferStrideAnother);
146 2 : Load(secondOutputOffset);
147 2 : Load(secondOutputStride);
148 :
149 2 : return;
150 : }
151 :
152 2 : void CcuContextAlltoAllMesh2D::ExchangeInfoAndSync()
153 : {
154 : // 交换信息并做同步，前同步固定用1,2,3号信号
155 2 : uint16_t selfBit = 1 << localId;
156 2 : uint16_t allBit = ((1 << localSize) - 1) & (~(1 << localId));
157 :
158 8 : for (auto t : transports) {
159 6 : if (t == nullptr) {
160 0 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D::Algorithm transport ptr is null"));
161 : }
162 6 : WriteVariableWithSignal(*t, bufferA[localId], 0, CKE_ID_1, selfBit); // index = 0，传递第一轮output信息
163 6 : WriteVariableWithSignal(*t, output[localId], 1, CKE_ID_2, selfBit); // index = 1，传递第二轮output信息
164 6 : WriteVariableWithSignal(*t, token[localId], 2, CKE_ID_3, selfBit); // index = 2，传递token信息
165 : }
166 2 : GroupWait(*transportGroup, CKE_ID_1, allBit);
167 2 : GroupWait(*transportGroup, CKE_ID_2, allBit);
168 2 : GroupWait(*transportGroup, CKE_ID_3, allBit);
169 :
170 2 : return;
171 : }
172 :
173 6 : void CcuContextAlltoAllMesh2D::RankSync(uint32_t signalIndex)
174 : {
175 : // 与远端做同步
176 6 : uint16_t selfBit = 1 << localId;
177 6 : uint16_t allBit = ((1 << localSize) - 1) & (~(1 << localId));
178 :
179 24 : for (auto t : transports) {
180 18 : if (t == nullptr) {
181 0 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D::Algorithm transport ptr is null"));
182 : }
183 18 : RemotePost(*t, signalIndex, selfBit);
184 : }
185 6 : GroupWait(*transportGroup, signalIndex, allBit);
186 :
187 6 : return;
188 : }
189 :
190 4 : void CcuContextAlltoAllMesh2D::AxisSync(uint32_t signalIndex)
191 : {
192 4 : const uint32_t DIE_NUM = 2; // 2个die
193 4 : if (signalIndex > 1) {
194 0 : THROW<InvalidParamsException>(
195 0 : StringFormat("[CcuContextAlltoAllMesh2D] Unexpected SignalInex[%u]", signalIndex));
196 : }
197 4 : LocalCtxPost(anotherAxisSignal, 1 << (axisId + signalIndex * DIE_NUM));
198 4 : LocalWait(localAxisSignal, 1 << (1 - axisId + signalIndex * DIE_NUM));
199 4 : return;
200 : }
201 :
202 4 : void CcuContextAlltoAllMesh2D::FirstStepOneSlice(uint16_t sliceId)
203 : {
204 4 : if (sliceId == anotherId) {
205 : // 当前分片属于对端，直接写到对端output
206 2 : uint32_t transIdx = 0; // 约定transport中的link按照rankId从小到大的顺序排列
207 6 : for (uint32_t peerId = 0; peerId < localSize; peerId++) {
208 4 : if (peerId == localId) {
209 2 : LocalPost(firstSignal[sliceId], (1 << peerId));
210 : } else {
211 2 : Write(
212 2 : *(transports[transIdx]), outputAddrs[peerId], inputAddrs[peerId], firstTransportSize,
213 2 : firstSignal[sliceId], (1 << peerId));
214 2 : transIdx++;
215 : }
216 4 : inputAddrs[peerId].addr += firstInputStrideAnother; // 给每个对端的下一片slice的input地址，增加对应偏移
217 4 : bufferAddrs[peerId].addr += firstBufferStride; // 跳过对端buffer中不需要转发的那一片
218 : }
219 : } else {
220 : // 当前分片需要对端转发，写到对端的bufferX/bufferY
221 2 : uint32_t transIdx = 0;
222 6 : for (uint32_t peerId = 0; peerId < localSize; peerId++) {
223 4 : if (peerId == localId) {
224 2 : LocalPost(firstSignal[sliceId], (1 << peerId)); // 对于本die经过转发无法到达的对端，只设置标记不发送数据
225 2 : continue;
226 : }
227 2 : Write(
228 2 : *(transports[transIdx]), bufferAddrs[peerId], inputAddrs[peerId], firstTransportSize,
229 2 : firstSignal[sliceId], (1 << peerId));
230 2 : inputAddrs[peerId].addr += firstInputStrideAnother;
231 2 : bufferAddrs[peerId].addr += firstBufferStride; // 给对端用于转发的分片，每片相对前片加localSize*sliceSize
232 2 : transIdx++;
233 : }
234 : }
235 :
236 4 : return;
237 : }
238 :
239 2 : void CcuContextAlltoAllMesh2D::FirstStep()
240 : {
241 2 : CcuRep::Memory lgSrc = CreateMemory();
242 2 : CcuRep::Memory lgDst = CreateMemory();
243 :
244 : // 统一处理token，访问第i个对端需要使用对应的token
245 6 : for (uint16_t i = 0; i < localSize; i++) {
246 4 : inputAddrs[i].token = token[i];
247 4 : bufferAddrs[i].token = token[i];
248 4 : outputAddrs[i].token = token[i];
249 : }
250 2 : lgSrc.token = token[localId];
251 2 : lgDst.token = token[localId];
252 : // 本rank的input内存块用一组mem地址来分割
253 2 : inputAddrs[0].addr = input;
254 2 : inputAddrs[0].addr += baseOffset;
255 2 : inputAddrs[0].addr += firstChunkOffset;
256 4 : for (uint16_t i = 1; i < localSize; i++) {
257 : // 准备发送给rank0对应分片的地址即为input首地址，后续rank的偏移依次递增
258 2 : inputAddrs[i].addr = inputAddrs[i - 1].addr + firstInputStrideLocal;
259 : }
260 6 : for (uint16_t i = 0; i < localSize; i++) {
261 : // output offset
262 4 : outputAddrs[i].addr = output[i];
263 4 : outputAddrs[i].addr += baseOffset;
264 4 : outputAddrs[i].addr += firstChunkOffset;
265 4 : outputAddrs[i].addr += firstOutputOffset; // 第一轮直接发送给对端的slice的偏移
266 : // buffer offset
267 4 : bufferAddrs[i].addr = bufferA[i];
268 4 : bufferAddrs[i].addr += firstBufferOffset; // 发送给对端用于转发的分片，第一片的起始偏移，后续每片步进相同长度
269 : }
270 : // 准备LG搬运的地址
271 2 : lgSrc.addr = inputAddrs[localId].addr;
272 2 : lgDst.addr = outputAddrs[localId].addr;
273 2 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) {
274 2 : if (sliceId == anotherId) {
275 2 : break;
276 : }
277 0 : lgSrc.addr += firstInputStrideAnother; // 在input中找到自身对应的那个分片，跳出循环
278 : }
279 :
280 : {
281 : // 当第一轮的搬运量为零时，跳过搬运
282 2 : CcuRep::Condition cond(this, firstTransportSize != 0);
283 :
284 6 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) { // sliceId等于dstRank在另一个维度上的id
285 4 : FirstStepOneSlice(sliceId);
286 : }
287 2 : }
288 : // Loopgroup做本地搬运，必然不为零
289 2 : if (axisId == 0) {
290 1 : LocalCopyByLoopGroup(lgDst, lgSrc, goSize_);
291 : }
292 :
293 : // 检查第一轮的数据是否已发完
294 : {
295 : // 当第一轮的搬运量非零时，检查相应完成标记
296 2 : CcuRep::Condition cond(this, firstTransportSize != 0);
297 6 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) {
298 4 : LocalWait(firstSignal[sliceId], (1 << localSize) - 1); // 等待第一轮所有分片都发完
299 : }
300 2 : }
301 :
302 4 : return;
303 2 : }
304 :
305 2 : void CcuContextAlltoAllMesh2D::SecondStep()
306 : {
307 : {
308 : // 当第二轮的搬运量为零时，跳过搬运
309 2 : CcuRep::Condition cond(this, secondTransportSize != 0);
310 :
311 : // 地址计算
312 : // input offset，本rank的input内存块用一组GSA来分割
313 2 : inputAddrs[0].addr = input;
314 2 : inputAddrs[0].addr += baseOffset;
315 2 : bufferAddrs[0].addr = bufferB;
316 2 : inputAddrs[0].addr += secondChunkOffset;
317 2 : inputAddrs[0].addr += secondInputOffset; // 一共从input发送localSize-1个分片（跳过自己），用localSize个input地址
318 4 : for (uint16_t i = 1; i < localSize; i++) {
319 2 : inputAddrs[i].addr = inputAddrs[i - 1].addr + secondInputStride;
320 : // 每轮给每个对端从buffer发送1个分片，共anotherSize-1轮
321 2 : bufferAddrs[i].addr = bufferAddrs[i - 1].addr + secondBufferStrideLocal;
322 : }
323 : // output offset
324 6 : for (uint16_t i = 0; i < localSize; i++) {
325 : // 给每个对端的output写anotherSize个分片，这些分片的src的rankId从offset开始，以stride步进
326 4 : outputAddrs[i].addr = output[i];
327 4 : outputAddrs[i].addr += baseOffset;
328 4 : outputAddrs[i].addr += secondChunkOffset;
329 4 : outputAddrs[i].addr += secondOutputOffset;
330 : }
331 :
332 : // 从input与buffer中给每个对端发anotherSize个分片
333 6 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) {
334 4 : uint32_t transIdx = 0;
335 12 : for (uint32_t peerId = 0; peerId < localSize; peerId++) {
336 8 : if (peerId == localId) {
337 4 : LocalPost(secondSignal[sliceId], (1 << peerId)); // 给自己的分片在第一轮已经发过，第二轮只设置标记
338 4 : continue;
339 : }
340 4 : if (sliceId == anotherId) {
341 : // 从input发出
342 2 : Write(
343 2 : *(transports[transIdx]), outputAddrs[peerId], inputAddrs[peerId], secondTransportSize,
344 2 : secondSignal[sliceId], (1 << peerId));
345 : } else {
346 : // 从buffer发出
347 2 : Write(
348 2 : *(transports[transIdx]), outputAddrs[peerId], bufferAddrs[peerId], secondTransportSize,
349 2 : secondSignal[sliceId], (1 << peerId));
350 : }
351 4 : transIdx++;
352 4 : outputAddrs[peerId].addr += secondOutputStride;
353 4 : bufferAddrs[peerId].addr += secondBufferStrideAnother;
354 : }
355 : }
356 6 : for (uint16_t sliceId = 0; sliceId < anotherSize; sliceId++) {
357 4 : LocalWait(secondSignal[sliceId], (1 << localSize) - 1); // 等待第二轮所有分片都发完
358 : }
359 2 : }
360 :
361 2 : return;
362 : }
363 :
364 1 : void CcuContextAlltoAllMesh2D::CreateLocalCopyLoop()
365 : {
366 1 : std::string opStr = "a2a_localcpy_loopgroup";
367 3 : for (uint32_t index = 0; index < 2; index++) { // 需要2个Loop
368 2 : CcuRep::LoopBlock lb(this, "a2a_localcpy_loop_" + std::to_string(index));
369 2 : CcuRep::Memory src = CreateMemory();
370 2 : CcuRep::Variable len = CreateVariable();
371 2 : CcuRep::Memory dst = CreateMemory();
372 2 : lb(src, dst, len);
373 :
374 2 : CcuRep::CcuBuffer buf = moRes.ccuBuffer[index * moConfig.msInterleave];
375 2 : CcuRep::MaskSignal sem = moRes.maskSignal[index];
376 :
377 2 : LocalCopy(buf, src, len, sem);
378 2 : LocalWait(sem);
379 2 : LocalCopy(dst, buf, len, sem);
380 2 : LocalWait(sem);
381 2 : }
382 2 : return;
383 1 : }
384 :
385 1 : void CcuContextAlltoAllMesh2D::LocalCopyByLoopGroup(CcuRep::Memory dst, CcuRep::Memory src, GroupOpSize& goPara)
386 : {
387 1 : std::string opStr = "a2a_localcpy_loopgroup";
388 1 : CreateLocalCopyLoop();
389 :
390 : {
391 1 : CcuRep::Condition cond(this, goPara.loopParam != 0);
392 :
393 1 : CcuRep::Variable loopParam = CreateVariable();
394 1 : loopParam = CcuRep::GetLoopParam(0, moConfig.memSlice * moConfig.loopCount, 0);
395 1 : loopParam += goPara.loopParam;
396 :
397 1 : CcuRep::Variable sliceSize = CreateVariable();
398 1 : sliceSize = moConfig.memSlice;
399 2 : auto lc = Loop("a2a_localcpy_loop_0")(src, dst, sliceSize);
400 :
401 1 : CcuRep::Variable paraCfg = CreateVariable();
402 1 : paraCfg = CcuRep::GetParallelParam(moConfig.loopCount - 1, 0, 1);
403 1 : CcuRep::Variable offsetCfg = CreateVariable();
404 1 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
405 6 : LoopGroup({lc}, {loopParam}, paraCfg, offsetCfg);
406 1 : }
407 :
408 : {
409 1 : CcuRep::Condition cond(this, goPara.parallelParam != 0);
410 :
411 1 : dst.addr += goPara.addrOffset;
412 1 : src.addr += goPara.addrOffset;
413 2 : auto lc0 = Loop("a2a_localcpy_loop_0")(src, dst, goPara.residual);
414 :
415 1 : src.addr += goPara.residual;
416 1 : dst.addr += goPara.residual;
417 1 : CcuRep::Variable sliceSize = CreateVariable();
418 1 : sliceSize = moConfig.memSlice;
419 2 : auto lc1 = Loop("a2a_localcpy_loop_1")(src, dst, sliceSize);
420 :
421 1 : CcuRep::Variable loopCfg0 = CreateVariable();
422 1 : loopCfg0 = CcuRep::GetLoopParam(0, 0, 1);
423 1 : CcuRep::Variable loopCfg1 = CreateVariable();
424 1 : loopCfg1 = CcuRep::GetLoopParam(0, 0, 1);
425 1 : CcuRep::Variable offsetCfg = CreateVariable();
426 1 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
427 8 : LoopGroup({lc0, lc1}, {loopCfg0, loopCfg1}, goPara.parallelParam, offsetCfg);
428 1 : }
429 5 : }
430 :
431 2 : void CcuContextAlltoAllMesh2D::Algorithm()
432 : {
433 : // 初始化寄存器资源 & 加载外部输入参数
434 6 : HCCL_INFO("[CcuContextAlltoAllMesh2D] AllgatherMesh1D Algorithm Init Begins.");
435 2 : InitResources();
436 2 : LoadArgs();
437 :
438 : // 第一轮，X方向发a，Y方向发后b，到对端的块均放在output，要沿X转发的b块放在对端的bufferX，根据转发目的、自身locId两级偏移
439 6 : HCCL_INFO("[CcuContextAlltoAllMesh2D] Algorithm first step begins.");
440 2 : ExchangeInfoAndSync();
441 2 : FirstStep();
442 2 : RankSync(CKE_ID_0);
443 2 : AxisSync(FST_AXIS_ID);
444 :
445 : // 第二轮，从input和buffer中将剩余的本端分片以及待转发分片发给对端；其中给每个对端发1个本端分片，localSize-1个转发分片
446 6 : HCCL_INFO("[CcuContextAlltoAllMesh2D] Algorithm second step begins.");
447 2 : RankSync(CKE_ID_1);
448 2 : SecondStep();
449 2 : RankSync(CKE_ID_0);
450 2 : AxisSync(SEC_AXIS_ID);
451 :
452 6 : HCCL_INFO("[CcuContextAlltoAllMesh2D] Algorithm Ends.");
453 2 : return;
454 : }
455 :
456 2 : void CcuContextAlltoAllMesh2D::CalculateArgs(const CcuTaskArgAlltoAllMesh2D* taskArg)
457 : {
458 2 : if (taskArg == nullptr) {
459 0 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D::taskArg ptr is null"));
460 : }
461 :
462 2 : uint64_t sendStride = taskArg->sendStride;
463 2 : uint64_t recvStride = taskArg->recvStride;
464 2 : uint64_t aSize = taskArg->aSize;
465 2 : uint64_t bSize = taskArg->bSize;
466 2 : uint64_t sendLength = taskArg->sendLength;
467 :
468 2 : uint64_t sliceSize = aSize + bSize;
469 2 : uint64_t srcStride = sendLength + sendStride;
470 2 : uint64_t dstStride = sendLength + recvStride;
471 :
472 : // 根据axisId决定bufferA与bufferB的地址，暂定a与b的大小相等
473 2 : if (axisId == 0) {
474 2 : firstTransportSizeValue = aSize;
475 2 : firstChunkOffsetValue = 0;
476 2 : firstInputStrideLocalValue = srcStride;
477 2 : firstInputStrideAnotherValue = dimSize[0] * srcStride;
478 :
479 2 : secondTransportSizeValue = bSize;
480 2 : secondChunkOffsetValue = aSize;
481 2 : secondInputOffsetValue = dimId[1] * dimSize[0] * srcStride;
482 2 : secondInputStrideValue = srcStride;
483 2 : secondOutputOffsetValue = dimId[0] * dstStride;
484 2 : secondOutputStrideValue = dimSize[0] * dstStride;
485 : } else {
486 0 : firstTransportSizeValue = bSize;
487 0 : firstChunkOffsetValue = aSize;
488 0 : firstInputStrideLocalValue = dimSize[0] * srcStride;
489 0 : firstInputStrideAnotherValue = srcStride;
490 :
491 0 : secondTransportSizeValue = aSize;
492 0 : secondChunkOffsetValue = 0;
493 0 : secondInputOffsetValue = dimId[0] * srcStride;
494 0 : secondInputStrideValue = dimSize[0] * srcStride;
495 0 : secondOutputOffsetValue = dimId[1] * dimSize[0] * dstStride;
496 0 : secondOutputStrideValue = dstStride;
497 : }
498 :
499 2 : firstBufferOffsetValue = dimId[axisId] * sliceSize;
500 2 : firstBufferStrideValue = dimSize[axisId] * sliceSize;
501 2 : firstOutputOffsetValue = rankId_ * dstStride;
502 :
503 2 : secondBufferStrideLocalValue = dimSize[1 - axisId] * sliceSize;
504 2 : secondBufferStrideAnotherValue = sliceSize;
505 :
506 2 : return;
507 : }
508 :
509 8 : std::vector<uint64_t> CcuContextAlltoAllMesh2D::GeneArgs(const CcuTaskArg& arg)
510 : {
511 8 : const CcuTaskArgAlltoAllMesh2D* taskArg = dynamic_cast<const CcuTaskArgAlltoAllMesh2D*>(&arg);
512 8 : if (taskArg == nullptr) {
513 12 : THROW<NullPtrException>(StringFormat("CcuContextAlltoAllMesh2D::taskArg ptr is null"));
514 : }
515 :
516 : // input&output&buffer地址
517 2 : uint64_t inputAddr = taskArg->inputAddr;
518 2 : uint64_t outputAddr = taskArg->outputAddr;
519 2 : uint64_t scratchAddr = taskArg->scratchAddr;
520 2 : uint64_t tokenInfo = taskArg->token;
521 2 : uint64_t sliceSizeValue = taskArg->aSize + taskArg->bSize;
522 :
523 : // scratch的前rankSize*sliceSize大小为bufferY，后一块为bufferX
524 : // die0第一轮写到对端的bufferY，第二轮从本端bufferX发送；die1第一轮写到对端的bufferX，第二轮从本端bufferY发送
525 2 : uint64_t bufferAAddr = 0;
526 2 : uint64_t bufferBAddr = 0;
527 2 : if (axisId == 0) {
528 2 : bufferAAddr = scratchAddr; // 需要交换给对端，是bufferY
529 2 : bufferBAddr = scratchAddr + dimSize[0] * dimSize[1] * sliceSizeValue; // bufferX rankSize * sliceSize
530 : } else {
531 0 : bufferAAddr = scratchAddr + dimSize[0] * dimSize[1] * sliceSizeValue; // bufferX
532 0 : bufferBAddr = scratchAddr; // 不需要交换给对端，是bufferY
533 : }
534 :
535 : // loopgroup按照sliceSize大小做本地搬运，只die0执行
536 2 : auto goSize = CalGoSize(taskArg->aSize + taskArg->bSize);
537 2 : CalculateArgs(taskArg);
538 :
539 6 : HCCL_INFO(
540 : "[CcuContextAlltoAllMesh2D][GeneArgs] RankId[%u]--AxisId[%u], inputAddr[%llu], outputAddr[%llu], \
541 : bufferA[%llu], bufferB[%llu], goSize--[%llu][%llu][%llu][%llu], sendStride[%llu], recvStride[%llu], \
542 : sendRecvSize[%llu], sendLength[%llu], aSize[%llu], bSize[%llu], baseOffset[%llu]",
543 : rankId_, axisId, inputAddr, outputAddr, bufferAAddr, bufferBAddr, goSize[0], goSize[1], goSize[2], goSize[3],
544 : taskArg->sendStride, taskArg->recvStride, sliceSizeValue, taskArg->sendLength, taskArg->aSize, taskArg->bSize,
545 : taskArg->baseOffset);
546 :
547 6 : HCCL_INFO(
548 : "[CcuContextAlltoAllMesh2D][CalculateArgs] firstTransportSize[%llu], firstChunkOffset[%llu], \
549 : firstInputStrideLocal[%llu], firstInputStrideAnother[%llu], firstBufferOffset[%llu], firstBufferStride[%llu], \
550 : firstOutputOffset[%llu], secondTransportSize[%llu], secondChunkOffset[%llu], secondInputOffset[%llu], \
551 : secondInputStride[%llu], secondBufferStrideLocal[%llu], secondBufferStrideAnother[%llu], \
552 : secondOutputOffset[%llu], secondOutputStride[%llu]",
553 : firstTransportSizeValue, firstChunkOffsetValue, firstInputStrideLocalValue, firstInputStrideAnotherValue,
554 : firstBufferOffsetValue, firstBufferStrideValue, firstOutputOffsetValue, secondTransportSizeValue,
555 : secondChunkOffsetValue, secondInputOffsetValue, secondInputStrideValue, secondBufferStrideLocalValue,
556 : secondBufferStrideAnotherValue, secondOutputOffsetValue, secondOutputStrideValue);
557 :
558 : return {
559 : inputAddr,
560 : outputAddr,
561 : tokenInfo,
562 : bufferAAddr,
563 : bufferBAddr,
564 : sliceSizeValue,
565 2 : goSize[0],
566 2 : goSize[1],
567 2 : goSize[2],
568 2 : goSize[3],
569 2 : taskArg->baseOffset,
570 2 : firstTransportSizeValue,
571 2 : firstChunkOffsetValue,
572 2 : firstInputStrideLocalValue,
573 2 : firstInputStrideAnotherValue,
574 2 : firstBufferOffsetValue,
575 2 : firstBufferStrideValue,
576 2 : firstOutputOffsetValue,
577 2 : secondTransportSizeValue,
578 2 : secondChunkOffsetValue,
579 2 : secondInputOffsetValue,
580 2 : secondInputStrideValue,
581 2 : secondBufferStrideLocalValue,
582 2 : secondBufferStrideAnotherValue,
583 2 : secondOutputOffsetValue,
584 6 : secondOutputStrideValue};
585 2 : }
586 : } // namespace Hccl
|
