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_broadcast_mesh2d_mem2mem.h"
12 : #include "ccu_instruction_broadcast_mesh2d_mem2mem.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr int INPUT_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 : constexpr int X_AXIS_ID = 0;
26 : constexpr int Y_AXIS_ID = 1;
27 :
28 0 : CcuContextBroadcastMeshMem2Mem2D::CcuContextBroadcastMeshMem2Mem2D(const CcuCtxArg &arg,
29 : const std::vector<CcuTransport *> &transports,
30 0 : const CcuTransportGroup &group)
31 0 : : CcuContextAlgBase(arg, transports, group)
32 : {
33 0 : const CcuCtxArgBroadcastMeshMem2mem2D *ctxArg = dynamic_cast<const CcuCtxArgBroadcastMeshMem2mem2D *>(&arg);
34 0 : if (ctxArg == nullptr) {
35 0 : THROW<NullPtrException>(StringFormat("CcuContextBroadcastMeshMem2Mem2D::ctxArg ptr is null"));
36 : }
37 0 : rankId_ = ctxArg->rankId_;
38 0 : dimSize_ = ctxArg->dimSize_;
39 0 : axisId_ = ctxArg->axisId_;
40 0 : if (dimSize_.size() != 2 || axisId_ > 1 || dimSize_[0] == 0) { // 2D 拓扑校验
41 0 : THROW<NullPtrException>(
42 0 : StringFormat("[CcuContextBroadcastMeshMem2Mem2D] dimSize[%zu] or axisId[%u] or dimSize[0] [%u] is invalid",
43 0 : dimSize_.size(), axisId_, dimSize_[0]));
44 : }
45 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
46 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
47 0 : localId_ = dimId_[axisId_];
48 0 : localSize_ = dimSize_[axisId_];
49 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] RankId[%u], DimSize0[%u], DimSize1[%u], localId[%u], lcoalSize[%u]",
50 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_);
51 0 : dataType_ = ctxArg->op_.dataType;
52 :
53 0 : rootId_ = ctxArg->op_.root;
54 0 : rootDimId_.emplace_back(rootId_ % dimSize_[0]);
55 0 : rootDimId_.emplace_back(rootId_ / dimSize_[0]);
56 :
57 0 : rootLocalId_ = rootDimId_[axisId_];
58 :
59 0 : localAxisSignalName_ = "CcuContextBroadcastMeshMem2Mem2DAxisSync_" + std::to_string(axisId_);
60 0 : anotherAxisSignalName_ = "CcuContextBroadcastMeshMem2Mem2DAxisSync_" + std::to_string(1 - axisId_);
61 0 : }
62 :
63 0 : void CcuContextBroadcastMeshMem2Mem2D::InitResources()
64 : {
65 0 : localAxisSignal_ = CreateMaskSignal();
66 0 : anotherAxisSignal_ = CreateMaskSignal();
67 0 : xAxisSize_ = CreateVariable();
68 0 : yAxisSize_ = CreateVariable();
69 0 : yAxisOffset_ = CreateVariable();
70 0 : dataSize_ = CreateVariable();
71 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
72 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
73 :
74 0 : uint32_t transportIdx = 0;
75 0 : if (transports.size() == 0) {
76 0 : THROW<NullPtrException>(StringFormat("CcuContextBroadcastMeshMem2Mem2D transports is empty"));
77 : }
78 0 : for (uint32_t peerId = 0; peerId < localSize_; peerId++) {
79 0 : if (peerId == localId_) {
80 0 : input_.push_back(CreateVariable());
81 0 : token_.push_back(CreateVariable());
82 : } else {
83 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] MyRank[%u], PeerId[%u], TransportId[%u]", localId_, peerId,
84 : transportIdx);
85 0 : CHK_PRT_RET(transports[transportIdx] == nullptr,
86 : HCCL_ERROR("[CcuContextBroadcastMeshMem2Mem2D] Algorithm transport ptr is null"), );
87 0 : input_.push_back(
88 0 : CreateVariable((*transports[transportIdx]), INPUT_XN_ID)); // 获取transport中id=1的Var来传递output
89 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
90 0 : transportIdx++;
91 : }
92 : }
93 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] InitResources finished");
94 : }
95 :
96 0 : void CcuContextBroadcastMeshMem2Mem2D::PreSync()
97 : {
98 0 : uint16_t selfBit = 1 << localId_;
99 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
100 :
101 0 : for (auto t : transports) {
102 0 : WriteVariableWithSignal(*t, input_[localId_], INPUT_XN_ID, CKE_IDX_1, selfBit); // index = 1，传递output信息
103 0 : WriteVariableWithSignal(*t, token_[localId_], TOKEN_XN_ID, CKE_IDX_2, selfBit); // index = 2，传递token信息
104 : }
105 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit); // index = 1，传递output信息
106 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit); // index = 2，传递token信息
107 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] PreSync run finished");
108 0 : }
109 :
110 0 : void CcuContextBroadcastMeshMem2Mem2D::PostSync(uint32_t signalIndex)
111 : {
112 0 : uint16_t selfBit = 1 << localId_;
113 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
114 :
115 0 : for (auto t : transports) {
116 0 : RemotePost(*t, signalIndex, selfBit);
117 : }
118 0 : GroupWait(*transportGroup, signalIndex, allBit);
119 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] PostSync run finished");
120 0 : }
121 :
122 0 : void CcuContextBroadcastMeshMem2Mem2D::AxisSync(uint32_t signalIndex)
123 : {
124 0 : const uint32_t DIE_NUM = 2;
125 0 : if (signalIndex > 1) {
126 0 : THROW<InvalidParamsException>(
127 0 : StringFormat("[CcuContextBroadcastMeshMem2Mem2D] Unexpected SignalInex[%u]", signalIndex));
128 : }
129 0 : LocalCtxPost(anotherAxisSignal_, 1 << (axisId_ + signalIndex * DIE_NUM));
130 0 : LocalWait(localAxisSignal_, 1 << (1 - axisId_ + signalIndex * DIE_NUM));
131 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] AxisSync run finished");
132 0 : return;
133 : }
134 :
135 0 : void CcuContextBroadcastMeshMem2Mem2D::LoadArgs()
136 : {
137 0 : Load(input_[localId_]);
138 0 : Load(yAxisOffset_);
139 0 : Load(token_[localId_]);
140 0 : Load(xAxisSize_);
141 0 : Load(yAxisSize_);
142 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] LoadArgs run finished");
143 0 : }
144 :
145 0 : void CcuContextBroadcastMeshMem2Mem2D::Step1BroadcastForRoot()
146 : {
147 0 : std::vector<CcuRep::Memory> dst;
148 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
149 0 : dst.push_back(CreateMemory());
150 : }
151 0 : CcuRep::Memory src = CreateMemory();
152 0 : src.addr = input_[localId_];
153 0 : src.token = token_[localId_];
154 :
155 0 : dataSize_ = (axisId_ == X_AXIS_ID) ? xAxisSize_ : yAxisSize_;
156 0 : CCU_IF(dataSize_ != 0)
157 : {
158 0 : uint32_t transportId = 0;
159 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
160 :
161 0 : if (axisId_ == Y_AXIS_ID) {
162 0 : src.addr += yAxisOffset_;
163 : }
164 :
165 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
166 0 : if (rankIdx == localId_) {
167 0 : LocalPost(localMask, 1 << rankIdx);
168 0 : continue;
169 : }
170 0 : dst[rankIdx].addr = input_[rankIdx];
171 0 : dst[rankIdx].token = token_[rankIdx];
172 0 : if (axisId_ == Y_AXIS_ID) {
173 0 : dst[rankIdx].addr += yAxisOffset_;
174 : }
175 :
176 0 : Write(*transports[transportId], dst[rankIdx], src, dataSize_, localMask, 1 << rankIdx);
177 0 : transportId++;
178 : }
179 0 : LocalWait(localMask, (1 << localSize_) - 1);
180 0 : }
181 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] Step1BroadcastForRoot run finished");
182 0 : }
183 :
184 0 : void CcuContextBroadcastMeshMem2Mem2D::Step2BroadcastForRoot()
185 : {
186 0 : std::vector<CcuRep::Memory> dst;
187 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
188 0 : dst.push_back(CreateMemory());
189 : }
190 0 : CcuRep::Memory src = CreateMemory();
191 0 : src.addr = input_[localId_];
192 0 : src.token = token_[localId_];
193 :
194 0 : dataSize_ = (axisId_ == Y_AXIS_ID) ? xAxisSize_ : yAxisSize_;
195 0 : CCU_IF(dataSize_ != 0)
196 : {
197 0 : uint32_t transportId = 0;
198 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
199 0 : if (axisId_ == X_AXIS_ID) {
200 0 : src.addr += yAxisOffset_;
201 : }
202 :
203 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
204 0 : if (rankIdx == localId_) {
205 0 : LocalPost(localMask, 1 << rankIdx);
206 0 : continue;
207 : }
208 :
209 0 : dst[rankIdx].addr = input_[rankIdx];
210 0 : dst[rankIdx].token = token_[rankIdx];
211 :
212 0 : if (axisId_ == X_AXIS_ID) {
213 0 : dst[rankIdx].addr += yAxisOffset_;
214 : }
215 :
216 0 : Write(*transports[transportId], dst[rankIdx], src, dataSize_, localMask, 1 << rankIdx);
217 0 : transportId++;
218 : }
219 0 : LocalWait(localMask, (1 << localSize_) - 1);
220 0 : }
221 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] Step2BroadcastForRoot run finished");
222 0 : }
223 :
224 0 : void CcuContextBroadcastMeshMem2Mem2D::Step2Broadcast()
225 : {
226 0 : std::vector<CcuRep::Memory> dst;
227 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
228 0 : dst.push_back(CreateMemory());
229 : }
230 0 : CcuRep::Memory src = CreateMemory();
231 0 : src.addr = input_[localId_];
232 0 : src.token = token_[localId_];
233 : // 当前rank和root在同一行, 只有y轴需要再执行一次bcast; 当前rank和root在同列, 只有x轴需要再执行一次bcast
234 0 : if (dimId_[1] == rootDimId_[1] && axisId_ == Y_AXIS_ID) {
235 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] rankId[%u] is on the same row as rootId[%u]", rankId_, rootId_);
236 0 : CCU_IF(xAxisSize_ != 0)
237 : {
238 0 : u32 transportId = 0;
239 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
240 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
241 0 : if (rankIdx == localId_) {
242 0 : LocalPost(localMask, 1 << rankIdx);
243 0 : continue;
244 : }
245 0 : dst[rankIdx].addr = input_[rankIdx];
246 0 : dst[rankIdx].token = token_[rankIdx];
247 0 : Write(*transports[transportId], dst[rankIdx], src, xAxisSize_, localMask, 1 << rankIdx);
248 0 : transportId++;
249 : }
250 0 : LocalWait(localMask, (1 << localSize_) - 1);
251 0 : }
252 0 : } else if (dimId_[0] == rootDimId_[0] && axisId_ == X_AXIS_ID) {
253 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] rankId[%u] is on the same column as rootId[%u]", rankId_, rootId_);
254 0 : CCU_IF(yAxisSize_ != 0)
255 : {
256 0 : src.addr += yAxisOffset_;
257 0 : u32 transportId = 0;
258 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
259 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
260 0 : if (rankIdx == localId_) {
261 0 : LocalPost(localMask, 1 << rankIdx);
262 0 : continue;
263 : }
264 0 : dst[rankIdx].addr = input_[rankIdx];
265 0 : dst[rankIdx].addr += yAxisOffset_;
266 0 : dst[rankIdx].token = token_[rankIdx];
267 0 : Write(*transports[transportId], dst[rankIdx], src, yAxisSize_, localMask, 1 << rankIdx);
268 0 : transportId++;
269 : }
270 0 : LocalWait(localMask, (1 << localSize_) - 1);
271 0 : }
272 : }
273 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] Step2Broadcast run finished");
274 0 : }
275 :
276 0 : void CcuContextBroadcastMeshMem2Mem2D::Algorithm()
277 : {
278 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] BroadcastMesh2D run");
279 0 : InitResources();
280 0 : LoadArgs();
281 :
282 0 : PreSync();
283 :
284 0 : if (rankId_ == rootId_) {
285 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] Algorithm root run begins.");
286 0 : Step1BroadcastForRoot();
287 0 : PostSync(CKE_IDX_3);
288 0 : AxisSync(FST_AXIS_ID);
289 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] Algorithm second step begins.");
290 0 : PostSync(CKE_IDX_4);
291 0 : Step2BroadcastForRoot();
292 0 : } else if (dimId_[0] == rootDimId_[0] || dimId_[1] == rootDimId_[1]) { // 当前rank和root在同一行或同一列
293 0 : if (dimId_[axisId_] != rootDimId_[axisId_]) {
294 0 : PostSync(CKE_IDX_3);
295 : }
296 0 : AxisSync(FST_AXIS_ID);
297 0 : if (dimId_[axisId_] != rootDimId_[axisId_]) {
298 0 : PostSync(CKE_IDX_4);
299 : }
300 0 : Step2Broadcast();
301 : } else {
302 0 : AxisSync(FST_AXIS_ID);
303 : }
304 0 : PostSync(CKE_IDX_0);
305 0 : AxisSync(SEC_AXIS_ID);
306 :
307 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D] BroadcastMesh2D end");
308 0 : return;
309 : }
310 :
311 0 : std::vector<uint64_t> CcuContextBroadcastMeshMem2Mem2D::GeneArgs(const CcuTaskArg &arg)
312 : {
313 0 : const CcuTaskArgBroadcastMeshMem2mem2D *taskArg = dynamic_cast<const CcuTaskArgBroadcastMeshMem2mem2D *>(&arg);
314 0 : if (taskArg == nullptr) {
315 0 : THROW<NullPtrException>(StringFormat("CcuContextBroadcastMeshMem2Mem2D::taskArg ptr is null"));
316 : }
317 0 : uint64_t inputAddr = taskArg->inputAddr_;
318 0 : uint64_t tokenInfo = taskArg->token_;
319 0 : uint64_t yAxisOffset = taskArg->xAxisSize_;
320 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
321 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
322 0 : auto xAxisGoSize = CalGoSize(xAxisSize);
323 0 : auto yAxisGoSize = CalGoSize(yAxisSize);
324 0 : HCCL_INFO("[CcuContextBroadcastMeshMem2Mem2D][GeneArgs] RankId[%u]--AxisId[%u], inputAddr[%llu],xAxisSize[%llu], "
325 : "yAxisSize[%llu],yAxisOffset[%llu]",
326 : rankId_, axisId_, inputAddr, xAxisSize, yAxisSize, yAxisOffset);
327 0 : return {inputAddr, yAxisOffset, tokenInfo, xAxisSize, yAxisSize};
328 0 : }
329 : } // namespace Hccl
|
