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 "reduce_scatter_local_reduce_pub.h"
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 1 : ReduceScatterLocalReduce::ReduceScatterLocalReduce(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
17 :
18 2 : ReduceScatterLocalReduce::~ReduceScatterLocalReduce() {}
19 :
20 1 : HcclResult ReduceScatterLocalReduce::Prepare(
21 : DeviceMem& inputMem, DeviceMem& outputMem, DeviceMem& scratchMem, const u64 count, const HcclDataType dataType,
22 : const Stream& stream, const HcclReduceOp reductionOp, const u32 root, const std::vector<Slice>& slices,
23 : const u64 baseOffset, const u64 reduceAttrBitMap, std::vector<Stream>& meshStreams,
24 : std::vector<std::shared_ptr<LocalNotify>>& meshSignal, std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux,
25 : u32 userRank, const HcomCollOpInfo* opInfo)
26 : {
27 1 : reduceAttr_ = reduceAttrBitMap;
28 1 : userRank_ = userRank;
29 1 : meshStreams_ = meshStreams;
30 1 : meshSignalPtr_ = &meshSignal;
31 1 : meshSignalAuxPtr_ = &meshSignalAux;
32 1 : opInfo_ = opInfo;
33 4 : return AlgTemplateBase::Prepare(
34 2 : inputMem, outputMem, scratchMem, count, dataType, stream, reductionOp, root, slices, baseOffset);
35 : }
36 :
37 0 : HcclResult ReduceScatterLocalReduce::MainRecordSub(u32 streamNum)
38 : {
39 0 : u32 totalTask = streamNum;
40 0 : CHK_PRT_RET(
41 : (totalTask > meshSignalAuxPtr_->size()),
42 : HCCL_ERROR(
43 : "[ReduceScatterLocalReduce][MainRecordSub]totalTask[%u] is over range of meshSignalAux[%zu]", totalTask,
44 : meshSignalAuxPtr_->size()),
45 : HCCL_E_PARA);
46 0 : for (u32 signalIndex = 0; signalIndex < totalTask; signalIndex++) {
47 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
48 : }
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 0 : HcclResult ReduceScatterLocalReduce::SubWaitMain(u32 streamNum)
53 : {
54 0 : u32 totalTask = streamNum;
55 0 : CHK_PRT_RET(
56 : (totalTask > meshSignalAuxPtr_->size() || totalTask > meshStreams_.size()),
57 : HCCL_ERROR(
58 : "[ReduceScatterLocalReduce][SubWaitMain]totalTask[%u] is over range of meshSignalAux[%zu]"
59 : "or meshStreams_[%zu]",
60 : totalTask, meshSignalAuxPtr_->size(), meshStreams_.size()),
61 : HCCL_E_PARA);
62 0 : for (u32 streamIndex = 0; streamIndex < totalTask; streamIndex++) {
63 0 : CHK_RET(LocalNotify::Wait(
64 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
65 : }
66 0 : return HCCL_SUCCESS;
67 : }
68 :
69 0 : HcclResult ReduceScatterLocalReduce::MainWaitSub(u32 streamNum)
70 : {
71 0 : u32 totalTask = streamNum;
72 0 : CHK_PRT_RET(
73 : (totalTask > meshSignalPtr_->size()),
74 : HCCL_ERROR(
75 : "[ReduceScatterLocalReduce][MainWaitSub]totalTask[%u] is over range of meshSignal[%zu]", totalTask,
76 : meshSignalPtr_->size()),
77 : HCCL_E_PARA);
78 0 : for (u32 signalIndex = 0; signalIndex < totalTask; signalIndex++) {
79 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
80 : }
81 0 : return HCCL_SUCCESS;
82 : }
83 :
84 0 : HcclResult ReduceScatterLocalReduce::SubRecordMain(u32 streamNum)
85 : {
86 0 : u32 totalTask = streamNum;
87 0 : CHK_PRT_RET(
88 : (totalTask > meshSignalPtr_->size() || totalTask > meshStreams_.size()),
89 : HCCL_ERROR(
90 : "[ReduceScatterLocalReduce][SubWaitMain]totalTask[%u] is over range of meshSignal[%zu]"
91 : "or meshStreams_[%zu]",
92 : totalTask, meshSignalPtr_->size(), meshStreams_.size()),
93 : HCCL_E_PARA);
94 0 : for (u32 streamIndex = 0; streamIndex < totalTask; streamIndex++) {
95 0 : CHK_RET(LocalNotify::Post(
96 : meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
97 : }
98 0 : return HCCL_SUCCESS;
99 : }
100 :
101 : // 计算每片数据的offset
102 0 : HcclResult ReduceScatterLocalReduce::PrepareOffset(u32 rankSize)
103 : {
104 0 : Slice temp;
105 0 : u32 unitSize = SIZE_TABLE[dataType_];
106 0 : u64 totalSize = (opInfo_->count) * unitSize;
107 0 : slices_.clear();
108 0 : slices_.reserve(rankSize);
109 0 : if (rankSize == 0) {
110 0 : HCCL_ERROR("[Prepare][Offset]data slice prepare, sliceNum is 0");
111 0 : return HCCL_E_PARA;
112 : }
113 0 : for (u32 i = 0; i < rankSize; i++) {
114 0 : if (count_ * SIZE_TABLE[dataType_] > HCCL_SMALL_COUNT_32_KB) {
115 0 : temp.offset = (i * totalSize) % HCCL_MIN_SLICE_ALIGN_910B;
116 : } else {
117 0 : temp.offset = 0;
118 : }
119 0 : temp.size = 0;
120 0 : slices_.push_back(temp);
121 : }
122 0 : return HCCL_SUCCESS;
123 : }
124 :
125 0 : HcclResult ReduceScatterLocalReduce::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
126 : {
127 0 : HcclResult ret = HCCL_SUCCESS;
128 0 : CHK_SMART_PTR_NULL(dispatcher_);
129 0 : CHK_PTR_NULL(stream_.ptr());
130 0 : HCCL_INFO(
131 : "ReduceScatterLocalReduce run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu].", rank, rankSize,
132 : inputMem_.ptr(), outputMem_.ptr(), count_);
133 :
134 0 : if (links.size() < rankSize) {
135 0 : HCCL_ERROR(
136 : "[ReduceScatterLocalReduce][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(),
137 : rankSize);
138 0 : return HCCL_E_INTERNAL;
139 : }
140 :
141 0 : ret = PrepareOffset(rankSize);
142 0 : CHK_PRT_RET(
143 : ret != HCCL_SUCCESS,
144 : HCCL_ERROR(
145 : "[ReduceScatterLocalReduce][RunAsync]rank[%u] count[%llu] failed in PrepareOffset step", rank, count_),
146 : ret);
147 :
148 0 : ret = RunReduceScatter(rank, rankSize, links);
149 0 : CHK_PRT_RET(
150 : ret != HCCL_SUCCESS,
151 : HCCL_ERROR(
152 : "[ReduceScatterLocalReduce][RunAsync]rank[%u] count[%llu] failed in ReduceScatter step", rank, count_),
153 : ret);
154 :
155 0 : ret = RunLocalReduce(rank, rankSize);
156 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceScatterlocalReduce]rank[%u] LocalReduce failed", rank), ret);
157 :
158 0 : HCCL_INFO("ReduceScatterLocalReduce finished: rank[%u] ranksize[%u]", rank, rankSize);
159 0 : return HCCL_SUCCESS;
160 : }
161 :
162 0 : HcclResult ReduceScatterLocalReduce::CalAlign(u64 totalSize, u32 rankSize, u64& alignSize)
163 : {
164 0 : auto maxIt = std::max_element(slices_.begin(), slices_.end(), [](const Slice& slice1, const Slice& slice2) {
165 0 : return slice1.offset < slice2.offset;
166 : });
167 :
168 0 : u64 maxOffset = maxIt->offset;
169 :
170 0 : alignSize = RoundUpWithDivisor(totalSize, HCCL_MIN_SLICE_ALIGN_910B);
171 0 : if (alignSize * (rankSize - 1) > (outputMem_.size() - maxOffset)) {
172 0 : alignSize = RoundUpWithDivisor(totalSize, HCCL_MIN_SLICE_ALIGN_ONCHIP);
173 : }
174 0 : if (alignSize * (rankSize - 1) > (outputMem_.size() - maxOffset)) {
175 0 : alignSize = totalSize;
176 : }
177 0 : return HCCL_SUCCESS;
178 : }
179 :
180 0 : HcclResult ReduceScatterLocalReduce::RunReduceScatter(u32 rank, u32 rankSize, const std::vector<LINK>& links)
181 : {
182 0 : HCCL_INFO(
183 : "ReduceScatterLocalReduce run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu].", rank, rankSize,
184 : inputMem_.ptr(), outputMem_.ptr(), count_);
185 :
186 : // 数据准备
187 0 : u64 unitSize = SIZE_TABLE[dataType_];
188 0 : u64 totalSize = count_ * unitSize;
189 0 : u64 alignSize = totalSize;
190 0 : CHK_RET(CalAlign(totalSize, rankSize, alignSize));
191 0 : u64 offset = (opInfo_->count) * unitSize;
192 0 : DeviceMem UserMemIn = DeviceMem::create(opInfo_->inputAddr, offset * rankSize);
193 0 : DeviceMem CommMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
194 0 : DeviceMem UserMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize);
195 :
196 0 : DeviceMem src;
197 0 : DeviceMem dst;
198 :
199 0 : DeviceMem emptySrc = UserMemIn.range(0, 0);
200 0 : DeviceMem emptyDst = CommMemOut.range(0, 0);
201 :
202 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
203 :
204 0 : CHK_RET(MainRecordSub(rankSize - base));
205 0 : CHK_RET(SubWaitMain(rankSize - base));
206 :
207 0 : for (u32 round = 1; round < rankSize; round++) {
208 0 : u32 dstRank = (round + rank) % rankSize;
209 0 : Stream& subStream = (round == rankSize - 1) ? stream_ : meshStreams_[round - 1];
210 0 : CHK_RET(links[dstRank]->TxAck(subStream));
211 0 : CHK_RET(links[dstRank]->RxAck(subStream));
212 : }
213 :
214 0 : CHK_RET(SubRecordMain(rankSize - base));
215 0 : CHK_RET(MainWaitSub(rankSize - base));
216 :
217 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
218 :
219 0 : CHK_RET(MainRecordSub(meshStreams_.size()));
220 0 : CHK_RET(SubWaitMain(meshStreams_.size()));
221 :
222 0 : for (u32 round = 1; round < rankSize; round++) {
223 0 : Stream& subStream = (round == rankSize - 1) ? stream_ : meshStreams_[round - 1];
224 :
225 0 : u32 dstRank = (round + rank) % rankSize;
226 :
227 0 : void* remMemPtr = nullptr;
228 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
229 :
230 0 : dst = DeviceMem::create(
231 0 : static_cast<u8*>(remMemPtr) + alignSize * (round - 1) + slices_[dstRank].offset, totalSize);
232 0 : src = UserMemIn.range(offset * dstRank, totalSize);
233 0 : CHK_RET(HcclD2DMemcpyAsync(
234 : dispatcher_, dst, src, subStream, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
235 :
236 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
237 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
238 : }
239 :
240 0 : src = UserMemIn.range(offset * rank, totalSize);
241 0 : dst = UserMemOut.range(0, totalSize);
242 0 : Stream& subStream = (meshStreams_.size() > 0) ? meshStreams_[meshStreams_.size() - 1] : stream_;
243 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStream));
244 :
245 0 : CHK_RET(SubRecordMain(meshStreams_.size()));
246 0 : CHK_RET(MainWaitSub(meshStreams_.size()));
247 :
248 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
249 0 : return HCCL_SUCCESS;
250 0 : }
251 :
252 0 : HcclResult ReduceScatterLocalReduce::RunLocalReduce(u32 rank, u32 rankSize)
253 : {
254 0 : u32 power = static_cast<u32>(log2(rankSize));
255 0 : u32 rankPower = static_cast<u32>(pow(base, power));
256 0 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
257 0 : u64 alignSize = totalSize;
258 0 : CHK_RET(CalAlign(totalSize, rankSize, alignSize));
259 0 : DeviceMem CommMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
260 0 : DeviceMem UserMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize);
261 0 : CommMemOut = CommMemOut.range(slices_[rank].offset, outputMem_.size() - slices_[rank].offset);
262 0 : DeviceMem src;
263 0 : DeviceMem dst;
264 0 : DeviceMem emptySrc = CommMemOut.range(0, 0);
265 0 : DeviceMem emptyDst = CommMemOut.range(0, 0);
266 0 : if (rankPower < rankSize) {
267 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
268 0 : CHK_RET(MainRecordSub(rankSize - rankPower - 1));
269 0 : CHK_RET(SubWaitMain(rankSize - rankPower - 1));
270 0 : for (u32 add = 0; add < (rankSize - rankPower); add++) {
271 0 : Stream& subStream = (add == 0) ? stream_ : meshStreams_[add - 1];
272 0 : src = CommMemOut.range(alignSize * (add + rankPower - 1), totalSize);
273 0 : if (add == 0) {
274 0 : dst = UserMemOut.range(0, totalSize);
275 : } else {
276 0 : dst = CommMemOut.range(alignSize * (add - 1), totalSize);
277 : }
278 0 : CHK_RET(HcclReduceAsync(
279 : dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, subStream,
280 : static_cast<void*>(dst.ptr()), INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP, INLINE_REDUCE_BIT));
281 : }
282 0 : CHK_RET(SubRecordMain(rankSize - rankPower - 1));
283 0 : CHK_RET(MainWaitSub(rankSize - rankPower - 1));
284 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
285 : }
286 0 : for (u32 round = 0; round < power; round++) {
287 0 : rankPower = static_cast<u32>(pow(base, power - round - 1));
288 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
289 0 : CHK_RET(MainRecordSub(rankPower - 1));
290 0 : CHK_RET(SubWaitMain(rankPower - 1));
291 0 : for (u32 add = 0; add < rankPower; add++) {
292 0 : Stream& subStream = (add == 0) ? stream_ : meshStreams_[add - 1];
293 0 : src = CommMemOut.range(alignSize * (add + rankPower - 1), totalSize);
294 0 : dst = (add == 0) ? UserMemOut.range(0, totalSize) : CommMemOut.range(alignSize * (add - 1), totalSize);
295 0 : CHK_RET(HcclReduceAsync(
296 : dispatcher_, static_cast<void*>(src.ptr()), count_, dataType_, reductionOp_, subStream,
297 : static_cast<void*>(dst.ptr()), INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP, INLINE_REDUCE_BIT));
298 : }
299 0 : CHK_RET(SubRecordMain(rankPower - 1));
300 0 : CHK_RET(MainWaitSub(rankPower - 1));
301 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyDst, emptySrc, stream_));
302 : }
303 0 : return HCCL_SUCCESS;
304 0 : }
305 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_LOCAL_REDUCE, ReduceScatterLocalReduce);
306 : } // namespace hccl
|
