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