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(const CcuCtxArg &arg, const std::vector<CcuTransport*> &transports,
27 0 : 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(
39 0 : "Invalid dimSize[0][%u]", dimSize_[0]));
40 : }
41 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
42 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
43 0 : localId_ = dimId_[axisId_];
44 0 : localSize_ = dimSize_[axisId_];
45 0 : oppsiteSize_ = dimSize_[1 - axisId_];
46 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] RankId[%u], DimSize0[%u], DimSize1[%u], localId[%u], lcoalSize[%u], oppsiteSize[%u]",
47 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_, oppsiteSize_);
48 0 : dataType_ = ctxArg->op_.dataType;
49 0 : outputDataType_ = ctxArg->op_.outputDataType;
50 0 : if (outputDataType_ == DataType::INVALID) {
51 0 : outputDataType_ = dataType_;
52 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] outputDataType is [INVALID], set outputDataType to[%s]",
53 : outputDataType_.Describe().c_str());
54 : }
55 0 : reduceOp_ = ctxArg->op_.reduceOp;
56 0 : localAxisSignalName_ = "CcuContextReduceScatterMesh2DAxisSync_" + std::to_string(axisId_);
57 0 : anotherAxisSignalName_ = "CcuContextReduceScatterMesh2DAxisSync_" + std::to_string(1 - axisId_);
58 0 : }
59 :
60 0 : void CcuContextReduceScatterMesh2D::InitResources()
61 : {
62 0 : step0BaseOffset_ = CreateVariable();
63 0 : step0AddOffset_ = CreateVariable();
64 0 : step1AddOffset_ = CreateVariable();
65 0 : xAxisGroupOpSize_ = CreateGroupOpSize();
66 0 : yAxisGroupOpSize_ = CreateGroupOpSize();
67 0 : localAxisSignal_ = CreateMaskSignal();
68 0 : anotherAxisSignal_ = CreateMaskSignal();
69 0 : yAxisOffset_ = CreateVariable();
70 :
71 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
72 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
73 :
74 0 : output_.push_back(CreateVariable());
75 0 : uint32_t transportIdx = 0;
76 0 : if (transports.size() == 0) {
77 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh2D transports is empty"));
78 : }
79 0 : for (uint64_t peerId = 0; peerId < localSize_; peerId++) {
80 0 : if (peerId == localId_) {
81 0 : input_.push_back(CreateVariable());
82 0 : token_.push_back(CreateVariable());
83 : } else {
84 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] MyRank[%u], PeerId[%llu], TransportId[%u]",
85 : localId_, peerId, transportIdx);
86 0 : CHK_PRT_RET(transports[transportIdx] == nullptr || transportIdx >= transports.size(),
87 : HCCL_ERROR("[CcuContextReduceScatterMesh2D] Algorithm transport ptr is null or transportIdx is out of bounds"),);
88 0 : input_.push_back(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("[CcuContextReduceScatterMesh2D] InitResources finished");
94 : }
95 :
96 0 : void CcuContextReduceScatterMesh2D::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("[CcuContextReduceScatterMesh2D] PreSync run finished");
108 0 : }
109 :
110 0 : void CcuContextReduceScatterMesh2D::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("[CcuContextReduceScatterMesh2D] PostSync run finished");
120 0 : }
121 :
122 0 : void CcuContextReduceScatterMesh2D::AxisSync(uint32_t signalIndex)
123 : {
124 0 : const uint32_t DIE_NUM = 2;
125 0 : if (signalIndex > 1) {
126 0 : THROW<InvalidParamsException>(StringFormat(
127 : "[CcuContextReduceScatterMesh2D] 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("[CcuContextReduceScatterMesh2D] AxisSync run finished");
132 0 : return;
133 : }
134 :
135 0 : void CcuContextReduceScatterMesh2D::LoadArgs()
136 : {
137 0 : Load(input_[localId_]);
138 0 : Load(output_[0]);
139 0 : Load(token_[localId_]);
140 0 : Load(step0BaseOffset_);
141 0 : Load(step0AddOffset_);
142 0 : Load(step1AddOffset_);
143 0 : Load(yAxisOffset_);
144 0 : Load(xAxisGroupOpSize_);
145 0 : Load(yAxisGroupOpSize_);
146 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] LoadArgs run finished");
147 0 : }
148 :
149 0 : void CcuContextReduceScatterMesh2D::Step1Reduce()
150 : {
151 0 : std::vector<CcuRep::Memory> src;
152 0 : std::vector<CcuRep::Memory> tempSrc;
153 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
154 0 : src.push_back(CreateMemory());
155 0 : tempSrc.push_back(CreateMemory());
156 : }
157 0 : CcuRep::Memory dst = CreateMemory();
158 0 : CcuRep::Memory tempDst = CreateMemory();
159 0 : uint32_t dstId = 0;
160 0 : uint32_t curId = 0;
161 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
162 0 : if (localIdx != localId_) {
163 0 : curId = dstId;
164 0 : dstId++;
165 : } else {
166 0 : curId = localSize_ - 1;
167 : }
168 0 : src[curId].addr = input_[localIdx];
169 0 : src[curId].token = token_[localIdx];
170 : }
171 0 : dst.addr = input_[localId_];
172 0 : dst.token = token_[localId_];
173 0 : for (uint32_t oppsiteIdx = 0; oppsiteIdx < oppsiteSize_; oppsiteIdx++) {
174 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
175 0 : if (oppsiteIdx == 0) {
176 0 : src[localIdx].addr += step0BaseOffset_;
177 : } else {
178 0 : src[localIdx].addr += step0AddOffset_;
179 : }
180 : }
181 0 : if (oppsiteIdx == 0) {
182 0 : dst.addr += step0BaseOffset_;
183 : } else {
184 0 : dst.addr += step0AddOffset_;
185 : }
186 0 : tempDst.addr = dst.addr;
187 0 : tempDst.token = dst.token;
188 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
189 0 : tempSrc[rankIdx].addr = src[rankIdx].addr;
190 0 : tempSrc[rankIdx].token = src[rankIdx].token;
191 : }
192 0 : if (axisId_ == 0) {
193 0 : GroupReduce(transports, tempDst, tempSrc, xAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
194 : } else {
195 0 : GroupReduce(transports, tempDst, tempSrc, yAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
196 : }
197 : }
198 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Step1Reduce run finished");
199 0 : }
200 :
201 0 : void CcuContextReduceScatterMesh2D::Step2Reduce()
202 : {
203 0 : std::vector<CcuRep::Memory> src;
204 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
205 0 : src.push_back(CreateMemory());
206 : }
207 0 : CcuRep::Memory dst = CreateMemory();
208 0 : dst.addr = output_[0];
209 0 : dst.token = token_[localId_];
210 0 : uint32_t dstId = 0;
211 0 : uint32_t curId = 0;
212 0 : for (uint16_t localIdx = 0; localIdx < localSize_; localIdx++) {
213 0 : if (localIdx != localId_) {
214 0 : curId = dstId;
215 0 : dstId++;
216 : } else {
217 0 : curId = localSize_ - 1;
218 : }
219 0 : src[curId].addr = input_[localIdx];
220 0 : src[curId].addr += step1AddOffset_;
221 0 : src[curId].token = token_[localIdx];
222 : }
223 0 : if (axisId_ == 0) {
224 0 : dst.addr += yAxisOffset_;
225 0 : GroupReduce(transports, dst, src, yAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
226 : } else {
227 0 : GroupReduce(transports, dst, src, xAxisGroupOpSize_, dataType_, outputDataType_, reduceOp_);
228 : }
229 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Step2Reduce run finished");
230 0 : }
231 :
232 0 : void CcuContextReduceScatterMesh2D::Algorithm()
233 : {
234 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] ReduceScatterMesh2D run");
235 :
236 0 : InitResources();
237 0 : LoadArgs();
238 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Algorithm first step begins.");
239 0 : PreSync();
240 :
241 0 : Step1Reduce();
242 0 : PostSync(CKE_IDX_3);
243 0 : AxisSync(FST_AXIS_ID);
244 :
245 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] Algorithm second step begins.");
246 0 : PostSync(CKE_IDX_4);
247 :
248 0 : Step2Reduce();
249 0 : PostSync(CKE_IDX_0);
250 0 : AxisSync(SEC_AXIS_ID);
251 :
252 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] ReduceScatterMesh2D end");
253 0 : return;
254 : }
255 :
256 0 : std::vector<uint64_t> CcuContextReduceScatterMesh2D::GeneArgs(const CcuTaskArg &arg)
257 : {
258 0 : const CcuTaskArgReduceScatterMesh2D *taskArg = dynamic_cast<const CcuTaskArgReduceScatterMesh2D *>(&arg);
259 0 : if (taskArg == nullptr) {
260 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh2D::taskArg ptr is null"));
261 : }
262 0 : uint64_t inputAddr = taskArg->inputAddr_;
263 0 : uint64_t outputAddr = taskArg->outputAddr_;
264 0 : uint64_t tokenInfo = taskArg->token_;
265 0 : uint64_t outputSize = taskArg->outputSize_;
266 0 : uint64_t offset = taskArg->offset_;
267 0 : uint64_t yAxisOffset = taskArg->xAxisSize_;
268 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
269 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
270 :
271 : // 计算不同die的数据
272 : uint64_t step0BaseOffset =
273 0 : axisId_ == 0 ? dimId_[0] * outputSize + offset : dimId_[1] * dimSize_[0] * outputSize + offset + xAxisSize;
274 0 : uint64_t step0AddOffset = axisId_ == 0 ? dimSize_[0] * outputSize : outputSize;
275 0 : uint64_t step1AddOffset = rankId_ * outputSize + offset + (axisId_ == 0 ? xAxisSize : 0);
276 0 : auto xAxisGoSize = CalGoSize(xAxisSize);
277 0 : auto yAxisGoSize = CalGoSize(yAxisSize);
278 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] GeneArgs: inputAddr[%llu], outputAddr[%llu],"
279 : "step0BaseOffset[%llu], step0AddOffset[%llu], step1AddOffset[%llu]",
280 : inputAddr, outputAddr, step0BaseOffset, step0AddOffset, step1AddOffset);
281 : return {inputAddr, outputAddr, tokenInfo, step0BaseOffset, step0AddOffset, step1AddOffset, yAxisOffset,
282 0 : xAxisGoSize[0], xAxisGoSize[1], xAxisGoSize[2], xAxisGoSize[3],
283 0 : yAxisGoSize[0], yAxisGoSize[1], yAxisGoSize[2], yAxisGoSize[3]};
284 0 : }
285 : }
|
