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