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 "reduce_scatter_halving_doubling.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : ReduceScatterHalvingDoubling::ReduceScatterHalvingDoubling(const HcclDispatcher dispatcher)
16 0 : : AlgTemplateBase(dispatcher)
17 0 : {}
18 :
19 0 : ReduceScatterHalvingDoubling::~ReduceScatterHalvingDoubling() {}
20 :
21 0 : HcclResult ReduceScatterHalvingDoubling::Prepare(
22 : DeviceMem& inputMem, DeviceMem& outputMem, DeviceMem& scratchMem, const u64 count, const HcclDataType dataType,
23 : const Stream& stream, const HcclReduceOp reductionOp, const u32 root, const std::vector<Slice>& slices,
24 : const u64 baseOffset, const u32 blockSize, const u64 reduceAttrBitMap, const UserMemType hdInputMemType,
25 : const UserMemType hdOutputMemType)
26 : {
27 0 : blockSize_ = blockSize;
28 0 : reduceAttr_ = reduceAttrBitMap;
29 0 : hdInputMemType_ = hdInputMemType;
30 0 : hdOutputMemType_ = hdOutputMemType;
31 0 : return AlgTemplateBase::Prepare(
32 0 : inputMem, outputMem, scratchMem, count, dataType, stream, reductionOp, root, slices, baseOffset);
33 : }
34 :
35 0 : u32 ReduceScatterHalvingDoubling::GetBlockStep(u32 blocksize) const
36 : {
37 : // 求以2为底数的对数计算
38 0 : u32 step = 0;
39 0 : while ((blocksize >> (step + 1)) != 0) {
40 0 : step++;
41 : }
42 :
43 0 : return step;
44 : }
45 :
46 : HcclResult
47 0 : ReduceScatterHalvingDoubling::CalculateSlices(const u64 size, const u32 sliceNum, std::vector<Slice>& slicesOut)
48 : {
49 0 : CHK_PRT_RET((sliceNum == 0), HCCL_ERROR("[Calculate][Slices]calculate_slices failed"), HCCL_E_INTERNAL);
50 :
51 : // 不对size, count和slice_num做检查, 默认满足reduce-scatter的要求
52 0 : std::vector<Slice> slices(sliceNum);
53 :
54 0 : u64 sliceSize = size / sliceNum;
55 :
56 0 : for (u32 i = 0; i < sliceNum; i++) {
57 0 : slices[i].size = sliceSize;
58 0 : slices[i].offset = i * sliceSize;
59 : }
60 :
61 0 : slicesOut = std::move(slices);
62 0 : return HCCL_SUCCESS;
63 0 : }
64 :
65 0 : HcclResult ReduceScatterHalvingDoubling::CalcStepSlices(
66 : const std::vector<Slice>& inputSlices, const u32 stepNum, const u32 rank, const SliceType type,
67 : std::vector<Slice>& slicesOut)
68 : {
69 0 : std::vector<Slice> slice(stepNum);
70 :
71 0 : for (u32 step = 0; step < stepNum; step++) {
72 : // reduce-scatter操作, halving_bitmask从高往低循环, size倍减
73 0 : u32 halvingBitmask = (1 << (stepNum - step - 1));
74 0 : u32 peerRank = rank ^ halvingBitmask;
75 :
76 : // 计算tx_slice/rx_slice
77 0 : u32 sliceId = (type == SliceType::SLICE_TYPE_TX) ? (peerRank & (~(halvingBitmask - 1))) :
78 0 : (rank & (~(halvingBitmask - 1)));
79 :
80 0 : slice[step].offset = inputSlices[sliceId].offset;
81 0 : HcclResult ret = Sum(inputSlices, sliceId, halvingBitmask, slice[step].size);
82 0 : CHK_PRT_RET(
83 : ret != HCCL_SUCCESS,
84 : HCCL_ERROR("[Calc][StepSlices]rank[%u] ReduceScatter Halving Doubling sum error", rank), ret);
85 0 : HCCL_DEBUG(
86 : "rank[%u] step[%u] type[%u] slice[%u].offset[%llu] slice[%u].size[%llu] ", rank, step, type, step,
87 : slice[step].offset, step, slice[step].size);
88 : }
89 :
90 0 : slicesOut = std::move(slice);
91 0 : return HCCL_SUCCESS;
92 0 : }
93 :
94 0 : HcclResult ReduceScatterHalvingDoubling::RunSourceReducer(const LINK& link, const Slice& txSlice)
95 : {
96 0 : HcclResult ret = link->RxAck(stream_);
97 0 : CHK_PRT_RET(
98 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][SourceReducer]txSlice.size[%llu] rx ack run failed", txSlice.size), ret);
99 :
100 0 : DeviceMem txMem = inputMem_.range(txSlice.offset, txSlice.size);
101 :
102 0 : HCCL_DEBUG(
103 : "tx_slice.offset[%llu],base_offset[%llu],tx_slice.size[%llu]", txSlice.offset, baseOffset_, txSlice.size);
104 :
105 : // 发送到对端的output
106 0 : CHK_RET(senderInfo_->run(link, txSlice.offset + baseOffset_, txMem, stream_));
107 :
108 0 : return HCCL_SUCCESS;
109 0 : }
110 :
111 0 : HcclResult ReduceScatterHalvingDoubling::RunDestRducer(
112 : const LINK& link, const HcclDispatcher dispatcher, const Slice& rxSlice, const DstMemType reduceDst)
113 : {
114 : (void)reduceDst;
115 0 : DeviceMem rxMem = scratchMem_.range(rxSlice.offset, rxSlice.size);
116 0 : HCCL_DEBUG("rx_mem.offset[%llu],base_offset[%llu],rx_slice.size[%llu]", rxSlice.offset, baseOffset_, rxSlice.size);
117 :
118 0 : DeviceMem reduceSrcMem;
119 0 : DeviceMem reduceDstMem;
120 :
121 0 : if (DataUnitSize(dataType_) == 0) {
122 0 : HCCL_ERROR("[Run][DestRducer]DataUnitSize(data_type_) == 0, error");
123 0 : return HCCL_E_INTERNAL;
124 : }
125 :
126 0 : if (link->IsSpInlineReduce() && static_cast<bool>((INLINE_REDUCE_BITMASK & reduceAttr_))) {
127 0 : reduceSrcMem = inputMem_.range(rxSlice.offset, rxSlice.size);
128 : } else {
129 0 : reduceSrcMem = outputMem_.range(rxSlice.offset, rxSlice.size);
130 : }
131 0 : reduceDstMem = inputMem_.range(rxSlice.offset, rxSlice.size);
132 :
133 0 : HcclResult ret = reducerInfo_->run(
134 0 : dispatcher, link, rxSlice.offset + baseOffset_, reduceSrcMem, reduceDstMem, rxMem, stream_,
135 : DstMemType::RESULT_INPUT_MEM);
136 :
137 0 : CHK_PRT_RET(
138 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][DestRducer]offset[%llu] reduce_async failed", rxSlice.offset), ret);
139 :
140 0 : return HCCL_SUCCESS;
141 0 : }
142 :
143 0 : HcclResult ReduceScatterHalvingDoubling::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
144 : {
145 0 : CHK_SMART_PTR_NULL(dispatcher_);
146 0 : CHK_PTR_NULL(stream_.ptr());
147 0 : if (!outputMem_ || !inputMem_) {
148 0 : HCCL_ERROR("[ReduceScatterHalvingDoubling][RunAsync]rank[%u] run_async inputmem or outputmem is null", rank);
149 0 : return HCCL_E_PTR;
150 : }
151 :
152 0 : HCCL_INFO(
153 : "ReduceScatterHalvingDoubling run: rank[%u] totalrank[%u] "
154 : "inputMem[%p] outputMem[%p] count[%llu]",
155 : rank, rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
156 0 : HcclResult ret = HCCL_SUCCESS;
157 :
158 : // 仅一个rank, 则直接input拷贝到output
159 0 : if (rankSize == 1) {
160 0 : if (inputMem_ != outputMem_) {
161 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
162 : }
163 0 : return ret;
164 : }
165 :
166 : // 创建reducer & sender
167 0 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
168 0 : CHK_SMART_PTR_NULL(senderInfo_);
169 :
170 0 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
171 0 : CHK_SMART_PTR_NULL(reducerInfo_);
172 :
173 0 : bool bRetSize = (links.size() < rankSize);
174 0 : CHK_PRT_RET(
175 : bRetSize,
176 : HCCL_ERROR("[ReduceScatterHalvingDoubling][RunAsync]rank[%u] linksize[%llu] is error", rank, links.size()),
177 : HCCL_E_INTERNAL);
178 :
179 : // 检查是否已对数据分片
180 0 : if (slices_.size() != rankSize) {
181 0 : CHK_RET(CalculateSlices(inputMem_.size(), rankSize, slices_));
182 : }
183 :
184 : // 计算每个step的数据size
185 0 : u32 stepNum = GetBlockStep(rankSize);
186 0 : CHK_RET(CalcStepSlices(slices_, stepNum, rank, SliceType::SLICE_TYPE_TX, txSlices_));
187 :
188 0 : CHK_RET(CalcStepSlices(slices_, stepNum, rank, SliceType::SLICE_TYPE_RX, rxSlices_));
189 :
190 0 : CHK_RET(RunReduceScatter(rank, stepNum, dispatcher_, links));
191 :
192 0 : HCCL_INFO("ReduceScatterHalvingDoubling rank[%u] finished", rank);
193 0 : return HCCL_SUCCESS;
194 : }
195 :
196 0 : HcclResult ReduceScatterHalvingDoubling::RunReduceScatter(
197 : const u32 rank, const u32 stepNum, const HcclDispatcher dispatcher, const std::vector<LINK>& links)
198 : {
199 0 : HcclResult ret = HCCL_SUCCESS;
200 :
201 0 : for (u32 step = 0; step < stepNum; step++) {
202 : // reduce-scatter操作, peer_rank_bitmask从高往低循环
203 0 : u32 peerRankBitmask = 1 << (stepNum - step - 1);
204 0 : u32 peerRank = rank ^ peerRankBitmask;
205 :
206 0 : CHK_SMART_PTR_NULL(links[peerRank]);
207 :
208 0 : ret = links[peerRank]->TxAck(stream_);
209 0 : CHK_PRT_RET(
210 : ret != HCCL_SUCCESS,
211 : HCCL_ERROR(
212 : "[Run][ReduceScatter]rank[%u] tx ack to peerrank[%u] in step[%u] "
213 : "run failed",
214 : rank, peerRank, step),
215 : ret);
216 :
217 0 : HCCL_DEBUG(
218 : "rank[%u] send to peerrank[%u] in step[%u], silce.offset[%llu], slice.size[%llu]", rank, peerRank, step,
219 : txSlices_[step].offset, txSlices_[step].size);
220 : // 本rank作为reducer的发送侧的动作
221 0 : ret = RunSourceReducer(links[peerRank], txSlices_[step]);
222 0 : CHK_PRT_RET(
223 : ret != HCCL_SUCCESS,
224 : HCCL_ERROR(
225 : "[Run][ReduceScatter]rank[%u] to peerrank[%u] reducer_src_run "
226 : "failed",
227 : rank, peerRank),
228 : ret);
229 :
230 : // 本rank作为reducer的接收侧的动作, 结果除最后一轮存放至output外, 均存放至input
231 0 : DstMemType reduceDst = (step == stepNum - 1) ? DstMemType::RESULT_OUTPUT_MEM : DstMemType::RESULT_INPUT_MEM;
232 :
233 0 : HCCL_DEBUG(
234 : "rank[%u] Reduce from peerrank[%u] in step[%u], silce.offset[%llu], slice.size[%llu]", rank, peerRank, step,
235 : rxSlices_[step].offset, rxSlices_[step].size);
236 :
237 0 : ret = RunDestRducer(links[peerRank], dispatcher, rxSlices_[step], reduceDst);
238 0 : CHK_PRT_RET(
239 : ret != HCCL_SUCCESS,
240 : HCCL_ERROR(
241 : "[Run][ReduceScatter]rank[%u] to peerrank[%u] reducer_dst_run "
242 : "failed",
243 : rank, peerRank),
244 : ret);
245 0 : ret = links[peerRank]->RxWaitDone(stream_);
246 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]RxWaitDone failed"), ret);
247 0 : ret = links[peerRank]->TxWaitDone(stream_);
248 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]TxWaitDone failed"), ret);
249 : }
250 :
251 0 : DeviceMem reduceSrcMem = inputMem_.range(rxSlices_[stepNum - 1].offset, rxSlices_[stepNum - 1].size);
252 0 : DeviceMem reduceDstMem = outputMem_.range(rxSlices_[stepNum - 1].offset, rxSlices_[stepNum - 1].size);
253 0 : ret = HcclD2DMemcpyAsync(dispatcher_, reduceDstMem, reduceSrcMem, stream_);
254 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]HcclD2DMemcpyAsync failed"), ret);
255 :
256 0 : return ret;
257 0 : }
258 :
259 0 : HcclResult ReduceScatterHalvingDoubling::GetNslbAdjInfo(
260 : [[maybe_unused]] const u32 rank, [[maybe_unused]] const u32 rankSize,
261 : [[maybe_unused]] const std::vector<LINK>& links, [[maybe_unused]] AdjInfo& nslbAdjInfo)
262 : {
263 0 : return HCCL_SUCCESS;
264 : }
265 :
266 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_HD, ReduceScatterHalvingDoubling);
267 : } // namespace hccl
|