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_mesh.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : ReduceScatterMesh::ReduceScatterMesh(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : ReduceScatterMesh::~ReduceScatterMesh() {}
18 :
19 0 : HcclResult ReduceScatterMesh::Prepare(u64 reduceAttrBitMap, u32 streamIndex)
20 : {
21 0 : reduceAttr_ = reduceAttrBitMap;
22 0 : streamIndex_ = streamIndex;
23 0 : return HCCL_SUCCESS;
24 : }
25 :
26 0 : HcclResult ReduceScatterMesh::RunSourceReducer(const LINK& link, const Slice& txSlice, const Slice& dstSlice)
27 : {
28 : // 发送inputmem
29 0 : DeviceMem srcMem = inputMem_.range(txSlice.offset, txSlice.size);
30 0 : HCCL_INFO("rank[%u] inputSlice range[%llu], size[%llu]", interRank_, txSlice.offset, txSlice.size);
31 0 : CHK_RET(senderInfo_->run(link, baseOffset_ + dstSlice.offset, srcMem, stream_));
32 :
33 0 : return HCCL_SUCCESS;
34 0 : }
35 0 : HcclResult ReduceScatterMesh::RunDestRducer(const LINK& link, const Slice& rxSlice, const Slice& dstSlice)
36 : {
37 : // 使用scratchmem接收数据，并同inputmem数据做reduce
38 0 : DeviceMem dstMem = inputMem_.range(rxSlice.offset, rxSlice.size);
39 0 : DeviceMem srcMem = scratchMem_.range(dstSlice.offset, dstSlice.size);
40 0 : HCCL_INFO(
41 : "rank[%u] rxSlice offset[%llu], size[%llu] dstSlice offset[%llu] size[%llu]", interRank_, rxSlice.offset,
42 : rxSlice.size, dstSlice.offset, dstSlice.size);
43 :
44 : HcclResult ret
45 0 : = reducerInfo_->run(dispatcher_, link, baseOffset_ + rxSlice.offset, dstMem, dstMem, srcMem, stream_);
46 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][DestRducer]rank[%u] reducer info run failed", interRank_), ret);
47 :
48 0 : return HCCL_SUCCESS;
49 0 : }
50 :
51 0 : HcclResult ReduceScatterMesh::RunReduceScatter(
52 : const std::vector<LINK>& links, const std::vector<Slice>& inputSlices, const std::vector<Slice>& scratchSlices)
53 : {
54 0 : std::vector<u32> txRankOpOrder;
55 0 : std::vector<u32> rxRankOpOrder;
56 : // 计算默认的每轮接收的源rank和发送的目的rank
57 0 : for (u32 round = 1; round < interRankSize_; round++) {
58 0 : u32 srcRank = ForwardRank(interRank_, interRankSize_, round);
59 0 : u32 dstRank = BackwardRank(interRank_, interRankSize_, round);
60 0 : HCCL_INFO("<multiDie>RunReduceScatter:srcRank[%u] dstRank[%u]", srcRank, dstRank);
61 0 : rxRankOpOrder.push_back(srcRank);
62 0 : txRankOpOrder.push_back(dstRank);
63 : }
64 :
65 0 : HcclResult ret = HCCL_SUCCESS;
66 0 : for (u32 round = 1; round < interRankSize_; round++) {
67 : // 不同的stream依次轮训默认的顺序数组
68 0 : u32 orderIndex = (round + streamIndex_ - 1) % (interRankSize_ - 1);
69 0 : u32 srcRank = rxRankOpOrder[orderIndex];
70 0 : s32 dstRank = txRankOpOrder[orderIndex];
71 0 : CHK_SMART_PTR_NULL(links[srcRank]);
72 0 : HCCL_INFO("rank[%u] will tx_ack to rank[%u]", interRank_, srcRank);
73 0 : ret = links[srcRank]->TxAck(stream_);
74 0 : CHK_PRT_RET(
75 : ret != HCCL_SUCCESS,
76 : HCCL_ERROR("[Run][ReduceScatter]rank[%u] tx ack to rank[%u] failed", interRank_, srcRank), ret);
77 0 : CHK_SMART_PTR_NULL(links[dstRank]);
78 0 : HCCL_INFO("rank[%u] will rx_ack from rank[%d]", interRank_, dstRank);
79 :
80 0 : ret = links[dstRank]->RxAck(stream_);
81 0 : CHK_PRT_RET(
82 : ret != HCCL_SUCCESS,
83 : HCCL_ERROR("[Run][ReduceScatter]rank[%u] rx ack from rank[%d] failed", interRank_, dstRank), ret);
84 0 : HCCL_INFO(
85 : "rank:%u round[%u] send to rank:[%d], inputSlices offset[%llu]"
86 : "size[%llu] scratchSlice offset[%llu] size[%llu] ",
87 : interRank_, round, dstRank, inputSlices[dstRank].offset, inputSlices[dstRank].size,
88 : scratchSlices[dstRank].offset, scratchSlices[dstRank].size);
89 : // 发送数据
90 0 : ret = RunSourceReducer(links[dstRank], inputSlices[dstRank], scratchSlices[dstRank]);
91 0 : CHK_PRT_RET(
92 : ret != HCCL_SUCCESS,
93 : HCCL_ERROR("[Run][ReduceScatter]rank:%u round[%u] reducer src run failed", interRank_, round), ret);
94 0 : HCCL_INFO(
95 : "rank[%u] round[%u] rx from rank[%u], inSlicesoffset[%llu] size[%llu] "
96 : "scratchSlices offset[%llu] size[%llu]",
97 : interRank_, round, srcRank, inputSlices[interRank_].offset, inputSlices[interRank_].size,
98 : scratchSlices[interRank_].offset, scratchSlices[interRank_].size);
99 :
100 0 : ret = RunDestRducer(links[srcRank], inputSlices[interRank_], scratchSlices[interRank_]);
101 0 : CHK_PRT_RET(
102 : ret != HCCL_SUCCESS,
103 : HCCL_ERROR("[Run][ReduceScatter]rank[%u] round[%u] reducer dst run failed", interRank_, round), ret);
104 :
105 0 : ret = links[srcRank]->RxWaitDone(stream_);
106 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]RxWaitDone failed"), ret);
107 0 : ret = links[dstRank]->TxWaitDone(stream_);
108 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]TxWaitDone failed"), ret);
109 : }
110 0 : if (barrierSwitchOn_) {
111 0 : for (u32 round = 1; round < interRankSize_; round++) {
112 0 : u32 orderIndex = (round + streamIndex_ - 1) % (interRankSize_ - 1);
113 0 : u32 srcRank = rxRankOpOrder[orderIndex];
114 0 : s32 dstRank = txRankOpOrder[orderIndex];
115 :
116 0 : ret = ExecuteBarrier(links[srcRank], links[dstRank]);
117 0 : CHK_PRT_RET(
118 : ret != HCCL_SUCCESS,
119 : HCCL_ERROR(
120 : "[Run][ReduceScatter]rank[%u] run ReduceScatter executor barrier "
121 : "failed. srcRank:%u dstRank:%d",
122 : interRank_, srcRank, dstRank),
123 : ret);
124 : }
125 : }
126 :
127 0 : return HCCL_SUCCESS;
128 0 : }
129 :
130 : // reducescatter的入口函数
131 0 : HcclResult ReduceScatterMesh::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
132 : {
133 0 : CHK_SMART_PTR_NULL(dispatcher_);
134 0 : CHK_PTR_NULL(stream_.ptr());
135 0 : if (!outputMem_ || !inputMem_) {
136 0 : HCCL_ERROR("[ReduceScatterMesh][RunAsync]rank[%u] run_async inputmem or outputmem is null", rank);
137 0 : return HCCL_E_PTR;
138 : }
139 0 : HCCL_INFO(
140 : "ReduceScatterMesh run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
141 : inputMem_.ptr(), outputMem_.ptr(), count_);
142 :
143 0 : interRank_ = rank;
144 0 : interRankSize_ = rankSize;
145 :
146 : // 创建reducer & sender
147 0 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
148 0 : CHK_SMART_PTR_NULL(senderInfo_);
149 :
150 0 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
151 0 : CHK_SMART_PTR_NULL(reducerInfo_);
152 0 : if (interRankSize_ == 1) {
153 0 : if (inputMem_ != outputMem_) {
154 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
155 : }
156 0 : return HCCL_SUCCESS;
157 : }
158 :
159 0 : if (links.size() < rankSize) {
160 0 : HCCL_ERROR("[ReduceScatterMesh][RunAsync]rank[%u] linksize error", rank);
161 0 : return HCCL_E_INTERNAL;
162 : }
163 :
164 0 : if (streamIndex_ >= interRankSize_ - 1) {
165 0 : HCCL_ERROR(
166 : "[ReduceScatterMesh][RunAsync]rank[%u] stream index[%u] is out of range when ranksize[%u]", rank,
167 : streamIndex_, rankSize);
168 0 : return HCCL_E_INTERNAL;
169 : }
170 :
171 0 : u32 unitSize = DataUnitSize(dataType_);
172 0 : if (unitSize == 0) {
173 0 : HCCL_ERROR("[ReduceScatterMesh][RunAsync]rank[%u] unit data size is zero", rank);
174 0 : return HCCL_E_INTERNAL;
175 : }
176 :
177 0 : std::vector<Slice> scratchSlices(slices_);
178 0 : if (slices_.size() == 0) {
179 0 : slices_.resize(rankSize);
180 0 : scratchSlices.resize(rankSize);
181 :
182 : // 生成std::vector<Slice> slices_
183 0 : u64 sliceSize = count_ * unitSize;
184 :
185 0 : for (u32 i = 0; i < rankSize; i++) {
186 0 : slices_[i].size = sliceSize;
187 0 : slices_[i].offset = (i * sliceSize);
188 0 : scratchSlices[i].size = sliceSize;
189 0 : scratchSlices[i].offset = (inputMem_.size() > outputMem_.size()) ? 0 : (i * sliceSize);
190 0 : HCCL_DEBUG(
191 : "rank[%u], slices[%u].offset=[%llu] slices[%u].size=[%llu]", rank, i, slices_[i].offset, i,
192 : slices_[i].size);
193 : }
194 : }
195 : // 运行reduce-scatter, mesh算法
196 0 : CHK_RET(RunReduceScatter(links, slices_, scratchSlices));
197 :
198 0 : if (inputMem_ != outputMem_) {
199 0 : DeviceMem src = inputMem_.range(slices_[interRank_].offset, slices_[interRank_].size);
200 0 : HCCL_DEBUG(
201 : "rank[%u] copy result from to output[%p] offset[%llu] size[%llu] ", interRank_, outputMem_.ptr(),
202 : slices_[interRank_].offset, slices_[interRank_].size);
203 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, src, stream_);
204 0 : CHK_PRT_RET(
205 : ret != HCCL_SUCCESS,
206 : HCCL_ERROR(
207 : "[ReduceScatterMesh][RunAsync]rank[%u] memcpy async from mem[%p] "
208 : "to ouputmem[%p] failed",
209 : rank, src.ptr(), outputMem_.ptr()),
210 : ret);
211 0 : }
212 :
213 0 : HCCL_INFO("ReduceScatterMesh finished: rank[%u]", rank);
214 0 : return HCCL_SUCCESS;
215 0 : }
216 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_MESH, ReduceScatterMesh);
217 : } // namespace hccl
|
