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 <cmath>
12 : #include "alg_template_register.h"
13 : #include "all_reduce_hd_optim_pub.h"
14 : namespace hccl {
15 0 : AllReduceHDOptim::AllReduceHDOptim(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : AllReduceHDOptim::~AllReduceHDOptim() {}
18 :
19 0 : HcclResult AllReduceHDOptim::Prepare(
20 : u64 reduceAttrBitMap, std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
21 : std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 userRank, HcomCollOpInfo* opInfo, bool aicpu)
22 : {
23 0 : reduceAttr_ = reduceAttrBitMap;
24 0 : userRank_ = userRank;
25 0 : meshStreams_ = meshStreams;
26 0 : meshSignal_ = &meshSignal;
27 0 : meshSignalAux_ = &meshSignalAux;
28 0 : opInfo_ = opInfo;
29 0 : aicpu_ = aicpu;
30 0 : return HCCL_SUCCESS;
31 : }
32 :
33 0 : HcclResult AllReduceHDOptim::MainRecordSub(u32 streamNum)
34 : {
35 0 : if (aicpu_) {
36 0 : return HCCL_SUCCESS;
37 : }
38 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
39 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex], profilerInput_.stage));
40 : }
41 0 : return HCCL_SUCCESS;
42 : }
43 :
44 0 : HcclResult AllReduceHDOptim::SubWaitMain(u32 streamNum)
45 : {
46 0 : if (aicpu_) {
47 0 : return HCCL_SUCCESS;
48 : }
49 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
50 0 : CHK_RET(LocalNotify::Wait(
51 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
52 : }
53 0 : return HCCL_SUCCESS;
54 : }
55 :
56 0 : HcclResult AllReduceHDOptim::MainWaitSub(u32 streamNum)
57 : {
58 0 : if (aicpu_) {
59 0 : return HCCL_SUCCESS;
60 : }
61 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
62 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
63 : }
64 0 : return HCCL_SUCCESS;
65 : }
66 :
67 0 : HcclResult AllReduceHDOptim::SubRecordMain(u32 streamNum)
68 : {
69 0 : if (aicpu_) {
70 0 : return HCCL_SUCCESS;
71 : }
72 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
73 0 : CHK_RET(LocalNotify::Post(
74 : meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
75 : }
76 0 : return HCCL_SUCCESS;
77 : }
78 :
79 : // allreduce算法的函数入口
80 0 : HcclResult AllReduceHDOptim::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
81 : {
82 0 : HcclResult ret = HCCL_SUCCESS;
83 0 : CHK_SMART_PTR_NULL(dispatcher_);
84 0 : CHK_PTR_NULL(stream_.ptr());
85 0 : HCCL_INFO(
86 : "AllReduceHDOptim run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
87 : outputMem_.ptr(), outputMem_.ptr(), count_);
88 :
89 0 : if (links.size() < rankSize) {
90 0 : HCCL_ERROR(
91 : "[AllReduceHDOptim][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(),
92 : rankSize);
93 0 : return HCCL_E_INTERNAL;
94 : }
95 :
96 0 : if (meshStreams_.size() < base) {
97 0 : HCCL_ERROR(
98 : "[AllReduceHDOptim][RunAsync]rank[%u] meshStreams_[%llu] is less than need[%u]", rank, meshStreams_.size(),
99 : base);
100 0 : return HCCL_E_INTERNAL;
101 : }
102 0 : u32 totalSize = SIZE_TABLE[dataType_] * count_;
103 0 : userMemIn = DeviceMem::create(opInfo_->inputAddr, totalSize);
104 0 : userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize);
105 0 : emptyMem_ = outputMem_.range(0, 0);
106 0 : nSteps = static_cast<u32>(log2(rankSize));
107 0 : stepPow = static_cast<u32>(pow(base, nSteps));
108 :
109 0 : ret = RunPreCopy(rank, rankSize, links);
110 0 : CHK_PRT_RET(
111 : ret != HCCL_SUCCESS,
112 : HCCL_ERROR("[AllReduceHDOptim][RunAsync]rank[%u] count[%llu] failed RunPreCopy step", rank, count_), ret);
113 :
114 0 : if (rank < stepPow) {
115 0 : ret = RunAllReduceHDOptim(rank, rankSize, links);
116 0 : CHK_PRT_RET(
117 : ret != HCCL_SUCCESS,
118 : HCCL_ERROR(
119 : "[AllReduceHDOptim][RunAsync]rank[%u] count[%llu] failed RunAllReduceHDOptim step", rank, count_),
120 : ret);
121 : }
122 :
123 0 : if (stepPow != rankSize) {
124 0 : ret = RunFinalStep(rank, rankSize, links);
125 0 : CHK_PRT_RET(
126 : ret != HCCL_SUCCESS,
127 : HCCL_ERROR("[AllReduceHDOptim][RunAsync]rank[%u] count[%llu] failed RunFinalStep step", rank, count_), ret);
128 : }
129 :
130 0 : HCCL_INFO("AllReduceHDOptim finished: rank[%u] ranksize[%u]", rank, rankSize);
131 0 : return HCCL_SUCCESS;
132 : }
133 :
134 0 : HcclResult AllReduceHDOptim::RunPreCopy(u32 rank, u32 rankSize, const std::vector<LINK>& links)
135 : {
136 0 : u32 totalSize = SIZE_TABLE[dataType_] * count_;
137 :
138 0 : DeviceMem src = userMemIn.range(0, totalSize);
139 0 : DeviceMem dst = outputMem_.range(0, totalSize);
140 0 : DeviceMem nextDst = outputMem_.range(totalSize, totalSize);
141 :
142 0 : if (stepPow == rankSize) {
143 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, nextDst, src, stream_));
144 0 : return HCCL_SUCCESS;
145 : }
146 :
147 0 : u32 neighCur = rank ^ (1 << nSteps);
148 0 : if (neighCur < rankSize) {
149 : // reduce写
150 0 : if (rank >= pow(base, nSteps)) {
151 0 : CHK_PTR_NULL(links[neighCur]);
152 0 : CHK_RET(links[neighCur]->RxAck(stream_));
153 0 : void* remMemPtr = nullptr;
154 0 : CHK_RET(links[neighCur]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
155 0 : dst = DeviceMem::create(static_cast<u8*>(remMemPtr), totalSize);
156 0 : CHK_RET(HcclReduceAsync(
157 : dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, stream_,
158 : static_cast<void*>(dst.ptr()), links[neighCur]->GetRemoteRank(), links[neighCur]->GetLinkType(),
159 : INLINE_REDUCE_BIT));
160 0 : CHK_RET(links[neighCur]->TxDataSignal(stream_));
161 : } else {
162 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
163 0 : CHK_RET(links[neighCur]->TxAck(stream_));
164 0 : CHK_RET(links[neighCur]->RxDataSignal(stream_));
165 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, nextDst, dst, stream_));
166 : }
167 : } else {
168 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
169 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, nextDst, dst, stream_));
170 : }
171 :
172 0 : return HCCL_SUCCESS;
173 0 : }
174 :
175 0 : HcclResult AllReduceHDOptim::RunBetweenStep(
176 : u32 rank, u32 step, u32 neighBefore, u32 neighNext, u32 rankSize, const std::vector<LINK>& links)
177 : {
178 : (void)rank;
179 0 : u32 totalSize = SIZE_TABLE[dataType_] * count_;
180 :
181 0 : DeviceMem src;
182 0 : DeviceMem dst;
183 :
184 0 : if ((step == 1) && (stepPow == rankSize)) {
185 : // 二次幂整第一步写 串行同步
186 0 : CHK_RET(links[neighBefore]->TxDataSignal(stream_));
187 0 : CHK_RET(links[neighBefore]->RxDataSignal(stream_));
188 0 : CHK_RET(MainRecordSub(base));
189 0 : CHK_RET(SubWaitMain(base));
190 0 : } else {
191 0 : CHK_RET(MainRecordSub(base));
192 0 : CHK_RET(SubWaitMain(base));
193 0 : CHK_RET(links[neighBefore]->TxDataSignal(stream_));
194 0 : CHK_RET(links[neighBefore]->RxDataSignal(stream_));
195 : }
196 :
197 0 : src = outputMem_.range(step * totalSize, totalSize);
198 0 : if ((step == (nSteps - 1)) && (static_cast<u32>(pow(base, nSteps)) == rankSize)) {
199 0 : dst = userMemOut.range(0, totalSize);
200 : } else {
201 0 : dst = outputMem_.range((step + 1) * totalSize, totalSize);
202 : }
203 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, aicpu_ ? stream_ : meshStreams_[0]));
204 :
205 0 : CHK_RET(links[neighNext]->TxAck(aicpu_ ? stream_ : meshStreams_[1]));
206 0 : CHK_RET(links[neighNext]->RxAck(aicpu_ ? stream_ : meshStreams_[1]));
207 :
208 0 : CHK_RET(SubRecordMain(base));
209 0 : CHK_RET(MainWaitSub(base));
210 0 : return HCCL_SUCCESS;
211 0 : }
212 :
213 0 : HcclResult AllReduceHDOptim::RunAllReduceHDOptim(u32 rank, u32 rankSize, const std::vector<LINK>& links)
214 : {
215 0 : u32 unitSize = SIZE_TABLE[dataType_];
216 0 : u32 totalSize = unitSize * count_;
217 :
218 0 : DeviceMem src;
219 0 : DeviceMem dst;
220 :
221 0 : u32 neighCur = rank ^ (1 << 0);
222 0 : CHK_RET(links[neighCur]->TxAck(stream_));
223 0 : CHK_RET(links[neighCur]->RxAck(stream_));
224 :
225 0 : u32 neighNext = 0;
226 0 : void* remMemPtr = nullptr;
227 0 : for (u32 step = 1; step <= nSteps; step++) {
228 0 : if ((step != nSteps) || (stepPow != rankSize)) {
229 0 : dst = outputMem_.range(step * totalSize, totalSize);
230 : } else {
231 0 : dst = userMemOut.range(0, totalSize);
232 : }
233 0 : CHK_RET(links[neighCur]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
234 0 : src = DeviceMem::create(static_cast<u8*>(remMemPtr) + (step - 1) * totalSize, totalSize);
235 0 : if ((step == 1) && (stepPow == rankSize)) {
236 : // 二次幂整第一步写
237 0 : src = userMemIn.range(0, totalSize);
238 0 : dst = DeviceMem::create(static_cast<u8*>(remMemPtr) + totalSize, totalSize);
239 : }
240 :
241 0 : CHK_RET(HcclReduceAsync(
242 : dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, stream_,
243 : static_cast<void*>(dst.ptr()), links[neighCur]->GetRemoteRank(), links[neighCur]->GetLinkType(),
244 : INLINE_REDUCE_BIT));
245 :
246 0 : if (step != nSteps) {
247 0 : neighNext = rank ^ (1 << (step));
248 0 : CHK_RET(RunBetweenStep(rank, step, neighCur, neighNext, rankSize, links));
249 0 : neighCur = neighNext;
250 : }
251 : }
252 :
253 0 : CHK_RET(links[neighCur]->TxDataSignal(stream_));
254 0 : CHK_RET(links[neighCur]->RxDataSignal(stream_));
255 :
256 0 : return HCCL_SUCCESS;
257 0 : }
258 :
259 0 : HcclResult AllReduceHDOptim::RunFinalStep(u32 rank, u32 rankSize, const std::vector<LINK>& links)
260 : {
261 0 : u32 unitSize = SIZE_TABLE[dataType_];
262 0 : u32 totalSize = unitSize * count_;
263 :
264 0 : DeviceMem src = outputMem_.range(nSteps * totalSize, totalSize);
265 0 : DeviceMem dst = userMemOut.range(0, totalSize);
266 :
267 0 : u32 neighCur = rank ^ (1 << nSteps);
268 0 : if (neighCur >= rankSize) {
269 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
270 0 : } else if (rank < pow(base, nSteps)) {
271 0 : CHK_RET(links[neighCur]->TxAck(stream_));
272 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
273 0 : CHK_RET(links[neighCur]->RxDataSignal(stream_));
274 : } else {
275 0 : CHK_RET(links[neighCur]->RxAck(stream_));
276 0 : void* remMemPtr = nullptr;
277 0 : CHK_RET(links[neighCur]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
278 0 : src = DeviceMem::create(static_cast<u8*>(remMemPtr) + nSteps * totalSize, totalSize);
279 0 : CHK_RET(HcclD2DMemcpyAsync(
280 : dispatcher_, dst, src, stream_, links[neighCur]->GetRemoteRank(), links[neighCur]->GetLinkType()));
281 0 : CHK_RET(links[neighCur]->TxDataSignal(stream_));
282 : }
283 0 : return HCCL_SUCCESS;
284 0 : }
285 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_HD_OPTIM, AllReduceHDOptim);
286 : } // namespace hccl
|
