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