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