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 "coll_reduce_scatter_aiv_rdma_executor.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 : constexpr u32 A_X_SIZE = 16;
16 :
17 0 : CollReduceScatterAivRdmaExecutor::CollReduceScatterAivRdmaExecutor(
18 0 : const HcclDispatcher dispatcher, std::unique_ptr<TopoMatcher>& topoMatcher)
19 0 : : CollReduceScatterExecutor(dispatcher, topoMatcher)
20 : {
21 0 : DMAReduceFlag_ = false;
22 0 : desc_.isAivMode = true;
23 0 : }
24 :
25 0 : void CollReduceScatterAivRdmaExecutor::ParseParam(const OpParam& param)
26 : {
27 0 : tag_ = param.tag;
28 0 : root_ = param.root;
29 0 : opType_ = param.opType;
30 : // 记录图模式总数据量
31 0 : totalSize_ = topoAttr_.userRankSize * param.DataDes.count * SIZE_TABLE[param.DataDes.dataType];
32 0 : }
33 :
34 0 : HcclResult CollReduceScatterAivRdmaExecutor::CalcScratchMemSize(u64& scratchMemSize)
35 : {
36 0 : if (workflowMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
37 0 : scratchMemSize = 0U;
38 : } else {
39 0 : scratchMemSize = totalSize_;
40 : }
41 0 : HCCL_INFO(
42 : "[CollReduceScatterAivRdmaExecutor][CalcScratchMemSize] tag[%s] scratchMemSize[%llu]", tag_.c_str(),
43 : scratchMemSize);
44 0 : return HCCL_SUCCESS;
45 : }
46 :
47 0 : HcclResult CollReduceScatterAivRdmaExecutor::CalcCommInfo(std::vector<LevelNSubCommTransport>& opTransport)
48 : {
49 0 : TransportMemType inputType = TransportMemType::RESERVED;
50 0 : TransportMemType outputType = TransportMemType::RESERVED;
51 0 : CHK_RET(CalcTransportMemType(inputType, outputType));
52 0 : CHK_RET(CalcLevel0CommInfo(inputType, outputType, opTransport));
53 0 : CHK_RET(CalcLevel1CommInfo(inputType, outputType, opTransport));
54 0 : return HCCL_SUCCESS;
55 : }
56 :
57 : HcclResult
58 0 : CollReduceScatterAivRdmaExecutor::CalcTransportMemType(TransportMemType& inputType, TransportMemType& outputType)
59 : {
60 : // 使用AIVIN,标记区在AIVIN末尾,单算子模式用CCLOUT,图模式用USEROUT
61 0 : inputType = TransportMemType::AIV_INPUT;
62 0 : if (workflowMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
63 0 : outputType = TransportMemType::CCL_OUTPUT;
64 : } else {
65 0 : outputType = TransportMemType::SCRATCH;
66 : }
67 :
68 0 : HCCL_INFO(
69 : "[CollReduceScatterAivRdmaExecutor][CalcTransportMemType] tag[%s] inputType[%d], outputType[%d]", tag_.c_str(),
70 : inputType, outputType);
71 0 : return HCCL_SUCCESS;
72 : }
73 :
74 0 : HcclResult CollReduceScatterAivRdmaExecutor::CalcLevel0CommInfo(
75 : TransportMemType inputType, TransportMemType outputType, std::vector<LevelNSubCommTransport>& opTransport)
76 : {
77 0 : CommParaInfo commParaLevel0(COMM_LEVEL0, CommType::COMM_TAG_MESH);
78 0 : commParaLevel0.meshSinglePlane = true;
79 0 : CHK_RET(CalcCommPlaneInfo(tag_, commParaLevel0, opTransport[COMM_LEVEL0], inputType, outputType));
80 0 : return HCCL_SUCCESS;
81 0 : }
82 :
83 0 : HcclResult CollReduceScatterAivRdmaExecutor::CalNumBlocks(
84 : u32& numBlocks, u32 rankSize, [[maybe_unused]] u64 dataSize, [[maybe_unused]] HcclCMDType cmdType)
85 : {
86 0 : numBlocks = rankSize; // 多机场景,单算子ReduceScatter使用rankSize个aiv
87 0 : u32 bestNumBlocks = numBlocks;
88 :
89 0 : CHK_PRT_RET(
90 : numBlocks_ < numBlocks,
91 : HCCL_WARNING(
92 : "[CollReduceScatterAivRdmaExecutor][CalNumBlocks]aivCore[%u] is invalid, at least need [%u].", numBlocks_,
93 : numBlocks),
94 : HCCL_E_PARA);
95 :
96 0 : HCCL_INFO(
97 : "[CollReduceScatterAivRdmaExecutor][CalNumBlocks] numBlocks is set to [%u], limit[%u], recommanded[%u]",
98 : numBlocks, numBlocks_, bestNumBlocks);
99 0 : return HCCL_SUCCESS;
100 : }
101 :
102 0 : HcclResult CollReduceScatterAivRdmaExecutor::Orchestrate(OpParam& param, AlgResourceResponse& algRes)
103 : {
104 0 : HCCL_INFO("[CollReduceScatterAivRdmaExecutor][Orchestrate]start");
105 :
106 0 : HcclUs startut = TIME_NOW();
107 0 : tag_ = param.tag;
108 0 : algResResp_ = &algRes;
109 :
110 : // OutputMem单算子模式用CCLOUT,图模式用USEROUT
111 0 : ExecMem execMem;
112 0 : execMem.count = param.DataDes.count;
113 0 : execMem.inputPtr = param.inputPtr;
114 0 : execMem.outputPtr = param.outputPtr;
115 0 : execMem.inputMem = algRes.aivInputMem;
116 0 : if (workflowMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
117 0 : execMem.outputMem = algRes.cclOutputMem;
118 0 : execMem.scratchMem = algRes.cclOutputMem;
119 : } else {
120 0 : execMem.outputMem = algRes.paramOutputMem;
121 0 : execMem.scratchMem = algRes.scratchMem;
122 : }
123 0 : HcclResult ret = KernelRun(param, execMem);
124 :
125 0 : CHK_PRT_RET(
126 : ret != HCCL_SUCCESS,
127 : HCCL_ERROR(
128 : "[CollReduceScatterAivRdmaExecutor]errNo[0x%016llx] tag[%s] executor kernel run failed",
129 : HCCL_ERROR_CODE(ret), param.tag.c_str()),
130 : ret);
131 :
132 0 : HCCL_INFO(
133 : "tag[%s], ReduceScatter executor orchestrate success, take time [%lld]us.", param.tag.c_str(),
134 : DURATION_US(TIME_NOW() - startut));
135 0 : return HCCL_SUCCESS;
136 0 : }
137 :
138 0 : HcclResult CollReduceScatterAivRdmaExecutor::KernelRun(const OpParam& param, ExecMem& execMem)
139 : {
140 0 : HCCL_CONFIG_INFO(HCCL_ALG, "[CollReduceScatterAivRdmaExecutor][KernelRun]ReduceScatter aiv enter");
141 :
142 0 : HcclWorkflowMode workflow = workflowMode_;
143 0 : bool isOpbase = (workflow == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE);
144 :
145 : // 获取通信域信息
146 0 : CHK_RET(CheckCommSize(COMM_LEVEL0, COMM_INDEX_0 + 1));
147 0 : SubCommInfo outerCommInfo = GetSubCommInfo(COMM_LEVEL0, COMM_INDEX_0);
148 0 : u32 commIndex = outerCommInfo.localRank;
149 0 : CHK_RET(CheckCommSize(COMM_LEVEL1, commIndex + 1));
150 0 : SubCommInfo innerCommInfo = GetSubCommInfo(COMM_LEVEL1, commIndex);
151 :
152 : /* 第一步 节点内重排序RS */
153 : // 数据准备,按照server内rankSize切片
154 0 : u32 perDataSize = SIZE_TABLE[param.DataDes.dataType];
155 0 : u64 perRankSize = param.DataDes.count * perDataSize;
156 0 : std::vector<hccl::LINK> intraLinks = outerCommInfo.links; // 机间
157 0 : std::vector<hccl::LINK> interLinks = innerCommInfo.links; // 机内
158 0 : u32 intraRankSize = outerCommInfo.localRankSize;
159 0 : u32 intraRankId = outerCommInfo.localRank;
160 :
161 : // reduce scatter阶段,inputMem0-31m做数据区,32M开始后的1M做标记区
162 : void* dataBuffers[MAX_RANK_SIZE];
163 : void* flagBuffers[MAX_RANK_SIZE]; // 标记区的具体偏移在kernel中决定
164 0 : CHK_RET(PrepareAivBuffers(
165 : intraRankSize, intraRankId, 0, execMem.inputMem, execMem.inputMem, intraLinks, dataBuffers, flagBuffers,
166 : UserMemType::INPUT_MEM, UserMemType::INPUT_MEM, 0, HCCL_MID_COUNT_32_MB));
167 :
168 0 : u32 serverNum = innerCommInfo.localRankSize;
169 : // 先做本地拷贝到AIVIN再跨片拷贝;output统一为reduceScatterInput的位置,即buffer中原位
170 0 : AivOpArgs opArgs{
171 : HcclCMDType::HCCL_CMD_REDUCE_SCATTER,
172 0 : execMem.inputPtr,
173 0 : execMem.outputPtr,
174 0 : execMem.count,
175 0 : param.DataDes.dataType,
176 0 : param.reduceType,
177 : 0,
178 0 : isOpbase};
179 : AivTopoArgs topoArgs{
180 : intraRankId,
181 : intraRankSize,
182 0 : topoAttr_.isDiffDeviceModule ? topoAttr_.devicePhyId : A_X_SIZE,
183 : 0,
184 : serverNum,
185 0 : topoAttr_.deviceType,
186 0 : algoAttr_.identifier};
187 : u32 numBlocks;
188 0 : CHK_PRT_RET(
189 : CalNumBlocks(numBlocks, intraRankSize) != HCCL_SUCCESS, HCCL_ERROR("[%s] CalNumBlocks failed", __func__),
190 : HCCL_E_PARA);
191 0 : numBlocks_ = numBlocks;
192 0 : AivResourceArgs resourceArgs{param.tag, param.stream.ptr(), dataBuffers, flagBuffers, execMem.inputMem.size(),
193 0 : numBlocks_, param.aivTag};
194 0 : AivAlgArgs algArgs{0};
195 0 : algArgs.execTimeOut = topoMatcher_->GetExecTimeOutConfig();
196 0 : algArgs.execTimeOutSet = true;
197 0 : struct AivProfilingInfo aivProfilingInfo;
198 0 : aivProfilingInfo.counter = opCounter_;
199 :
200 0 : CHK_RET(ExecuteKernelLaunch(opArgs, topoArgs, resourceArgs, algArgs, aivProfilingInfo));
201 : /* 第二步 节点间RS */
202 0 : auto autoSelectedAlgTypeLevel1 = static_cast<u32>(algType_.algoLevel1);
203 0 : ReduceType reduceType
204 0 : = ((param.reduceType != HCCL_REDUCE_PROD) && (param.DataDes.dataType != HCCL_DATA_TYPE_INT64)) ?
205 : ReduceType::INLINE_REDUCE :
206 : ReduceType::TBE_REDUCE;
207 0 : auto opMeta = HcclOpMetaInfo::GetOneForReduceScatter(
208 0 : autoSelectedAlgTypeLevel1, param.DataDes.dataType, reduceType, false, false, CopyPattern::BCOPY, false, 0,
209 : true);
210 0 : CHK_RET(InitTask(dispatcher_, const_cast<Stream&>(param.stream), opMeta.isEnableCache, opMeta.GetCacheKey()));
211 :
212 0 : u32 innerRankSize = innerCommInfo.localRankSize;
213 0 : DeviceMem inputMem = execMem.inputMem;
214 0 : if (innerRankSize > 1) {
215 : // execMem.inputMem要改成按平面制定的初始位置
216 0 : u64 reduceAttr = GetReduceAttr(execMem.inputMem, execMem.outputMem, param.DataDes.dataType, param.reduceType);
217 0 : std::unique_ptr<ExecutorBase> innerExecutor;
218 0 : std::vector<Slice> dataSegsSlice;
219 0 : dataSegsSlice.resize(innerRankSize);
220 0 : for (u32 i = 0; i < innerRankSize; i++) {
221 0 : dataSegsSlice[i].size = perRankSize;
222 0 : dataSegsSlice[i].offset = (commIndex * innerRankSize + i) * perRankSize;
223 : }
224 0 : u64 count = param.DataDes.count;
225 0 : u64 baseOffset = 0;
226 0 : DeviceMem inputMem = execMem.inputMem;
227 0 : DeviceMem scratchMem = execMem.scratchMem;
228 0 : if (algType_.algoLevel1 == AlgTypeLevel1::ALG_LEVEL1_RING) {
229 0 : innerExecutor = AlgTemplateRegistry::Instance().GetAlgTemplate(
230 0 : TemplateType::TEMPLATE_REDUCESCATTER_RING, dispatcher_);
231 0 : CHK_SMART_PTR_NULL(innerExecutor);
232 0 : CHK_RET(innerExecutor->Prepare(reduceAttr));
233 0 : HCCL_INFO("ReduceScatter mesh: using ring algo inter-server.");
234 0 : } else if (algType_.algoLevel1 == AlgTypeLevel1::ALG_LEVEL1_NHR) {
235 : innerExecutor
236 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCESCATTER_NHR, dispatcher_);
237 0 : CHK_SMART_PTR_NULL(innerExecutor);
238 0 : CHK_RET(innerExecutor->Prepare(reduceAttr, false));
239 0 : innerExecutor->CloseBarrier();
240 0 : HCCL_INFO("ReduceScatter mesh: using nhr algo inter-server.");
241 0 : } else if (algType_.algoLevel1 == AlgTypeLevel1::ALG_LEVEL1_NHR_V1) {
242 0 : innerExecutor = AlgTemplateRegistry::Instance().GetAlgTemplate(
243 0 : TemplateType::TEMPLATE_REDUCESCATTER_NHR_V1, dispatcher_);
244 0 : CHK_SMART_PTR_NULL(innerExecutor);
245 0 : CHK_RET(innerExecutor->Prepare(reduceAttr));
246 0 : HCCL_INFO("ReduceScatter mesh: using nhr_v1 algo inter-server.");
247 0 : } else if (algType_.algoLevel1 == AlgTypeLevel1::ALG_LEVEL1_NB) {
248 : innerExecutor
249 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCESCATTER_NB, dispatcher_);
250 0 : CHK_SMART_PTR_NULL(innerExecutor);
251 0 : CHK_RET(innerExecutor->Prepare(reduceAttr));
252 0 : HCCL_INFO("ReduceScatter mesh: using nonuniform-bruck algo inter-server.");
253 : } else {
254 0 : count = count * innerRankSize;
255 0 : baseOffset = commIndex * innerRankSize * perRankSize;
256 0 : inputMem = execMem.inputMem.range(commIndex * innerRankSize * perRankSize, perRankSize * innerRankSize);
257 0 : scratchMem = execMem.scratchMem.range(commIndex * innerRankSize * perRankSize, perRankSize * innerRankSize);
258 0 : innerExecutor = AlgTemplateRegistry::Instance().GetAlgTemplate(
259 0 : TemplateType::TEMPLATE_REDUCESCATTER_RECURSIVE_HD, dispatcher_);
260 0 : CHK_SMART_PTR_NULL(innerExecutor);
261 0 : CHK_RET(innerExecutor->Prepare(reduceAttr));
262 0 : HCCL_INFO("ReduceScatter mesh: using halving-doubling algo inter-server.");
263 : }
264 0 : CHK_RET(innerExecutor->Prepare(
265 : inputMem, inputMem, scratchMem, count, param.DataDes.dataType, param.stream, param.reduceType,
266 : LEVEL0_BRIDGE_RANK_ID, dataSegsSlice, baseOffset));
267 0 : CHK_RET(innerExecutor->RegisterProfiler(
268 : (innerRankSize << PROF_RANKSIZE_OFFSET_OF_PLANEID) + innerCommInfo.localRank, PROF_STAGE_0,
269 : HCCL_EXEC_STEP_NOT_SET, param.stream));
270 :
271 0 : CHK_RET(RunTemplate(innerExecutor, innerCommInfo));
272 0 : HCCL_INFO("[CollReduceScatterAivRdmaExecutor] rdma stage run success.");
273 0 : }
274 : /* 第三步 最后D2D拷贝 */
275 :
276 : // 如果使用CCL buffer,需要将CCL buffer in中的结果拷贝到user buffer out
277 : DeviceMem srcMem
278 0 : = execMem.inputMem.range(perRankSize * (commIndex * serverNum + innerCommInfo.localRank), perRankSize);
279 0 : DeviceMem dstMem = DeviceMem::create(execMem.outputPtr, perRankSize);
280 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, const_cast<Stream&>(param.stream)));
281 :
282 0 : CHK_RET(LaunchTask(dispatcher_, const_cast<Stream&>(param.stream)));
283 :
284 0 : HCCL_INFO("[CollReduceScatterAivRdmaExecutor][KernelRun]ReduceScatter aiv run success");
285 0 : return HCCL_SUCCESS;
286 0 : }
287 :
288 : REGISTER_EXEC("ReduceScatterAivRdmaExecutor", ReduceScatterAivRdma, CollReduceScatterAivRdmaExecutor);
289 :
290 : } // namespace hccl
|