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