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