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_gather_mesh2d.h"
12 : #include "ccu_instruction_all_gather_mesh2d.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr int OUTPUT_XN_ID = 1;
17 : constexpr int TOKEN_XN_ID = 2;
18 : constexpr int CKE_IDX_0 = 0;
19 : constexpr int CKE_IDX_1 = 1;
20 : constexpr int CKE_IDX_2 = 2;
21 : constexpr int CKE_IDX_3 = 3;
22 : constexpr int CKE_IDX_4 = 4;
23 : constexpr int FST_AXIS_ID = 0;
24 : constexpr int SEC_AXIS_ID = 1;
25 :
26 0 : CcuContextAllGatherMesh2D::CcuContextAllGatherMesh2D(
27 0 : const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
28 0 : : CcuContextAlgBase(arg, transports, group)
29 : {
30 0 : xAxisSize_ = CreateVariable();
31 0 : yAxisSize_ = CreateVariable();
32 0 : offset_ = CreateVariable();
33 0 : sliceSize_ = CreateVariable();
34 0 : firstInOffset_ = CreateVariable();
35 0 : firstOutOffset_ = CreateVariable();
36 0 : goASize_ = CreateGroupOpSize();
37 0 : goBSize_ = CreateGroupOpSize();
38 0 : secondInOutBaseOffset_ = CreateVariable();
39 0 : secondInOutStepOffset_ = CreateVariable();
40 0 : localAxisSignal_ = CreateMaskSignal();
41 :
42 0 : const CcuCtxArgAllGatherMesh2D* ctxArg = dynamic_cast<const CcuCtxArgAllGatherMesh2D*>(&arg);
43 0 : if (ctxArg == nullptr) {
44 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMesh2D::ctxArg ptr is null"));
45 : }
46 0 : rankId_ = ctxArg->rankId_;
47 0 : dimSize_ = ctxArg->dimSize_;
48 0 : axisId_ = ctxArg->axisId_;
49 0 : uint32_t max_dimSize = 2;
50 0 : if (dimSize_.size() != max_dimSize or axisId_ > 1) {
51 0 : THROW<NullPtrException>(
52 0 : StringFormat("CcuContextAllGatherMesh2D::dimSize[%u] or axisId[%u] is invalid", dimSize_.size(), axisId_));
53 : }
54 0 : CHK_PRT_THROW(
55 : dimSize_[0] == 0 || dimSize_[1] == 0,
56 : HCCL_ERROR("[CcuContextAllGatherMesh2D] dimSize0[%llu] or dimSize1[%llu] is zero", dimSize_[0], dimSize_[1]),
57 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
58 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
59 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
60 0 : localId_ = dimId_[axisId_];
61 0 : localSize_ = dimSize_[axisId_];
62 0 : localAxisSignalName_ = "CcuContextAllGatherMesh2DAxisSync_" + std::to_string(axisId_);
63 0 : anotherAxisSignalName_ = "CcuContextAllGatherMesh2DAxisSync_" + std::to_string(1 - axisId_);
64 0 : HCCL_INFO(
65 : "[CcuContextAllGatherMesh2D] RankId[%u], DimSize: D0[%u]--D1[%u], localId[%u], lcoalSize[%u]", rankId_,
66 : dimSize_[0], dimSize_[1], localId_, localSize_);
67 0 : }
68 :
69 0 : void CcuContextAllGatherMesh2D::InitResources()
70 : {
71 0 : input_.push_back(CreateVariable());
72 :
73 0 : uint32_t transportIdx = 0;
74 0 : for (uint32_t peerId = 0; peerId < localSize_; peerId++) {
75 0 : if (peerId == localId_) {
76 0 : output_.push_back(CreateVariable());
77 0 : token_.push_back(CreateVariable());
78 : } else {
79 0 : HCCL_INFO(
80 : "[CcuContextAllGatherMesh2D] MyRank[%u], peerId[%u], transportIdx[%u]", rankId_, peerId, transportIdx);
81 0 : CHK_PRT_RET(
82 : transports[transportIdx] == nullptr,
83 : HCCL_ERROR("[CcuContextAllGatherMesh2D] Algorithm transport ptr is null"), );
84 0 : output_.push_back(CreateVariable((*transports[transportIdx]), OUTPUT_XN_ID));
85 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
86 0 : transportIdx++;
87 : }
88 : }
89 :
90 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
91 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
92 0 : return;
93 : }
94 :
95 0 : void CcuContextAllGatherMesh2D::LoadArgs()
96 : {
97 0 : Load(input_[0]);
98 0 : Load(output_[localId_]);
99 0 : Load(token_[localId_]);
100 0 : Load(xAxisSize_);
101 0 : Load(yAxisSize_);
102 0 : Load(offset_);
103 0 : Load(sliceSize_);
104 0 : Load(firstInOffset_);
105 0 : Load(firstOutOffset_);
106 0 : Load(secondInOutBaseOffset_);
107 0 : Load(secondInOutStepOffset_);
108 0 : Load(goASize_);
109 0 : Load(goBSize_);
110 :
111 0 : return;
112 : }
113 :
114 0 : void CcuContextAllGatherMesh2D::ExchangeInfoAndSync()
115 : {
116 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] ExchangeInfoAndSync run begins");
117 0 : uint16_t selfBit = 1 << localId_;
118 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
119 :
120 0 : for (auto t : transports) {
121 0 : if (t == nullptr) {
122 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMesh2D::Algorithm transport ptr is null"));
123 : }
124 0 : WriteVariableWithSignal(*t, output_[localId_], OUTPUT_XN_ID, CKE_IDX_1, selfBit); // index = 1，传递output信息
125 0 : WriteVariableWithSignal(*t, token_[localId_], TOKEN_XN_ID, CKE_IDX_2, selfBit); // index = 2，传递token信息
126 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] change addr success");
127 : }
128 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit);
129 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit);
130 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] ExchangeInfoAndSync run finished");
131 0 : return;
132 : }
133 :
134 0 : void CcuContextAllGatherMesh2D::RankSync(uint32_t signalIndex)
135 : {
136 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] RankSync run begins");
137 0 : uint16_t selfBit = 1 << localId_;
138 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
139 :
140 0 : for (auto t : transports) {
141 0 : if (t == nullptr) {
142 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMesh2D::Algorithm transport ptr is null"));
143 : }
144 0 : RemotePost(*t, signalIndex, selfBit);
145 : }
146 0 : GroupWait(*transportGroup, signalIndex, allBit);
147 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] RankSync run ends");
148 0 : return;
149 : }
150 :
151 0 : void CcuContextAllGatherMesh2D::AxisSync(uint32_t signalIndex)
152 : {
153 0 : const uint32_t DIE_NUM = 2;
154 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] AxisSync run begins");
155 0 : LocalCtxPost(anotherAxisSignal_, 1 << (axisId_ + (signalIndex * DIE_NUM)));
156 0 : LocalWait(localAxisSignal_, 1 << ((1 - axisId_) + (signalIndex * DIE_NUM)));
157 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] AxisSync run ends");
158 0 : return;
159 : }
160 :
161 0 : void CcuContextAllGatherMesh2D::FirstStep()
162 : {
163 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] firstStep run begins");
164 0 : CcuRep::Memory src = CreateMemory();
165 0 : src.addr = input_[0];
166 0 : src.addr += firstInOffset_;
167 0 : src.token = token_[localId_];
168 :
169 0 : std::vector<CcuRep::Memory> dst;
170 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
171 0 : dst.push_back(CreateMemory());
172 : }
173 :
174 0 : uint32_t dstId = 0;
175 0 : uint32_t curId = 0;
176 :
177 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
178 0 : if (rankIdx != localId_) {
179 0 : curId = dstId;
180 0 : dstId++;
181 : } else {
182 0 : curId = localSize_ - 1;
183 : }
184 0 : dst[curId].addr = output_[rankIdx];
185 0 : dst[curId].addr += firstOutOffset_;
186 0 : dst[curId].token = token_[rankIdx];
187 : }
188 0 : if (axisId_ == 0) {
189 0 : GroupBroadcast(transports, dst, src, goASize_);
190 : } else {
191 0 : GroupBroadcast(transports, dst, src, goBSize_);
192 : }
193 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] firstStep run ends");
194 0 : return;
195 0 : }
196 :
197 0 : void CcuContextAllGatherMesh2D::SecondStep()
198 : {
199 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] secodeStep run begins");
200 0 : uint64_t anotherSize = dimSize_[1 - axisId_];
201 :
202 0 : CcuRep::Memory src = CreateMemory();
203 0 : src.addr = output_[localId_];
204 0 : src.token = token_[localId_];
205 :
206 0 : uint32_t dstId = 0;
207 0 : uint32_t curId = 0;
208 :
209 0 : std::vector<CcuRep::Memory> dst;
210 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
211 0 : dst.push_back(CreateMemory());
212 : }
213 :
214 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
215 0 : if (rankIdx != localId_) {
216 0 : curId = dstId;
217 0 : dstId++;
218 : } else {
219 0 : curId = localSize_ - 1;
220 : }
221 0 : dst[curId].addr = output_[rankIdx];
222 0 : dst[curId].token = token_[rankIdx];
223 : }
224 :
225 0 : CcuRep::Memory tmpSrc = CreateMemory();
226 0 : std::vector<CcuRep::Memory> tmpDst;
227 0 : for (uint32_t r = 0; r < localSize_; r++) {
228 0 : tmpDst.push_back(CreateMemory());
229 : }
230 :
231 0 : for (uint64_t m = 0; m < anotherSize; m++) {
232 0 : if (m == 0) {
233 0 : src.addr += secondInOutBaseOffset_;
234 : } else {
235 0 : src.addr += secondInOutStepOffset_;
236 : }
237 0 : for (uint32_t r = 0; r < localSize_; r++) {
238 0 : if (m == 0) {
239 0 : dst[r].addr += secondInOutBaseOffset_;
240 : } else {
241 0 : dst[r].addr += secondInOutStepOffset_;
242 : }
243 : }
244 :
245 0 : tmpSrc.addr = src.addr;
246 0 : tmpSrc.token = src.token;
247 0 : for (uint32_t r = 0; r < localSize_; r++) {
248 0 : tmpDst[r].addr = dst[r].addr;
249 0 : tmpDst[r].token = dst[r].token;
250 : }
251 :
252 0 : if (axisId_ == 0) {
253 0 : GroupBroadcast(transports, tmpDst, tmpSrc, goBSize_);
254 : } else {
255 0 : GroupBroadcast(transports, tmpDst, tmpSrc, goASize_);
256 : }
257 : }
258 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] secondStep run ends");
259 0 : return;
260 0 : }
261 :
262 0 : void CcuContextAllGatherMesh2D::Algorithm()
263 : {
264 : // 初始化寄存器资源 & 加载外部输入参数
265 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] AllgatherMesh2D Algorithm Init Begins.");
266 0 : InitResources();
267 0 : LoadArgs();
268 :
269 : // 第一轮
270 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] Algorithm first step begins.");
271 0 : ExchangeInfoAndSync();
272 0 : FirstStep();
273 0 : RankSync(CKE_IDX_3);
274 0 : AxisSync(FST_AXIS_ID);
275 :
276 : // 第二轮
277 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] Algorithm second step begins.");
278 0 : RankSync(CKE_IDX_4);
279 0 : SecondStep();
280 0 : RankSync(CKE_IDX_0);
281 0 : AxisSync(SEC_AXIS_ID);
282 :
283 0 : HCCL_INFO("[CcuContextAllGatherMesh2D] Algorithm Ends.");
284 0 : return;
285 : }
286 :
287 0 : std::vector<uint64_t> CcuContextAllGatherMesh2D::GeneArgs(const CcuTaskArg& arg)
288 : {
289 0 : const CcuTaskArgAllGatherMesh2D* taskArg = dynamic_cast<const CcuTaskArgAllGatherMesh2D*>(&arg);
290 0 : if (taskArg == nullptr) {
291 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMesh2D::taskArg ptr is null"));
292 : }
293 :
294 : // input&output&buffer地址
295 0 : uint64_t inputAddr = taskArg->inputAddr_;
296 0 : uint64_t outputAddr = taskArg->outputAddr_;
297 0 : uint64_t offset = taskArg->offset_;
298 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
299 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
300 0 : uint64_t sliceSize = xAxisSize + yAxisSize;
301 0 : uint64_t tokenValue = taskArg->token_;
302 :
303 0 : auto goSizeAxis = CalGoSize(xAxisSize);
304 0 : auto goSizeBxis = CalGoSize(yAxisSize);
305 :
306 : uint64_t firstInOffset;
307 0 : uint64_t firstOutOffset = 0;
308 0 : uint64_t secondInOutBaseOffset = 0;
309 0 : uint64_t secondInOutStepOffset = 0;
310 :
311 0 : for (uint32_t i = 0; i < rankId_; i++) {
312 0 : firstOutOffset += offset;
313 : }
314 :
315 0 : if (axisId_ == 0) {
316 0 : firstInOffset = 0;
317 0 : for (uint32_t i = 0; i < dimId_[0]; i++) {
318 0 : secondInOutBaseOffset += offset;
319 : }
320 0 : secondInOutBaseOffset += xAxisSize;
321 0 : for (uint64_t i = 0; i < dimSize_[0]; i++) {
322 0 : secondInOutStepOffset += offset;
323 : }
324 : } else {
325 0 : firstInOffset = xAxisSize;
326 0 : firstOutOffset += xAxisSize;
327 0 : for (uint32_t i = 0; i < dimId_[1]; i++) {
328 0 : for (uint64_t j = 0; j < dimSize_[0]; j++) {
329 0 : secondInOutBaseOffset += offset;
330 : }
331 : }
332 0 : secondInOutStepOffset = offset;
333 : }
334 :
335 0 : HCCL_INFO(
336 : "[CcuContextAllGatherMesh2D][GeneArgs] RankId[%u]--AxisId[%u], inputAddr[%llu], outputAddr[%llu], "
337 : "aSize[%llu], bSize[%llu], offset[%llu], sliceSize[%llu], firstInOffset[%llu], firstOutOffset[%llu], "
338 : "secondInOutBaseOffset[%llu], secondInOutStepOffset[%llu]",
339 : rankId_, axisId_, inputAddr, outputAddr, xAxisSize, yAxisSize, offset, sliceSize, firstInOffset, firstOutOffset,
340 : secondInOutBaseOffset, secondInOutStepOffset);
341 :
342 : return {
343 : inputAddr,
344 : outputAddr,
345 : tokenValue,
346 : xAxisSize,
347 : yAxisSize,
348 : offset,
349 : sliceSize,
350 : firstInOffset,
351 : firstOutOffset,
352 : secondInOutBaseOffset,
353 : secondInOutStepOffset,
354 0 : goSizeAxis[0],
355 0 : goSizeAxis[1],
356 0 : goSizeAxis[2],
357 0 : goSizeAxis[3],
358 0 : goSizeBxis[0],
359 0 : goSizeBxis[1],
360 0 : goSizeBxis[2],
361 0 : goSizeBxis[3]};
362 0 : }
363 :
364 : } // namespace Hccl
|
