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