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_mem2mem.h"
12 : #include "ccu_instruction_all_gather_mesh2d_mem2mem.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; // 约定：尾同步均使用第0个CKE
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 : CcuContextAllGatherMeshMem2Mem2D::CcuContextAllGatherMeshMem2Mem2D(const CcuCtxArg &arg,
27 : const std::vector<CcuTransport *> &transports,
28 0 : const CcuTransportGroup &group)
29 0 : : CcuContextAlgBase(arg, transports, group)
30 : {
31 0 : xAxisSize_ = CreateVariable();
32 0 : yAxisSize_ = CreateVariable();
33 0 : offset_ = CreateVariable();
34 0 : sliceSize_ = CreateVariable();
35 0 : firstOutOffset_ = CreateVariable();
36 0 : firstInOffset_ = CreateVariable();
37 0 : secondInOutBaseOffset_ = CreateVariable();
38 0 : secondInOutStepOffset_ = CreateVariable();
39 0 : goASize_ = CreateGroupOpSize();
40 0 : goBSize_ = CreateGroupOpSize();
41 0 : localAxisSignal_ = CreateMaskSignal();
42 :
43 0 : const CcuCtxArgAllGatherMeshMem2Mem2D *ctxArg = dynamic_cast<const CcuCtxArgAllGatherMeshMem2Mem2D *>(&arg);
44 0 : if (ctxArg == nullptr) {
45 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMeshMem2Mem2D::ctxArg ptr is null"));
46 : }
47 0 : rankId_ = ctxArg->rankId_;
48 0 : dimSize_ = ctxArg->dimSize_;
49 0 : axisId_ = ctxArg->axisId_;
50 : // DIM_SIZE由头文件传入，=2
51 0 : if (dimSize_.size() != DIM_SIZE or axisId_ > 1) {
52 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMeshMem2Mem2D::dimSize[%u] or axisId[%u] is invalid",
53 : dimSize_.size(), axisId_));
54 : }
55 0 : CHK_PRT_THROW(dimSize_[0] == 0 || dimSize_[1] == 0,
56 : HCCL_ERROR("[CcuContextAllGatherMeshMem2Mem2D] dimSize0[%llu] or dimSize1[%llu] is zero", dimSize_[0],
57 : dimSize_[1]),
58 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
59 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
60 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
61 0 : localId_ = dimId_[axisId_];
62 0 : localSize_ = dimSize_[axisId_];
63 0 : localAxisSignalName_ = "CcuContextAllGatherMeshMem2Mem2DAxisSync_" + std::to_string(axisId_);
64 0 : anotherAxisSignalName_ = "CcuContextAllGatherMeshMem2Mem2DAxisSync_" + std::to_string(1 - axisId_);
65 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] RankId[%u], DimSize: D0[%u]--D1[%u], localId[%u], lcoalSize[%u]",
66 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_);
67 0 : }
68 :
69 0 : void CcuContextAllGatherMeshMem2Mem2D::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("[CcuContextAllGatherMeshMem2Mem2D] MyRank[%u], peerId[%u], transportIdx[%u]", rankId_, peerId,
80 : transportIdx);
81 0 : CHK_PRT_RET(transports[transportIdx] == nullptr,
82 : HCCL_ERROR("[CcuContextAllGatherMeshMem2Mem2D] Algorithm transport ptr is null"), );
83 0 : output_.push_back(CreateVariable((*transports[transportIdx]), OUTPUT_XN_ID));
84 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
85 0 : transportIdx++;
86 : }
87 : }
88 :
89 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
90 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
91 0 : return;
92 : }
93 :
94 : // 需要与GeneArgs返回时的顺序相同
95 0 : void CcuContextAllGatherMeshMem2Mem2D::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 CcuContextAllGatherMeshMem2Mem2D::ExchangeInfoAndSync()
115 : {
116 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] 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("CcuContextAllGatherMeshMem2Mem2D::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("[CcuContextAllGatherMeshMem2Mem2D] change addr success");
127 : }
128 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit);
129 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit);
130 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] ExchangeInfoAndSync run finished.");
131 0 : return;
132 : }
133 :
134 0 : void CcuContextAllGatherMeshMem2Mem2D::RankSync(uint32_t signalIndex)
135 : {
136 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] 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("CcuContextAllGatherMeshMem2Mem2D::Algorithm transport ptr is null"));
143 : }
144 0 : RemotePost(*t, signalIndex, selfBit);
145 : }
146 0 : GroupWait(*transportGroup, signalIndex, allBit);
147 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] RankSync run ends.");
148 0 : return;
149 : }
150 :
151 0 : void CcuContextAllGatherMeshMem2Mem2D::AxisSync(uint32_t signalIndex)
152 : {
153 0 : constexpr uint32_t DIE_NUM = 2;
154 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] 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("[CcuContextAllGatherMeshMem2Mem2D] AxisSync run ends.");
158 0 : return;
159 : }
160 :
161 0 : void CcuContextAllGatherMeshMem2Mem2D::FirstStep()
162 : {
163 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] 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 0 : CcuRep::Variable tempSize = CreateVariable();
174 0 : tempSize = (axisId_ == 0) ? xAxisSize_ : yAxisSize_;
175 0 : GroupOpSize tempGoSize = CreateGroupOpSize();
176 0 : tempGoSize = (axisId_ == 0) ? goASize_ : goBSize_;
177 0 : CCU_IF(tempSize != 0) {
178 0 : u32 transportId = 0;
179 0 : CcuRep::MaskSignal locMask = CreateMaskSignal();
180 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
181 0 : dst[rankIdx].addr = output_[rankIdx];
182 0 : dst[rankIdx].addr += firstOutOffset_;
183 0 : dst[rankIdx].token = token_[rankIdx];
184 0 : if (rankIdx == localId_) {
185 0 : LocalPost(locMask, (1 << rankIdx));
186 : } else {
187 0 : Write(*transports[transportId], dst[rankIdx], src, tempSize, locMask, 1 << rankIdx);
188 0 : transportId++;
189 : }
190 : }
191 0 : GroupCopy(dst[localId_], src, tempGoSize);
192 0 : LocalWait(locMask, (1 << localSize_) - 1);
193 0 : }
194 :
195 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] firstStep run ends.");
196 0 : }
197 :
198 0 : void CcuContextAllGatherMeshMem2Mem2D::SecondStep()
199 : {
200 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] secodeStep run begins.");
201 0 : uint64_t anotherSize = dimSize_[1 - axisId_];
202 :
203 0 : CcuRep::Memory src = CreateMemory();
204 0 : src.addr = output_[localId_];
205 0 : src.token = token_[localId_];
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 : dst[rankIdx].addr = output_[rankIdx];
214 0 : dst[rankIdx].token = token_[rankIdx];
215 : }
216 :
217 0 : for (uint64_t localIdx = 0; localIdx < anotherSize; localIdx++) { // 分片数等于另外一个轴的大小
218 0 : if (localIdx == 0) {
219 0 : src.addr += secondInOutBaseOffset_;
220 : } else {
221 0 : src.addr += secondInOutStepOffset_;
222 : }
223 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
224 0 : if (localIdx == 0) {
225 0 : dst[rankIdx].addr += secondInOutBaseOffset_;
226 : } else {
227 0 : dst[rankIdx].addr += secondInOutStepOffset_;
228 : }
229 : }
230 0 : auto tempSize = (axisId_ == 0) ? yAxisSize_ : xAxisSize_;
231 0 : CCU_IF(tempSize != 0)
232 : {
233 0 : u32 transportId = 0;
234 0 : CcuRep::MaskSignal locMask = CreateMaskSignal();
235 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
236 0 : if (rankIdx == localId_) {
237 0 : LocalPost(locMask, (1 << rankIdx));
238 : } else {
239 0 : Write(*transports[transportId], dst[rankIdx], src, tempSize, locMask, 1 << rankIdx);
240 0 : transportId++;
241 : }
242 : }
243 0 : LocalWait(locMask, (1 << localSize_) - 1);
244 0 : }
245 0 : }
246 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] secondStep run ends.");
247 0 : }
248 :
249 0 : void CcuContextAllGatherMeshMem2Mem2D::Algorithm()
250 : {
251 : // 初始化寄存器资源 & 加载外部输入参数
252 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] AllgatherMesh2D Algorithm Init Begins.");
253 0 : InitResources();
254 0 : LoadArgs();
255 :
256 : // 第一轮
257 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] Algorithm first step begins.");
258 0 : ExchangeInfoAndSync();
259 0 : FirstStep();
260 0 : RankSync(CKE_IDX_3);
261 0 : AxisSync(FST_AXIS_ID);
262 :
263 : // 第二轮
264 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] Algorithm second step begins.");
265 0 : RankSync(CKE_IDX_4);
266 0 : SecondStep();
267 0 : RankSync(CKE_IDX_0);
268 0 : AxisSync(SEC_AXIS_ID);
269 :
270 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D] Algorithm Ends.");
271 0 : return;
272 : }
273 :
274 0 : std::vector<uint64_t> CcuContextAllGatherMeshMem2Mem2D::GeneArgs(const CcuTaskArg &arg)
275 : {
276 0 : const CcuTaskArgAllGatherMeshMem2Mem2D *taskArg = dynamic_cast<const CcuTaskArgAllGatherMeshMem2Mem2D *>(&arg);
277 0 : if (taskArg == nullptr) {
278 0 : THROW<NullPtrException>(StringFormat("CcuContextAllGatherMeshMem2Mem2D::taskArg ptr is null"));
279 : }
280 :
281 : // input&output&buffer地址
282 0 : uint64_t inputAddr = taskArg->inputAddr_;
283 0 : uint64_t outputAddr = taskArg->outputAddr_;
284 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
285 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
286 0 : uint64_t offset = taskArg->offset_;
287 0 : uint64_t sliceSize = xAxisSize + yAxisSize;
288 0 : uint64_t tokenValue = taskArg->token_;
289 :
290 0 : auto goSizeAxis = CalGoSize(xAxisSize);
291 0 : auto goSizeBxis = CalGoSize(yAxisSize);
292 :
293 : uint64_t firstInOffset;
294 0 : uint64_t firstOutOffset = 0;
295 0 : uint64_t secondInOutBaseOffset = 0;
296 0 : uint64_t secondInOutStepOffset = 0;
297 :
298 0 : for (uint32_t i = 0; i < rankId_; i++) {
299 0 : firstOutOffset += offset;
300 : }
301 :
302 0 : if (axisId_ == 0) {
303 0 : firstInOffset = 0;
304 0 : for (uint32_t i = 0; i < dimId_[0]; i++) {
305 0 : secondInOutBaseOffset += offset;
306 : }
307 0 : secondInOutBaseOffset += xAxisSize;
308 0 : for (uint64_t i = 0; i < dimSize_[0]; i++) {
309 0 : secondInOutStepOffset += offset;
310 : }
311 : } else {
312 0 : firstInOffset = xAxisSize;
313 0 : firstOutOffset += xAxisSize;
314 0 : for (uint32_t i = 0; i < dimId_[1]; i++) {
315 0 : for (uint64_t j = 0; j < dimSize_[0]; j++) {
316 0 : secondInOutBaseOffset += offset;
317 : }
318 : }
319 0 : secondInOutStepOffset = offset;
320 : }
321 :
322 0 : HCCL_INFO("[CcuContextAllGatherMeshMem2Mem2D][GeneArgs] RankId[%u]--AxisId[%u], inputAddr[%llu], outputAddr[%llu], "
323 : "aSize[%llu], bSize[%llu], offset[%llu], sliceSize[%llu], firstInOffset[%llu], firstOutOffset[%llu], "
324 : "secondInOutBaseOffset[%llu], secondInOutStepOffset[%llu]", rankId_, axisId_, inputAddr, outputAddr,
325 : xAxisSize, yAxisSize, offset, sliceSize, firstInOffset, firstOutOffset, secondInOutBaseOffset,
326 : secondInOutStepOffset);
327 :
328 : return {inputAddr, outputAddr, tokenValue, xAxisSize, yAxisSize, offset, sliceSize, firstInOffset, firstOutOffset,
329 0 : secondInOutBaseOffset, secondInOutStepOffset, goSizeAxis[0], goSizeAxis[1], goSizeAxis[2], goSizeAxis[3],
330 0 : goSizeBxis[0], goSizeBxis[1], goSizeBxis[2], goSizeBxis[3]};
331 0 : }
332 :
333 : } // namespace Hccl
|
