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