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_hd_stage_pub.h"
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 1 : ReduceScatterHDStage::ReduceScatterHDStage(const HcclDispatcher dispatcher)
17 1 : : AlgTemplateBase(dispatcher)
18 1 : {}
19 :
20 2 : ReduceScatterHDStage::~ReduceScatterHDStage()
21 2 : {}
22 :
23 1 : HcclResult ReduceScatterHDStage::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 ReduceScatterHDStage::MainRecordSub(u32 streamNum)
43 : {
44 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
45 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
46 : }
47 0 : return HCCL_SUCCESS;
48 : }
49 :
50 0 : HcclResult ReduceScatterHDStage::SubWaitMain(u32 streamNum)
51 : {
52 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
53 0 : CHK_RET(LocalNotify::Wait(
54 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
55 : }
56 0 : return HCCL_SUCCESS;
57 : }
58 :
59 0 : HcclResult ReduceScatterHDStage::MainWaitSub(u32 streamNum)
60 : {
61 0 : for (u32 signalIndex = 0; signalIndex < streamNum; signalIndex++) {
62 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
63 : }
64 0 : return HCCL_SUCCESS;
65 : }
66 :
67 0 : HcclResult ReduceScatterHDStage::SubRecordMain(u32 streamNum)
68 : {
69 0 : for (u32 streamIndex = 0; streamIndex < streamNum; streamIndex++) {
70 0 : CHK_RET(
71 : LocalNotify::Post(meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
72 : }
73 0 : return HCCL_SUCCESS;
74 : }
75 :
76 : // ringallreduce算法的函数入口
77 0 : HcclResult ReduceScatterHDStage::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
78 : {
79 0 : HcclResult ret = HCCL_SUCCESS;
80 0 : CHK_SMART_PTR_NULL(dispatcher_);
81 0 : CHK_PTR_NULL(stream_.ptr());
82 0 : HCCL_INFO("ReduceScatterHDStage run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]",
83 : rank,
84 : rankSize,
85 : inputMem_.ptr(),
86 : outputMem_.ptr(),
87 : count_);
88 :
89 0 : if (links.size() < rankSize) {
90 0 : HCCL_ERROR("[ReduceScatterHDStage][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]",
91 : rank,
92 : links.size(),
93 : rankSize);
94 0 : return HCCL_E_INTERNAL;
95 : }
96 :
97 0 : ret = PrepareSliceData(rankSize);
98 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
99 : HCCL_ERROR("[ReduceScatterHDStage][RunAsync]rank[%u] count[%llu] failed in PrepareSliceData "
100 : "step",
101 : rank,
102 : count_),
103 : ret);
104 :
105 0 : ret = RunReduceScatterStage(rank, rankSize, links);
106 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
107 : HCCL_ERROR("[ReduceScatterHDStage][RunAsync]rank[%u] count[%llu] failed"
108 : "step",
109 : rank,
110 : count_),
111 : ret);
112 :
113 0 : HCCL_INFO("ReduceScatterHDStage finished: rank[%u] ranksize[%u]", rank, rankSize);
114 0 : return HCCL_SUCCESS;
115 : }
116 :
117 0 : HcclResult ReduceScatterHDStage::PrepareSliceData(u32 rankSize)
118 : {
119 0 : Slice temp;
120 0 : u32 unitSize = SIZE_TABLE[dataType_];
121 0 : u64 totalSize = count_ * unitSize;
122 0 : u32 power = static_cast<u32>(log2(rankSize));
123 0 : u32 half = static_cast<u32>(pow(base, power - 1));
124 : u64 offset;
125 0 : for (u32 round = 1; round <= power; round++) {
126 0 : u32 sliceNum = rankSize / static_cast<u32>(pow(base, round));
127 0 : sliceMap_[round - 1].clear();
128 0 : sliceMap_[power - round].reserve(rankSize);
129 0 : for (u32 sliceGroup = 0; sliceGroup < pow(base, round); sliceGroup++) {
130 0 : for (u32 sliceCount = 0; sliceCount < sliceNum; sliceCount++) {
131 0 : temp.size = totalSize * sliceNum;
132 0 : offset = totalSize * sliceNum * sliceGroup;
133 0 : if (sliceGroup == 0) {
134 0 : temp.offset = offset;
135 : } else {
136 0 : if (round != 1) {
137 0 : temp.offset = (offset >= half * totalSize) ? (offset - half * totalSize) : offset;
138 : } else {
139 0 : temp.offset = offset;
140 : }
141 : }
142 0 : sliceMap_[round - 1].push_back(temp);
143 : }
144 : }
145 : }
146 0 : return HCCL_SUCCESS;
147 : }
148 :
149 0 : HcclResult ReduceScatterHDStage::RunReduceScatterStage(u32 rank, u32 rankSize, const std::vector<LINK> &links)
150 : {
151 0 : HCCL_INFO("RunReduceScatterStage run: rank[%u] totalrank[%u] outputMem[%p] count[%llu]",
152 : rank, rankSize, outputMem_.ptr(), count_);
153 0 : nSteps_ = static_cast<u32>(log2(rankSize));
154 0 : CHK_RET(RunReduceScatterStage1st(rank, rankSize, links));
155 0 : CHK_RET(RunReduceScatterRead(rank, rankSize, links));
156 :
157 0 : return HCCL_SUCCESS;
158 : }
159 :
160 0 : HcclResult ReduceScatterHDStage::RunReduceScatterStage1st(u32 rank, u32 rankSize, const std::vector<LINK> &links)
161 : {
162 0 : u32 unitSize = SIZE_TABLE[dataType_];
163 0 : u32 totalSize = unitSize * count_;
164 :
165 0 : DeviceMem UserMemIn = DeviceMem::create(opInfo_->inputAddr, rankSize * totalSize);
166 0 : DeviceMem CommMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
167 :
168 0 : DeviceMem src;
169 0 : DeviceMem dst;
170 :
171 0 : src = UserMemIn.range(sliceMap_[0][rank].offset, sliceMap_[0][rank].size);
172 0 : dst = CommMemOut.range(0, sliceMap_[0][rank].size);
173 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
174 :
175 0 : u32 dstRank = rank ^ (1 << (nSteps_ - 1));
176 0 : CHK_RET(links[dstRank]->TxAck(stream_));
177 0 : CHK_RET(links[dstRank]->RxAck(stream_));
178 :
179 0 : void *remMemPtr = nullptr;
180 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
181 0 : src = UserMemIn.range(sliceMap_[0][dstRank].offset, sliceMap_[0][dstRank].size);
182 0 : dst = DeviceMem::create(static_cast<u8 *>(remMemPtr), sliceMap_[0][dstRank].size);
183 0 : CHK_RET(HcclReduceAsync(dispatcher_, static_cast<void *>(src.ptr()), sliceMap_[0][dstRank].size / unitSize,
184 : dataType_, reductionOp_, stream_, static_cast<void *>(dst.ptr()),
185 : links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType(), INLINE_REDUCE_BIT));
186 :
187 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
188 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
189 0 : return HCCL_SUCCESS;
190 0 : }
191 :
192 0 : HcclResult ReduceScatterHDStage::RunBetweenStep(u32 rank, u32 neighCur, u32 neighNext, const std::vector<LINK> &links)
193 : {
194 : (void) rank;
195 0 : CHK_RET(MainRecordSub(1));
196 0 : CHK_RET(SubWaitMain(1));
197 :
198 0 : CHK_RET(links[neighCur]->TxDataSignal(meshStreams_[0]));
199 0 : CHK_RET(links[neighCur]->RxDataSignal(meshStreams_[0]));
200 :
201 0 : CHK_RET(links[neighNext]->TxAck(stream_));
202 0 : CHK_RET(links[neighNext]->RxAck(stream_));
203 :
204 0 : CHK_RET(SubRecordMain(1));
205 0 : CHK_RET(MainWaitSub(1));
206 :
207 0 : return HCCL_SUCCESS;
208 : }
209 :
210 0 : HcclResult ReduceScatterHDStage::RunReduceScatterRead(u32 rank, u32 rankSize, const std::vector<LINK> &links)
211 : {
212 : (void) rankSize;
213 0 : u32 unitSize = SIZE_TABLE[dataType_];
214 0 : u32 totalSize = unitSize * count_;
215 :
216 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, totalSize);
217 :
218 0 : void *remMemPtr = nullptr;
219 0 : DeviceMem dst;
220 0 : DeviceMem src;
221 : u32 dstRank;
222 0 : CHK_RET(links[rank ^ (1 << (nSteps_ - 1 - 1))]->TxAck(stream_));
223 0 : CHK_RET(links[rank ^ (1 << (nSteps_ - 1 - 1))]->RxAck(stream_));
224 0 : for (u32 step = 1; step < nSteps_; step++) {
225 0 : dstRank = rank ^ (1 << (nSteps_ - 1 - step));
226 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
227 0 : dst = outputMem_.range(sliceMap_[step][rank].offset, sliceMap_[step][rank].size);
228 0 : src = DeviceMem::create(static_cast<u8 *>(remMemPtr) + sliceMap_[step][rank].offset,
229 0 : sliceMap_[step][rank].size);
230 0 : CHK_RET(HcclReduceAsync(dispatcher_, static_cast<void *>(src.ptr()), sliceMap_[step][rank].size / unitSize,
231 : dataType_, reductionOp_, stream_, static_cast<void *>(dst.ptr()),
232 : links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType(), INLINE_REDUCE_BIT));
233 0 : if (step != (nSteps_ - 1)) {
234 0 : CHK_RET(RunBetweenStep(rank, dstRank, rank ^ (1 << (nSteps_ - 1 - step - 1)), links));
235 : }
236 : }
237 :
238 0 : CHK_RET(MainRecordSub(1));
239 0 : CHK_RET(SubWaitMain(1));
240 :
241 0 : CHK_RET(links[rank ^ (1 << 0)]->TxDataSignal(stream_));
242 0 : CHK_RET(links[rank ^ (1 << 0)]->RxDataSignal(stream_));
243 :
244 0 : src = outputMem_.range(sliceMap_[nSteps_ - 1][rank].offset, totalSize);
245 0 : dst = userMemOut.range(0, totalSize);
246 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, meshStreams_[0]));
247 :
248 0 : CHK_RET(SubRecordMain(1));
249 0 : CHK_RET(MainWaitSub(1));
250 :
251 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
252 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
253 :
254 0 : return HCCL_SUCCESS;
255 0 : }
256 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_HDSTAGE, ReduceScatterHDStage);
257 : } // namespace hccl
|
