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_reduce_scatter_mesh2d.h"
12 : #include "ccu_instruction_reduce_scatter_mesh2d.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr int CKE_IDX_0 = 0;
17 : constexpr int CKE_IDX_1 = 1;
18 : constexpr int CKE_IDX_2 = 2;
19 : constexpr int CKE_IDX_3 = 3;
20 : constexpr int CKE_IDX_4 = 4;
21 : constexpr int FST_AXIS_ID = 0;
22 : constexpr int SEC_AXIS_ID = 1;
23 : constexpr int INPUT_XN_ID = 1;
24 : constexpr int TOKEN_XN_ID = 2;
25 :
26 0 : CcuContextReduceScatterMesh2D::CcuContextReduceScatterMesh2D(
27 0 : const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
28 0 : : CcuContextAlgBase(arg, transports, group)
29 : {
30 0 : const CcuCtxArgReduceScatterMesh2D* ctxArg = dynamic_cast<const CcuCtxArgReduceScatterMesh2D*>(&arg);
31 0 : if (ctxArg == nullptr) {
32 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh2D::ctxArg ptr is null"));
33 : }
34 0 : rankId_ = ctxArg->rankId_;
35 0 : dimSize_ = ctxArg->dimSize_;
36 0 : axisId_ = ctxArg->axisId_;
37 0 : if (dimSize_[0] == 0) {
38 0 : THROW<InvalidParamsException>(StringFormat("Invalid dimSize[0][%u]", dimSize_[0]));
39 : }
40 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
41 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
42 0 : localId_ = dimId_[axisId_];
43 0 : localSize_ = dimSize_[axisId_];
44 0 : oppsiteSize_ = dimSize_[1 - axisId_];
45 0 : HCCL_INFO(
46 : "[CcuContextReduceScatterMesh2D] RankId[%u], DimSize0[%u], DimSize1[%u], localId[%u], lcoalSize[%u], "
47 : "oppsiteSize[%u]",
48 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_, oppsiteSize_);
49 0 : dataType_ = ctxArg->op_.dataType;
50 0 : outputDataType_ = ctxArg->op_.outputDataType;
51 0 : if (outputDataType_ == DataType::INVALID) {
52 0 : outputDataType_ = dataType_;
53 0 : HCCL_INFO(
54 : "[CcuContextReduceScatterMesh2D] outputDataType is [INVALID], set outputDataType to[%s]",
55 : outputDataType_.Describe().c_str());
56 : }
57 0 : reduceOp_ = ctxArg->op_.reduceOp;
58 0 : localAxisSignalName_ = "CcuContextReduceScatterMesh2DAxisSync_" + std::to_string(axisId_);
59 0 : anotherAxisSignalName_ = "CcuContextReduceScatterMesh2DAxisSync_" + std::to_string(1 - axisId_);
60 0 : }
61 :
62 0 : void CcuContextReduceScatterMesh2D::InitResources()
63 : {
64 0 : step0BaseOffset_ = CreateVariable();
65 0 : step0AddOffset_ = CreateVariable();
66 0 : step1AddOffset_ = CreateVariable();
67 0 : xAxisGroupOpSize_ = CreateGroupOpSize();
68 0 : yAxisGroupOpSize_ = CreateGroupOpSize();
69 0 : localAxisSignal_ = CreateMaskSignal();
70 0 : anotherAxisSignal_ = CreateMaskSignal();
71 0 : yAxisOffset_ = CreateVariable();
72 :
73 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
74 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
75 :
76 0 : output_.push_back(CreateVariable());
77 0 : uint32_t transportIdx = 0;
78 0 : if (transports.size() == 0) {
79 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh2D transports is empty"));
80 : }
81 0 : for (uint64_t peerId = 0; peerId < localSize_; peerId++) {
82 0 : if (peerId == localId_) {
83 0 : input_.push_back(CreateVariable());
84 0 : token_.push_back(CreateVariable());
85 : } else {
86 0 : HCCL_INFO(
87 : "[CcuContextReduceScatterMesh2D] MyRank[%u], PeerId[%llu], TransportId[%u]", localId_, peerId,
88 : transportIdx);
89 0 : CHK_PRT_RET(
90 : transports[transportIdx] == nullptr || transportIdx >= transports.size(),
91 : HCCL_ERROR("[CcuContextReduceScatterMesh2D] Algorithm transport ptr is null or transportIdx is out of "
92 : "bounds"), );
93 0 : input_.push_back(
94 0 : CreateVariable((*transports[transportIdx]), INPUT_XN_ID)); // 获取transport中id=1的Var来传递output
95 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
96 0 : transportIdx++;
97 : }
98 : }
99 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] InitResources finished");
100 : }
101 :
102 0 : void CcuContextReduceScatterMesh2D::PreSync()
103 : {
104 0 : uint16_t selfBit = 1 << localId_;
105 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
106 :
107 0 : for (auto t : transports) {
108 0 : WriteVariableWithSignal(*t, input_[localId_], INPUT_XN_ID, CKE_IDX_1, selfBit); // index = 1，传递output信息
109 0 : WriteVariableWithSignal(*t, token_[localId_], TOKEN_XN_ID, CKE_IDX_2, selfBit); // index = 2，传递token信息
110 : }
111 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit); // index = 1，传递output信息
112 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit); // index = 2，传递token信息
113 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] PreSync run finished");
114 0 : }
115 :
116 0 : void CcuContextReduceScatterMesh2D::PostSync(uint32_t signalIndex)
117 : {
118 0 : uint16_t selfBit = 1 << localId_;
119 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
120 :
121 0 : for (auto t : transports) {
122 0 : RemotePost(*t, signalIndex, selfBit);
123 : }
124 0 : GroupWait(*transportGroup, signalIndex, allBit);
125 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] PostSync run finished");
126 0 : }
127 :
128 0 : void CcuContextReduceScatterMesh2D::AxisSync(uint32_t signalIndex)
129 : {
130 0 : const uint32_t DIE_NUM = 2;
131 0 : if (signalIndex > 1) {
132 0 : THROW<InvalidParamsException>(
133 0 : StringFormat("[CcuContextReduceScatterMesh2D] Unexpected SignalInex[%u]", signalIndex));
134 : }
135 0 : LocalCtxPost(anotherAxisSignal_, 1 << (axisId_ + signalIndex * DIE_NUM));
136 0 : LocalWait(localAxisSignal_, 1 << (1 - axisId_ + signalIndex * DIE_NUM));
137 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] AxisSync run finished");
138 0 : return;
139 : }
140 :
141 0 : void CcuContextReduceScatterMesh2D::LoadArgs()
142 : {
143 0 : Load(input_[localId_]);
144 0 : Load(output_[0]);
145 0 : Load(token_[localId_]);
146 0 : Load(step0BaseOffset_);
147 0 : Load(step0AddOffset_);
148 0 : Load(step1AddOffset_);
149 0 : Load(yAxisOffset_);
150 0 : Load(xAxisGroupOpSize_);
151 0 : Load(yAxisGroupOpSize_);
152 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] LoadArgs run finished");
153 0 : }
154 :
155 0 : void CcuContextReduceScatterMesh2D::Step1Reduce()
156 : {
157 0 : std::vector<CcuRep::Memory> src;
158 0 : std::vector<CcuRep::Memory> tempSrc;
159 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
160 0 : src.push_back(CreateMemory());
161 0 : tempSrc.push_back(CreateMemory());
162 : }
163 0 : CcuRep::Memory dst = CreateMemory();
164 0 : CcuRep::Memory tempDst = CreateMemory();
165 0 : uint32_t dstId = 0;
166 0 : uint32_t curId = 0;
167 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
168 0 : if (localIdx != localId_) {
169 0 : curId = dstId;
170 0 : dstId++;
171 : } else {
172 0 : curId = localSize_ - 1;
173 : }
174 0 : src[curId].addr = input_[localIdx];
175 0 : src[curId].token = token_[localIdx];
176 : }
177 0 : dst.addr = input_[localId_];
178 0 : dst.token = token_[localId_];
179 0 : for (uint32_t oppsiteIdx = 0; oppsiteIdx < oppsiteSize_; oppsiteIdx++) {
180 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
181 0 : if (oppsiteIdx == 0) {
182 0 : src[localIdx].addr += step0BaseOffset_;
183 : } else {
184 0 : src[localIdx].addr += step0AddOffset_;
185 : }
186 : }
187 0 : if (oppsiteIdx == 0) {
188 0 : dst.addr += step0BaseOffset_;
189 : } else {
190 0 : dst.addr += step0AddOffset_;
191 : }
192 0 : tempDst.addr = dst.addr;
193 0 : tempDst.token = dst.token;
194 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
195 0 : tempSrc[rankIdx].addr = src[rankIdx].addr;
196 0 : tempSrc[rankIdx].token = src[rankIdx].token;
197 : }
198 0 : if (axisId_ == 0) {
199 0 : GroupReduce(transports, tempDst, tempSrc, xAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
200 : } else {
201 0 : GroupReduce(transports, tempDst, tempSrc, yAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
202 : }
203 : }
204 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Step1Reduce run finished");
205 0 : }
206 :
207 0 : void CcuContextReduceScatterMesh2D::Step2Reduce()
208 : {
209 0 : std::vector<CcuRep::Memory> src;
210 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
211 0 : src.push_back(CreateMemory());
212 : }
213 0 : CcuRep::Memory dst = CreateMemory();
214 0 : dst.addr = output_[0];
215 0 : dst.token = token_[localId_];
216 0 : uint32_t dstId = 0;
217 0 : uint32_t curId = 0;
218 0 : for (uint16_t localIdx = 0; localIdx < localSize_; localIdx++) {
219 0 : if (localIdx != localId_) {
220 0 : curId = dstId;
221 0 : dstId++;
222 : } else {
223 0 : curId = localSize_ - 1;
224 : }
225 0 : src[curId].addr = input_[localIdx];
226 0 : src[curId].addr += step1AddOffset_;
227 0 : src[curId].token = token_[localIdx];
228 : }
229 0 : if (axisId_ == 0) {
230 0 : dst.addr += yAxisOffset_;
231 0 : GroupReduce(transports, dst, src, yAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
232 : } else {
233 0 : GroupReduce(transports, dst, src, xAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
234 : }
235 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Step2Reduce run finished");
236 0 : }
237 :
238 0 : void CcuContextReduceScatterMesh2D::Algorithm()
239 : {
240 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] ReduceScatterMesh2D run");
241 :
242 0 : InitResources();
243 0 : LoadArgs();
244 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Algorithm first step begins.");
245 0 : PreSync();
246 :
247 0 : Step1Reduce();
248 0 : PostSync(CKE_IDX_3);
249 0 : AxisSync(FST_AXIS_ID);
250 :
251 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Algorithm second step begins.");
252 0 : PostSync(CKE_IDX_4);
253 :
254 0 : Step2Reduce();
255 0 : PostSync(CKE_IDX_0);
256 0 : AxisSync(SEC_AXIS_ID);
257 :
258 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] ReduceScatterMesh2D end");
259 0 : return;
260 : }
261 :
262 0 : std::vector<uint64_t> CcuContextReduceScatterMesh2D::GeneArgs(const CcuTaskArg& arg)
263 : {
264 0 : const CcuTaskArgReduceScatterMesh2D* taskArg = dynamic_cast<const CcuTaskArgReduceScatterMesh2D*>(&arg);
265 0 : if (taskArg == nullptr) {
266 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh2D::taskArg ptr is null"));
267 : }
268 0 : uint64_t inputAddr = taskArg->inputAddr_;
269 0 : uint64_t outputAddr = taskArg->outputAddr_;
270 0 : uint64_t tokenInfo = taskArg->token_;
271 0 : uint64_t outputSize = taskArg->outputSize_;
272 0 : uint64_t offset = taskArg->offset_;
273 0 : uint64_t yAxisOffset = taskArg->xAxisSize_;
274 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
275 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
276 :
277 : // 计算不同die的数据
278 : uint64_t step0BaseOffset
279 0 : = axisId_ == 0 ? dimId_[0] * outputSize + offset : dimId_[1] * dimSize_[0] * outputSize + offset + xAxisSize;
280 0 : uint64_t step0AddOffset = axisId_ == 0 ? dimSize_[0] * outputSize : outputSize;
281 0 : uint64_t step1AddOffset = rankId_ * outputSize + offset + (axisId_ == 0 ? xAxisSize : 0);
282 0 : auto xAxisGoSize = CalGoSize(xAxisSize);
283 0 : auto yAxisGoSize = CalGoSize(yAxisSize);
284 0 : HCCL_INFO(
285 : "[CcuContextReduceScatterMesh2D] GeneArgs: inputAddr[%llu], outputAddr[%llu],"
286 : "step0BaseOffset[%llu], step0AddOffset[%llu], step1AddOffset[%llu]",
287 : inputAddr, outputAddr, step0BaseOffset, step0AddOffset, step1AddOffset);
288 : return {inputAddr, outputAddr, tokenInfo, step0BaseOffset, step0AddOffset,
289 0 : step1AddOffset, yAxisOffset, xAxisGoSize[0], xAxisGoSize[1], xAxisGoSize[2],
290 0 : xAxisGoSize[3], yAxisGoSize[0], yAxisGoSize[1], yAxisGoSize[2], yAxisGoSize[3]};
291 0 : }
292 : } // namespace Hccl
|
