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_all_reduce_small_count_aiv_rdma_executor.h"
12 : #include "sender.h"
13 : #include "reducer.h"
14 :
15 : namespace hccl {
16 : constexpr s32 INTRA_RS_STEP = 0;
17 : constexpr s32 INTER_AR_STEP = 1;
18 : constexpr s32 INTRA_AG_STEP = 2;
19 : constexpr u32 A_X_AGGR_SIZE = 2;
20 : constexpr u32 A_X_SIZE = 16;
21 : constexpr u64 HALF_OFFSET = 16 * 1024 * 1024;
22 :
23 : u64 CollAllReduceSmallCountAivRdmaExecutor::allreduceSmallDataAivRdmaCount_ = 0;
24 :
25 0 : CollAllReduceSmallCountAivRdmaExecutor::CollAllReduceSmallCountAivRdmaExecutor(
26 0 : const HcclDispatcher dispatcher, std::unique_ptr<TopoMatcher>& topoMatcher)
27 0 : : CollAllReduceExecutor(dispatcher, topoMatcher)
28 : {
29 0 : DMAReduceFlag_ = false;
30 0 : desc_.isAivMode = true;
31 0 : desc_.aivTagNum = AIV_A2_ALL_REDUCE_RDMA_KERNEL_NUM;
32 0 : }
33 :
34 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::CalcStreamNum(u32& streamNum)
35 : {
36 0 : u32 totalStreamNum = topoAttr_.deviceNumPerAggregation > 1U ? topoAttr_.deviceNumPerAggregation - 1U : 1U;
37 0 : streamNum = totalStreamNum - 1U;
38 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][CalcStreamNum] tag[%s] streamNum[%u]", tag_.c_str(), streamNum);
39 0 : return HCCL_SUCCESS;
40 : }
41 :
42 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::CalcCommInfo(std::vector<LevelNSubCommTransport>& opTransport)
43 : {
44 0 : TransportMemType inputType = TransportMemType::RESERVED;
45 0 : TransportMemType outputType = TransportMemType::RESERVED;
46 0 : CHK_RET(CalcTransportMemType(inputType, outputType));
47 0 : CHK_RET(CalcLevel0CommInfo(inputType, outputType, opTransport));
48 0 : CHK_RET(CalcLevel1CommInfo(inputType, outputType, opTransport));
49 :
50 : // aiv+rdma小数据量在server间使用HD通信域，并在多机A+X场景下当未设置使用RDMA时，默认使用PCIE
51 0 : bool isSingleAX = topoAttr_.serverNum == 1 && topoAttr_.moduleNum == 2; // A+X单机跨module
52 0 : if (topoMatcher_->GetExternalInputIntraRoceSwitch() == 0 && isSingleAX) {
53 0 : std::vector<SingleSubCommTransport>& commTransportLevel1 = opTransport[COMM_LEVEL1];
54 0 : for (u32 ringIndex = 0; ringIndex < commTransportLevel1.size(); ringIndex++) {
55 0 : for (auto& transportRequest : commTransportLevel1[ringIndex].transportRequests) {
56 0 : transportRequest.isUsedRdma = false;
57 : }
58 : }
59 : }
60 0 : return HCCL_SUCCESS;
61 : }
62 :
63 : HcclResult
64 0 : CollAllReduceSmallCountAivRdmaExecutor::CalcTransportMemType(TransportMemType& inputType, TransportMemType& outputType)
65 : {
66 : // 小数据量：使用AIVIN+AIVOUT，标记区在AIVOUT，RS前从inputPtr到AIVIN做本地拷贝
67 0 : inputType = TransportMemType::AIV_INPUT;
68 0 : outputType = TransportMemType::AIV_OUTPUT;
69 0 : HCCL_INFO(
70 : "[CollAllReduceSmallCountAivRdmaExecutor][CalcTransportMemType] tag[%s] inputType[%d], outputType[%d]",
71 : tag_.c_str(), inputType, outputType);
72 0 : return HCCL_SUCCESS;
73 : }
74 :
75 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::CalcLevel0CommInfo(
76 : TransportMemType inputType, TransportMemType outputType, std::vector<LevelNSubCommTransport>& opTransport)
77 : {
78 0 : CommParaInfo commParaLevel0(COMM_LEVEL0, CommType::COMM_TAG_MESH);
79 0 : commParaLevel0.meshSinglePlane = true;
80 0 : CHK_RET(CalcCommPlaneInfo(tag_, commParaLevel0, opTransport[COMM_LEVEL0], inputType, outputType));
81 0 : return HCCL_SUCCESS;
82 0 : }
83 :
84 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::CalcLevel1CommInfo(
85 : TransportMemType inputType, TransportMemType outputType, std::vector<LevelNSubCommTransport>& opTransport)
86 : {
87 0 : CommParaInfo commParaInfo(COMM_LEVEL1, CommType::COMM_TAG_HALVING_DOUBLING);
88 0 : CHK_RET(CalcCommPlaneInfo(tag_, commParaInfo, opTransport[COMM_LEVEL1], inputType, outputType));
89 0 : return HCCL_SUCCESS;
90 0 : }
91 :
92 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::CalNumBlocks(
93 : u32& numBlocks, u32 rankSize, [[maybe_unused]] u64 dataSize, [[maybe_unused]] HcclCMDType cmdType)
94 : {
95 0 : numBlocks = rankSize; // 默认情况使用rankSize个AIV
96 0 : u32 bestNumBlocks = numBlocks;
97 :
98 0 : CHK_PRT_RET(
99 : numBlocks_ < numBlocks,
100 : HCCL_WARNING(
101 : "[CollAllReduceSmallCountAivRdmaExecutor][CalNumBlocks]aivCore[%u] is invalid, at least need [%u].",
102 : numBlocks_, numBlocks),
103 : HCCL_E_PARA);
104 :
105 0 : HCCL_INFO(
106 : "[CollAllReduceSmallCountAivRdmaExecutor][CalNumBlocks] numBlocks is set to [%u], limit[%u], recommanded[%u]",
107 : numBlocks, numBlocks_, bestNumBlocks);
108 0 : return HCCL_SUCCESS;
109 : }
110 :
111 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::Orchestrate(OpParam& param, AlgResourceResponse& algRes)
112 : {
113 0 : HcclUs startut = TIME_NOW();
114 0 : allreduceSmallDataAivRdmaCount_ += 1;
115 0 : HCCL_INFO(
116 : "[CollAllReduceSmallCountAivRdmaExecutor][Orchestrate] AllreduceSmallCountAivRdma has been called [%llu].",
117 : allreduceSmallDataAivRdmaCount_);
118 0 : tag_ = param.tag;
119 0 : algResResp_ = &algRes;
120 :
121 : // 小数据量：使用AIVIN+AIVOUT，标记区在AIVOUT
122 0 : ExecMem execMem;
123 0 : execMem.count = param.DataDes.count;
124 0 : execMem.inputPtr = param.inputPtr;
125 0 : execMem.outputPtr = param.outputPtr;
126 0 : execMem.inputMem = algRes.aivInputMem;
127 0 : execMem.outputMem = algRes.aivOutputMem;
128 0 : HcclResult ret = KernelRun(param, execMem);
129 :
130 0 : CHK_PRT_RET(
131 : ret != HCCL_SUCCESS,
132 : HCCL_ERROR(
133 : "[CollAllReduceSmallCountAivRdmaExecutor]errNo[0x%016llx] tag[%s] executor kernel run failed",
134 : HCCL_ERROR_CODE(ret), param.tag.c_str()),
135 : ret);
136 :
137 0 : HCCL_INFO(
138 : "tag[%s], AllReduce executor orchestrate success, take time [%lld]us.", param.tag.c_str(),
139 : DURATION_US(TIME_NOW() - startut));
140 0 : return HCCL_SUCCESS;
141 0 : }
142 :
143 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::GetAdjInfo(
144 : [[maybe_unused]] AlgResourceResponse& algRes, [[maybe_unused]] AdjInfo& adjInfo)
145 : {
146 0 : return HCCL_SUCCESS;
147 : }
148 :
149 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::InterServerHDOneshot(
150 : const OpParam& param, ExecMem& execMem, u32& outputOffset, u64 sliceCount, u32 dbOffset, u32 interRankSize,
151 : u32 interRankId, bool isOpbase, std::vector<LINK>& interLinks)
152 : {
153 : (void)isOpbase;
154 0 : u64 reduceAttr = GetReduceAttr(execMem.inputMem, execMem.outputMem, param.DataDes.dataType, param.reduceType);
155 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][InterServerHDOneshot]reduceAttr is [%llu].", reduceAttr);
156 0 : std::unique_ptr<Sender> senderInfo;
157 0 : std::unique_ptr<Reducer> reducerInfo;
158 0 : senderInfo.reset(new (std::nothrow) Sender(param.DataDes.dataType, param.reduceType, reduceAttr));
159 0 : CHK_SMART_PTR_NULL(senderInfo);
160 0 : reducerInfo.reset(new (std::nothrow) Reducer(param.DataDes.dataType, param.reduceType, reduceAttr));
161 0 : CHK_SMART_PTR_NULL(reducerInfo);
162 0 : u32 hdStepNum = static_cast<u32>(log2(interRankSize));
163 0 : HCCL_INFO(
164 : "[CollAllReduceSmallCountAivRdmaExecutor] Find interlink type for cross-aggregation link[%d].",
165 : interLinks[(interRankId + 1) % A_X_AGGR_SIZE + interRankId - interRankId % A_X_AGGR_SIZE]->GetLinkType());
166 0 : u32 sliceSize = sliceCount * SIZE_TABLE[param.DataDes.dataType];
167 0 : auto opMeta = HcclOpMetaInfo::GetOneForAllReduce(
168 0 : 0, param.DataDes.dataType, ReduceType::INLINE_REDUCE, true, 0, false, hccl::CopyPattern::BCOPY, 1, true);
169 0 : CHK_RET(InitTask(dispatcher_, const_cast<Stream&>(param.stream), opMeta.isEnableCache, opMeta.GetCacheKey()));
170 0 : for (u32 step = 1; step <= hdStepNum; step++) {
171 0 : u32 peerMask = 1 << (step - 1);
172 0 : u32 peer = interRankId ^ peerMask;
173 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][InterServerHDOneshot] Step %u, peer %u.", step, peer);
174 0 : u32 sliceForReadOffset = HCCL_SMALL_COUNT_2_MB + (step - 1) * sliceSize + dbOffset;
175 0 : u32 sliceForWriteOffset = HCCL_SMALL_COUNT_2_MB + step * sliceSize + dbOffset;
176 0 : DeviceMem src = execMem.inputMem.range(sliceForReadOffset, sliceSize);
177 0 : DeviceMem dst = execMem.inputMem.range(sliceForWriteOffset, sliceSize);
178 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, const_cast<Stream&>(param.stream)));
179 0 : interLinks[peer]->TxAck(const_cast<Stream&>(param.stream));
180 0 : interLinks[peer]->RxAck(const_cast<Stream&>(param.stream));
181 0 : if (interLinks[peer]->IsSupportTransportWithReduce()
182 0 : && ((interLinks[peer]->GetLinkType() == LinkType::LINK_STANDARD_ROCE)
183 0 : || static_cast<bool>((RDMA_REDUCE_BITMASK & reduceAttr)))) {
184 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][InterServerHDOneshot] inter use RDMA");
185 0 : CHK_RET(senderInfo->run(
186 : interLinks[peer], sliceForWriteOffset, src, const_cast<Stream&>(param.stream), UserMemType::INPUT_MEM));
187 0 : CHK_RET(reducerInfo->run(
188 : dispatcher_, interLinks[peer], 0, src, src, src, const_cast<Stream&>(param.stream),
189 : DstMemType::RESULT_INPUT_MEM, UserMemType::INPUT_MEM));
190 0 : } else if (interLinks[peer]->IsSpInlineReduce() && static_cast<bool>((INLINE_REDUCE_BITMASK & reduceAttr))) {
191 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][InterServerHDOneshot] inter use SDMA");
192 0 : CHK_RET(senderInfo->run(
193 : interLinks[peer], sliceForWriteOffset, src, const_cast<Stream&>(param.stream), UserMemType::INPUT_MEM));
194 0 : CHK_RET(reducerInfo->run(
195 : dispatcher_, interLinks[peer], sliceForReadOffset, dst, dst, dst, const_cast<Stream&>(param.stream),
196 : DstMemType::RESULT_INPUT_MEM, UserMemType::INPUT_MEM));
197 : } else {
198 0 : CHK_RET(interLinks[peer]->TxAsync(
199 : UserMemType::INPUT_MEM, HALF_OFFSET + sliceForWriteOffset, src.ptr(), src.size(),
200 : const_cast<Stream&>(param.stream)));
201 0 : DeviceMem localSrc = execMem.inputMem.range(HALF_OFFSET + sliceForWriteOffset, sliceSize);
202 0 : CHK_RET(reducerInfo->run(
203 : dispatcher_, interLinks[peer], 0, localSrc, dst, src, const_cast<Stream&>(param.stream),
204 : DstMemType::RESULT_INPUT_MEM, UserMemType::INPUT_MEM));
205 0 : }
206 0 : }
207 0 : CHK_RET(LaunchTask(dispatcher_, const_cast<Stream&>(param.stream)));
208 0 : outputOffset = HCCL_SMALL_COUNT_2_MB + hdStepNum * sliceSize;
209 0 : return HCCL_SUCCESS;
210 0 : }
211 :
212 0 : HcclResult CollAllReduceSmallCountAivRdmaExecutor::KernelRun(const OpParam& param, ExecMem& execMem)
213 : {
214 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][KernelRun]AllReduce aiv enter");
215 0 : HcclWorkflowMode workflow = workflowMode_;
216 0 : bool isOpbase = (workflow == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE);
217 :
218 : // 获取通信域信息
219 0 : CHK_RET(CheckCommSize(COMM_LEVEL0, COMM_INDEX_0 + 1));
220 0 : SubCommInfo level0CommInfo = GetSubCommInfo(COMM_LEVEL0, COMM_INDEX_0);
221 0 : u32 commIndex = level0CommInfo.localRank;
222 0 : CHK_RET(CheckCommSize(COMM_LEVEL1, commIndex + 1));
223 0 : SubCommInfo level1CommInfo = GetSubCommInfo(COMM_LEVEL1, commIndex);
224 :
225 : // 数据准备，按照server内rankSize切片
226 0 : u32 perDataSize = SIZE_TABLE[param.DataDes.dataType];
227 0 : u64 totalSize = param.DataDes.count * perDataSize;
228 0 : std::vector<Slice> dataSegsSlice; // 数据分成ranksize份，每份的起始偏移和大小
229 0 : u32 sliceNum = level0CommInfo.localRankSize;
230 0 : CHK_RET(PrepareSliceDataWithAlignSize(totalSize, sliceNum, 0, dataSegsSlice, perDataSize));
231 0 : CHK_PRT_RET(
232 : commIndex >= dataSegsSlice.size(),
233 : HCCL_ERROR(
234 : "[CollAllReduceSmallCountAivRdmaExecutor][Run]commIndex[%u] >= dataSegsSlice size[%zu]", commIndex,
235 : dataSegsSlice.size()),
236 : HCCL_E_INTERNAL);
237 0 : std::vector<hccl::LINK> intraLinks = level0CommInfo.links;
238 0 : std::vector<hccl::LINK> interLinks = level1CommInfo.links;
239 0 : u32 intraRankSize = level0CommInfo.localRankSize;
240 0 : u32 interRankSize = level1CommInfo.localRankSize;
241 0 : u32 intraRankId = level0CommInfo.localRank;
242 0 : u32 interRankId = level1CommInfo.localRank;
243 :
244 : // reduce scatter via AIV
245 : void* dataBuffers[MAX_RANK_SIZE];
246 : void* flagBuffers[MAX_RANK_SIZE]; // 标记区的具体偏移在kernel中决定
247 0 : CHK_RET(PrepareAivBuffers(
248 : intraRankSize, intraRankId, 0, execMem.inputMem, execMem.inputMem, intraLinks, dataBuffers, flagBuffers,
249 : UserMemType::INPUT_MEM, UserMemType::INPUT_MEM, 0, HCCL_MID_COUNT_32_MB));
250 : // RS总数据量最大1m，rs的结果存储到2m处
251 0 : void* rsOutput = static_cast<u8*>(execMem.inputMem.ptr()) + HCCL_SMALL_COUNT_2_MB;
252 :
253 0 : AivOpArgs opArgs{HcclCMDType::HCCL_CMD_ALLREDUCE, execMem.inputPtr, rsOutput, execMem.count,
254 0 : param.DataDes.dataType, param.reduceType, 0, isOpbase};
255 0 : AivTopoArgs topoArgs{intraRankId, intraRankSize, topoAttr_.isDiffDeviceModule ? topoAttr_.devicePhyId : A_X_SIZE};
256 0 : topoArgs.identify = algoAttr_.identifier;
257 : u32 numBlocks;
258 0 : CHK_PRT_RET(
259 : CalNumBlocks(numBlocks, intraRankSize) != HCCL_SUCCESS, HCCL_ERROR("[%s] CalNumBlocks failed", __func__),
260 : HCCL_E_PARA);
261 0 : numBlocks_ = numBlocks;
262 0 : AivResourceArgs resourceArgs{param.tag, param.stream.ptr(), dataBuffers, flagBuffers, execMem.inputMem.size(),
263 0 : numBlocks_, param.aivTag};
264 0 : AivAlgArgs algArgs{INTRA_RS_STEP, true};
265 0 : algArgs.execTimeOut = topoMatcher_->GetExecTimeOutConfig();
266 0 : algArgs.execTimeOutSet = true;
267 0 : struct AivProfilingInfo aivProfilingInfo;
268 0 : aivProfilingInfo.counter = opCounter_;
269 :
270 0 : CHK_RET(ExecuteKernelLaunch(opArgs, topoArgs, resourceArgs, algArgs, aivProfilingInfo));
271 :
272 : // use hd algo
273 0 : u32 arOutputOffset = 0; // 跨机allreduce的结果的位置，相对于inputMem的偏移
274 0 : CHK_RET(InterServerHDOneshot(
275 : param, execMem, arOutputOffset, dataSegsSlice[commIndex].size / perDataSize, 0, interRankSize, interRankId,
276 : isOpbase, interLinks));
277 0 : void* arOutput = static_cast<u8*>(execMem.inputMem.ptr()) + arOutputOffset;
278 :
279 : // AllGather via AIV
280 0 : CHK_RET(PrepareAivBuffers(
281 : intraRankSize, intraRankId, 0, execMem.inputMem, execMem.inputMem, intraLinks, dataBuffers, flagBuffers,
282 : UserMemType::INPUT_MEM, UserMemType::INPUT_MEM, HCCL_SMALL_COUNT_8_MB, HCCL_MID_COUNT_32_MB));
283 :
284 0 : opArgs.input = arOutput;
285 0 : opArgs.output = execMem.outputPtr;
286 0 : resourceArgs.buffersIn = dataBuffers;
287 0 : resourceArgs.buffersOut = flagBuffers;
288 0 : resourceArgs.aivTag = GetNextAivTag(resourceArgs.aivTag);
289 0 : algArgs.step = INTRA_AG_STEP;
290 :
291 0 : CHK_RET(ExecuteKernelLaunch(opArgs, topoArgs, resourceArgs, algArgs, aivProfilingInfo));
292 :
293 0 : HCCL_INFO("[CollAllReduceSmallCountAivRdmaExecutor][KernelRun]AllReduce aiv run success.");
294 0 : return HCCL_SUCCESS;
295 0 : }
296 :
297 : REGISTER_EXEC("AllReduceSmallCountAivRdmaExecutor", AllReduceSmallCountAivRdma, CollAllReduceSmallCountAivRdmaExecutor);
298 :
299 : } // namespace hccl
|
