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