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_all_reduce_mesh2d_one_shot.h"
12 : #include "ccu_instruction_all_reduce_mesh2d_one_shot.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr int INPUT_XN_ID = 0;
17 : constexpr int SCRATCH_XN_ID = 1;
18 : constexpr int TOKEN_XN_ID = 2;
19 : constexpr int CKE_IDX_0 = 0;
20 : constexpr int CKE_IDX_1 = 1;
21 : constexpr int CKE_IDX_2 = 2;
22 : constexpr int CKE_IDX_3 = 3;
23 : constexpr int CKE_IDX_4 = 4;
24 : constexpr uint32_t AXIS_NUM = 2;
25 :
26 0 : CcuContextAllReduceMesh2DOneShot::CcuContextAllReduceMesh2DOneShot(
27 0 : const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
28 0 : : CcuContextAlgBase(arg, transports, group)
29 : {
30 0 : const CcuCtxArgAllReduceMesh2DOneShot* ctxArg = dynamic_cast<const CcuCtxArgAllReduceMesh2DOneShot*>(&arg);
31 0 : if (ctxArg == nullptr) {
32 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DOneShot::ctxArg ptr is null"));
33 : }
34 0 : dimSize_ = ctxArg->dimSize_;
35 0 : axisId_ = ctxArg->axisId_;
36 0 : rankId_ = ctxArg->rankId_;
37 0 : dataType_ = ctxArg->op_.dataType;
38 0 : outputDataType_ = ctxArg->op_.outputDataType;
39 0 : reduceOp_ = ctxArg->op_.reduceOp;
40 0 : if (outputDataType_ == DataType::INVALID) {
41 0 : outputDataType_ = dataType_;
42 0 : HCCL_INFO(
43 : "[CcuContextAllReduceMesh2DOneShot] outputDataType is [INVALID], set outputDataType to[%s]",
44 : outputDataType_.Describe().c_str());
45 : }
46 :
47 0 : HCCL_INFO(
48 : "[CcuContextAllReduceMesh2DOneShot] Init, CtxArgs are dimSize0[%llu], dimSize1[%llu], axisId[%u], "
49 : "rankId[%llu], dataType[%s], outputDataType[%s], reduceOp[%s]",
50 : dimSize_[0], dimSize_[1], axisId_, rankId_, dataType_.Describe().c_str(), outputDataType_.Describe().c_str(),
51 : reduceOp_.Describe().c_str());
52 0 : uint32_t max_dimSize = 2;
53 0 : if (dimSize_.size() != max_dimSize or axisId_ > 1) {
54 0 : THROW<NullPtrException>(StringFormat(
55 : "[CcuContextAllReduceMesh2DOneShot] dimSize[%u] or axisId[%u] is invalid", dimSize_.size(), axisId_));
56 : }
57 0 : CHK_PRT_THROW(
58 : dimSize_[0] == 0 || dimSize_[1] == 0,
59 : HCCL_ERROR(
60 : "[CcuContextAllReduceMesh2DOneShot] dimSize0[%llu] or dimSize1[%llu] is zero", dimSize_[0], dimSize_[1]),
61 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
62 :
63 0 : myRankIdxInAxis_.push_back(rankId_ % dimSize_[0]); // 本 rank 在第 0 维上的 index
64 0 : myRankIdxInAxis_.push_back(rankId_ / dimSize_[0]); // 本 rank 在第 1 维上的 index
65 :
66 0 : myRankIdxInCurrentAxis_ = myRankIdxInAxis_[axisId_];
67 0 : currentAxisRankSize_ = dimSize_[axisId_];
68 :
69 : // 同步信号初始化
70 0 : currAxisSignalName_ = "CcuContextAllReduceMesh2DOneShotAxisSync_" + std::to_string(axisId_);
71 0 : otherAxisSignalName_ = "CcuContextAllReduceMesh2DOneShotAxisSync_" + std::to_string(1 - axisId_);
72 0 : currAxisSignal_ = CreateMaskSignal();
73 0 : ExportMaskSignal(currAxisSignal_, currAxisSignalName_);
74 0 : otherAxisSignal_ = ImportMaskSignal(otherAxisSignalName_);
75 :
76 0 : HCCL_INFO(
77 : "[CcuContextAllReduceMesh2DOneShot] Init, myRankIdx0[%llu], myRankIdx1[%llu], "
78 : "myRankIdxInCurrentAxis[%llu], currentAxisRankSize[%llu]",
79 : myRankIdxInAxis_[0], myRankIdxInAxis_[1], myRankIdxInCurrentAxis_, currentAxisRankSize_);
80 0 : }
81 :
82 0 : void CcuContextAllReduceMesh2DOneShot::Algorithm()
83 : {
84 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] AllReduceMesh2DOneShot run");
85 0 : uint16_t selfBit = 1 << myRankIdxInCurrentAxis_;
86 0 : uint16_t allBit = ((1 << currentAxisRankSize_) - 1) & (~(1 << myRankIdxInCurrentAxis_));
87 :
88 0 : InitVariables();
89 :
90 0 : LoadArgs();
91 :
92 : // 前同步
93 0 : for (auto t : transports) {
94 0 : WriteVariableWithSignal(*t, inputAddr_[myRankIdxInCurrentAxis_], INPUT_XN_ID, CKE_IDX_1, selfBit);
95 0 : WriteVariableWithSignal(*t, scratchAddr_[myRankIdxInCurrentAxis_], SCRATCH_XN_ID, CKE_IDX_2, selfBit);
96 0 : WriteVariableWithSignal(*t, token_[myRankIdxInCurrentAxis_], TOKEN_XN_ID, CKE_IDX_3, selfBit);
97 : }
98 :
99 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit);
100 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit);
101 0 : GroupWait(*transportGroup, CKE_IDX_3, allBit);
102 :
103 : // OneShot Step1
104 0 : CcuRep::Variable& Step1Offset = (axisId_ == 0) ? xSliceOffset_ : ySliceOffset_;
105 0 : GroupOpSize& Step1GoSize = (axisId_ == 0) ? xGoSize_ : yGoSize_;
106 :
107 0 : DoGroupReduce(inputAddr_, scratchAddr_[myRankIdxInCurrentAxis_], Step1Offset, Step1GoSize);
108 :
109 0 : DoAxisSync(0);
110 0 : DoGroupSync(CKE_IDX_4, selfBit, allBit);
111 0 : DoAxisSync(1);
112 :
113 : // OneShot Step2
114 0 : CcuRep::Variable& Step2Offset = (axisId_ == 0) ? ySliceOffset_ : xSliceOffset_;
115 0 : GroupOpSize& Step2GoSize = (axisId_ == 0) ? yGoSize_ : xGoSize_;
116 :
117 0 : DoGroupReduce(scratchAddr_, outputAddr_[myRankIdxInCurrentAxis_], Step2Offset, Step2GoSize);
118 :
119 0 : DoAxisSync(0);
120 0 : DoGroupSync(CKE_IDX_0, selfBit, allBit);
121 0 : DoAxisSync(1);
122 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] AllReduceMesh2DOneShot end");
123 0 : return;
124 : }
125 :
126 0 : void CcuContextAllReduceMesh2DOneShot::DoGroupSync(int ckeIdx, uint16_t selfBit, uint16_t allBit)
127 : {
128 0 : HCCL_INFO(
129 : "[CcuContextAllReduceMesh2DOneShot] DoGroupSync Starts, ckeIdx[%d], selfBit[%u], allBit[%u]", ckeIdx, selfBit,
130 : allBit);
131 0 : for (auto t : transports) {
132 0 : RemotePost(*t, ckeIdx, selfBit);
133 : }
134 0 : GroupWait(*transportGroup, ckeIdx, allBit);
135 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] DoGroupSync Ends");
136 0 : return;
137 : }
138 :
139 0 : void CcuContextAllReduceMesh2DOneShot::DoGroupReduce(
140 : std::vector<CcuRep::Variable>& srcBase, CcuRep::Variable& dstBase, CcuRep::Variable& offset, GroupOpSize& goSize)
141 : {
142 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] DoGroupReduce Starts");
143 : // 从轴上所有的对端的对应位置读取数据
144 0 : std::vector<CcuRep::Memory> srcAddrs;
145 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
146 0 : srcAddrs.push_back(CreateMemory());
147 : }
148 0 : uint32_t rmtId = 0;
149 0 : uint32_t curId = 0;
150 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
151 0 : if (rankIdx != myRankIdxInCurrentAxis_) {
152 0 : curId = rmtId;
153 0 : rmtId++;
154 : } else {
155 0 : curId = currentAxisRankSize_ - 1;
156 : }
157 0 : srcAddrs[curId].addr = srcBase[rankIdx];
158 0 : srcAddrs[curId].addr += offset;
159 0 : srcAddrs[curId].token = token_[rankIdx];
160 : }
161 : // Reduce 到本端
162 0 : CcuRep::Memory dstAddr = CreateMemory();
163 0 : dstAddr.addr = dstBase;
164 0 : dstAddr.addr += offset;
165 0 : dstAddr.token = token_[myRankIdxInCurrentAxis_];
166 0 : GroupReduce(transports, dstAddr, srcAddrs, goSize, dataType_, outputDataType_, reduceOp_);
167 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] DoGroupReduce Ends");
168 0 : return;
169 0 : }
170 :
171 0 : void CcuContextAllReduceMesh2DOneShot::DoAxisSync(uint32_t signalIdx)
172 : {
173 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] DoAxisSync Starts, signalIdx[%u]", signalIdx);
174 0 : uint32_t sendBit = 1 << axisId_;
175 0 : uint32_t waitBit = 1 << (1 - axisId_);
176 0 : sendBit = sendBit << (AXIS_NUM * signalIdx);
177 0 : waitBit = waitBit << (AXIS_NUM * signalIdx);
178 0 : LocalCtxPost(otherAxisSignal_, sendBit);
179 0 : LocalWait(currAxisSignal_, waitBit);
180 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] DoAxisSync Ends");
181 0 : return;
182 : }
183 :
184 0 : void CcuContextAllReduceMesh2DOneShot::InitVariables()
185 : {
186 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] InitVariables Starts");
187 : // 初始化资源
188 0 : uint16_t transportIdx = 0;
189 0 : if (transports.size() == 0) {
190 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DOneShot transports is empty"));
191 : }
192 : // 按照rank号从小到大遍历transports，遇到本rank就填充本地资源，否则依次取远端资源，要求给框架返回的Link同样是按顺序排列的
193 0 : for (uint64_t peerId = 0; peerId < currentAxisRankSize_; peerId++) {
194 0 : if (peerId == myRankIdxInCurrentAxis_) {
195 0 : inputAddr_.push_back(CreateVariable());
196 0 : scratchAddr_.push_back(CreateVariable());
197 0 : token_.push_back(CreateVariable());
198 : } else {
199 0 : HCCL_INFO(
200 : "[CcuContextAllReduceMesh2DOneShot] MyRank[%u], PeerId[%llu], TransportId[%u]", myRankIdxInCurrentAxis_,
201 : peerId, transportIdx);
202 0 : CHK_PRT_RET(
203 : transports[transportIdx] == nullptr || transportIdx >= transports.size(),
204 : HCCL_ERROR("[CcuContextAllReduceMesh2DOneShot] Algorithm transport ptr is null or transportIdx is out "
205 : "of bounds"), );
206 0 : inputAddr_.push_back(CreateVariable((*transports[transportIdx]), INPUT_XN_ID));
207 0 : scratchAddr_.push_back(CreateVariable((*transports[transportIdx]), SCRATCH_XN_ID));
208 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
209 0 : transportIdx++;
210 : }
211 0 : outputAddr_.push_back(CreateVariable());
212 : }
213 :
214 0 : xSliceOffset_ = CreateVariable();
215 0 : ySliceOffset_ = CreateVariable();
216 0 : xGoSize_ = CreateGroupOpSize();
217 0 : yGoSize_ = CreateGroupOpSize();
218 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] InitVariables Ends");
219 0 : return;
220 : }
221 :
222 0 : void CcuContextAllReduceMesh2DOneShot::LoadArgs()
223 : {
224 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] LoadArgs Starts");
225 0 : Load(inputAddr_[myRankIdxInCurrentAxis_]);
226 0 : Load(outputAddr_[myRankIdxInCurrentAxis_]);
227 0 : Load(token_[myRankIdxInCurrentAxis_]);
228 0 : Load(scratchAddr_[myRankIdxInCurrentAxis_]);
229 0 : Load(xSliceOffset_);
230 0 : Load(ySliceOffset_);
231 0 : Load(xGoSize_);
232 0 : Load(yGoSize_);
233 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] LoadArgs Eends");
234 0 : return;
235 : }
236 :
237 0 : std::vector<uint64_t> CcuContextAllReduceMesh2DOneShot::GeneArgs(const CcuTaskArg& arg)
238 : {
239 0 : HCCL_INFO("[CcuContextAllReduceMesh2DOneShot] GeneArgs Starts");
240 0 : const CcuTaskArgAllReduceMesh2DOneShot* taskArg = dynamic_cast<const CcuTaskArgAllReduceMesh2DOneShot*>(&arg);
241 0 : if (taskArg == nullptr) {
242 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DOneShot::taskArg ptr is null"));
243 : }
244 0 : uint64_t tokenInfo = taskArg->token_;
245 :
246 0 : uint64_t inputAddr = taskArg->inputAddr_;
247 0 : uint64_t outputAddr = taskArg->outputAddr_;
248 0 : uint64_t scratchAddr = taskArg->scratchAddr_;
249 :
250 0 : uint64_t xSliceOffset = taskArg->xSliceOffset_;
251 0 : uint64_t ySliceOffset = taskArg->ySliceOffset_;
252 :
253 0 : auto xGoSize = CalGoSize(taskArg->xSliceSize_);
254 0 : auto yGoSize = CalGoSize(taskArg->ySliceSize_);
255 :
256 0 : HCCL_INFO(
257 : "[CcuContextAllReduceMesh2DOneShot] GeneArgs, TaskArgs are inputAddr[%llu], "
258 : "outputAddr[%llu], scratchAddr[%llu], xSliceSize[%llu], ySliceSize[%llu], xSliceOffset[%llu], "
259 : "ySliceOffset[%llu]",
260 : inputAddr, outputAddr, scratchAddr, taskArg->xSliceSize_, taskArg->ySliceSize_, xSliceOffset, ySliceOffset);
261 :
262 0 : std::vector<uint64_t> taskArgList = {inputAddr, outputAddr, tokenInfo, scratchAddr, xSliceOffset, ySliceOffset};
263 : // push goSize
264 0 : for (auto goSize : {xGoSize, yGoSize}) {
265 0 : for (auto val : goSize) {
266 0 : taskArgList.push_back(val);
267 : }
268 0 : }
269 0 : return taskArgList;
270 0 : }
271 : } // namespace Hccl
|
