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 "scatter_ring_direct.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : ScatterRingDirect::ScatterRingDirect(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : ScatterRingDirect::~ScatterRingDirect() {}
18 :
19 0 : HcclResult ScatterRingDirect::Prepare(
20 : HcomCollOpInfo* opInfo, const u32 userRank, const std::vector<u32>& ringsOrders,
21 : const std::vector<Slice>& userMemInputSlices)
22 : {
23 0 : opInfo_ = opInfo;
24 0 : userRank_ = userRank;
25 0 : ringsOrder_ = ringsOrders;
26 0 : userMemInputSlices_ = userMemInputSlices;
27 0 : return HCCL_SUCCESS;
28 : }
29 :
30 : // reduce scatter ring direct算法的函数入口
31 0 : HcclResult ScatterRingDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
32 : {
33 : // 基本的检查
34 0 : CHK_RET(CheckParameters(rank, rankSize, links));
35 :
36 : // 判断rank_size == 1, 若inputMem_ != outputMem_，才需要搬运
37 0 : if (rankSize == 1) {
38 0 : CHK_RET(OneRankMemcpy());
39 0 : return HCCL_SUCCESS;
40 : }
41 : // 收集邻居信息
42 0 : CHK_RET(GetInitializedNeighborLinks(rank, rankSize, links));
43 : // 填充slice_
44 0 : CHK_RET(SetSlices(rank, rankSize));
45 :
46 : // 运行scatter, ring算法
47 0 : CHK_RET(RunScatter(rank, rankSize));
48 :
49 0 : if (barrierSwitchOn_) {
50 : // 执行barrier，保证数据收发完成
51 0 : CHK_RET(ExecuteBarrier(leftLink_, rightLink_));
52 : }
53 :
54 0 : HCCL_INFO("ScatterRingDirect finished: rank[%u]", rank);
55 0 : return HCCL_SUCCESS;
56 : }
57 :
58 0 : HcclResult ScatterRingDirect::CheckParameters(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
59 : {
60 0 : CHK_PTR_NULL(opInfo_);
61 0 : CHK_RET(CheckConcurrentDirectParameters(rank, rankSize, links));
62 : // 判断ringsOrder数量是否正确
63 0 : CHK_PRT_RET(
64 : ringsOrder_.size() != rankSize,
65 : HCCL_ERROR(
66 : "[ScatterRingDirect] ringsOrder size[%u] is not equal to rank size[%u]", ringsOrder_.size(), rankSize),
67 : HCCL_E_PARA);
68 : // 判断userMemInputSlices数量是否正确
69 0 : CHK_PRT_RET(
70 : userMemInputSlices_.size() != rankSize,
71 : HCCL_ERROR(
72 : "[ScatterRingDirect] userMemInputSlices size[%u] is not equal to rank size[%u]", userMemInputSlices_.size(),
73 : rankSize),
74 : HCCL_E_PARA);
75 0 : HCCL_INFO("ScatterRingDirect CheckParameters success");
76 0 : return HCCL_SUCCESS;
77 : }
78 :
79 0 : HcclResult ScatterRingDirect::OneRankMemcpy()
80 : {
81 0 : const Slice& srcSlice = userMemInputSlices_[0];
82 0 : const Slice& dstSlice = slices_[0];
83 0 : DeviceMem src;
84 0 : DeviceMem dst;
85 0 : if (opInfo_->inputAddr == nullptr) {
86 0 : src = inputMem_.range(srcSlice.offset, srcSlice.size);
87 : } else {
88 0 : src = DeviceMem::create(static_cast<u8*>(opInfo_->inputAddr) + srcSlice.offset, srcSlice.size);
89 : }
90 0 : if (opInfo_->outputAddr != nullptr) {
91 : // opInfo_->outputAddr != nullptr指示要将输出发送至user output
92 0 : u64 stepOffset = slices_[ringsOrder_[0]].offset;
93 0 : HCCL_DEBUG(
94 : "Memcpy operation: stream[main], rank[%u] starts to rcv offset[%llu], size[%llu] at userMemOut_", userRank_,
95 : stepOffset, dstSlice.size);
96 0 : dst = DeviceMem::create(static_cast<u8*>(opInfo_->outputAddr) + stepOffset, dstSlice.size);
97 : } else {
98 : // opInfo_->outputAddr == nullptr指示要将输出发送至CCL buffer
99 0 : HCCL_DEBUG(
100 : "Memcpy operation: stream[main], rank[%u] starts to rcv offset[%llu], size[%llu] at outputMem_", userRank_,
101 : dstSlice.offset, dstSlice.size);
102 0 : dst = outputMem_.range(dstSlice.offset, dstSlice.size);
103 : }
104 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
105 0 : return HCCL_SUCCESS;
106 0 : }
107 :
108 : HcclResult
109 0 : ScatterRingDirect::GetInitializedNeighborLinks(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
110 : {
111 : // 收集左邻居信息
112 0 : leftLink_ = links[(rank + rankSize - 1) % rankSize];
113 0 : CHK_SMART_PTR_NULL(leftLink_);
114 :
115 : // 收集右邻居信息
116 0 : rightLink_ = links[(rank + 1) % rankSize];
117 0 : CHK_SMART_PTR_NULL(rightLink_);
118 0 : HCCL_INFO("ScatterRingDirect finished to GetInitializedNeighborLinks");
119 0 : return HCCL_SUCCESS;
120 : }
121 :
122 0 : HcclResult ScatterRingDirect::SetSlices(const u32 rank, const u32 rankSize)
123 : {
124 0 : if (slices_.size() == 0) {
125 0 : slices_.resize(rankSize);
126 :
127 : // 生成std::vector<Slice> slices_
128 0 : u64 sliceSize = count_ * SIZE_TABLE[dataType_];
129 : ;
130 :
131 0 : for (u32 i = 0; i < rankSize; i++) {
132 0 : slices_[i].size = sliceSize;
133 : // 用于DMA消减过程中，消除src与dst不对位的风险
134 0 : slices_[i].offset = RoundUpWithDivisor(i * sliceSize, HCCL_MIN_SLICE_ALIGN);
135 :
136 0 : HCCL_DEBUG(
137 : "rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu]", rank, i, slices_[i].offset, i,
138 : slices_[i].size);
139 : }
140 : }
141 0 : if (UNLIKELY(HcclCheckLogLevel(HCCL_LOG_DEBUG))) {
142 0 : for (u32 i = 0; i < slices_.size(); i++) {
143 0 : HCCL_DEBUG(
144 : "[ScatterRingDirect][SetSlices]rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu]", rank, i,
145 : slices_[i].offset, i, slices_[i].size);
146 : }
147 : }
148 : // 最后一步搬到userMemOut_的offset, 不同的ring环offset不一样
149 0 : lastStepOffset_ = slices_[ringsOrder_[0]].offset;
150 0 : HCCL_INFO("ScatterRingDirect finished to SetSlices");
151 0 : return HCCL_SUCCESS;
152 : }
153 :
154 0 : HcclResult ScatterRingDirect::RunScatter(const u32 rank, const u32 rankSize)
155 : {
156 0 : HCCL_INFO("ScatterRingDirect starts, the input param rank[%u]", rank);
157 : // 空拷贝用于后续操作附着
158 0 : CHK_RET(ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
159 : // 例如rank[0,1,2,3]中，rank0的rxSliceIdx = 2，txSliceIdx = 3, subSliceIdx = 1
160 0 : u32 txSliceIdx = (rank + rankSize - 1) % rankSize;
161 0 : u32 rxSliceIdx = (rank + rankSize - DMA_REDUCE_TWO_OFFSET) % rankSize;
162 0 : u32 subSliceIdx = (rank + rankSize - DMA_REDUCE_TWO_OFFSET) % rankSize; // 只存在于根节点
163 0 : u32 stepsFromRank2Root = (rank + rankSize - root_) % rankSize;
164 0 : for (u32 step = 0; step < rankSize - 1; step++) {
165 0 : const Slice& subSlice = userMemInputSlices_[subSliceIdx];
166 0 : const Slice& cclSlice = slices_[subSliceIdx];
167 0 : const Slice& txSlice = slices_[txSliceIdx];
168 0 : const Slice& rxSlice = slices_[rxSliceIdx];
169 :
170 0 : CHK_RET(RunScatterOnRootRank(step, subSlice, cclSlice, rank, rankSize));
171 0 : CHK_RET(RunScatterOnOtherRank(stepsFromRank2Root, step, txSlice, rxSlice, rankSize));
172 :
173 : // 更新索引
174 0 : subSliceIdx = (subSliceIdx + rankSize - 1) % rankSize;
175 0 : txSliceIdx = (txSliceIdx + rankSize - 1) % rankSize;
176 0 : rxSliceIdx = (rxSliceIdx + rankSize - 1) % rankSize;
177 : }
178 0 : HCCL_INFO("ScatterRingDirect finished to RunScatter");
179 0 : return HCCL_SUCCESS;
180 : }
181 :
182 0 : HcclResult ScatterRingDirect::RunScatterOnOtherRank(
183 : const u32 stepsFromRank2Root, const u32 step, const Slice& txSlice, const Slice& rxSlice, const u32 rankSize)
184 : {
185 0 : bool needSend = stepsFromRank2Root <= step;
186 0 : bool needReceive = stepsFromRank2Root > 0 && stepsFromRank2Root <= (step + 1);
187 : // Ack
188 0 : if (needReceive) {
189 0 : CHK_RET(leftLink_->TxAck(stream_));
190 : }
191 0 : if (needSend) {
192 0 : CHK_RET(rightLink_->RxAck(stream_));
193 : }
194 :
195 0 : DeviceMem src;
196 : // 不同的rank会在不同的step开始持续发送操作，距离root节点越近，越早step开始发送操作
197 0 : if (needSend) {
198 0 : src = inputMem_.range(txSlice.offset, txSlice.size);
199 0 : CHK_RET(rightLink_->TxAsync(
200 : UserMemType::INPUT_MEM, txSlice.offset + baseOffset_, src.ptr(), txSlice.size, stream_));
201 : }
202 : // 不同的rank会在不同的step开始持续发送操作，距离root节点越近，越早step开始发送操作
203 0 : DeviceMem dst;
204 0 : if (needReceive) {
205 0 : HCCL_DEBUG(
206 : "MemcpyAsync operation: step[%u] stream[main], src rank[%u] starts to send offset[%llu] size[%llu] "
207 : "from leftMem_",
208 : step, leftLink_->GetRemoteRank(), rxSlice.offset, rxSlice.size);
209 0 : if (step == rankSize - DMA_REDUCE_TWO_OFFSET && opInfo_->outputAddr != nullptr) {
210 0 : HCCL_DEBUG(
211 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu], "
212 : "size[%llu] "
213 : "at userMemOut_",
214 : step, userRank_, lastStepOffset_, rxSlice.size);
215 0 : dst = DeviceMem::create(static_cast<u8*>(opInfo_->outputAddr) + lastStepOffset_, rxSlice.size);
216 : } else {
217 0 : HCCL_DEBUG(
218 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu], "
219 : "size[%llu] "
220 : "at inputMem_",
221 : step, userRank_, rxSlice.offset, rxSlice.size);
222 0 : dst = inputMem_.range(rxSlice.offset, rxSlice.size);
223 : }
224 0 : CHK_RET(
225 : leftLink_->RxAsync(UserMemType::INPUT_MEM, rxSlice.offset + baseOffset_, dst.ptr(), rxSlice.size, stream_));
226 : }
227 0 : return HCCL_SUCCESS;
228 0 : }
229 :
230 0 : HcclResult ScatterRingDirect::RunScatterOnRootRank(
231 : const u32 step, const Slice& subSlice, const Slice& cclSlice, const u32 rank, const u32 rankSize)
232 : {
233 0 : if (step == rankSize - DMA_REDUCE_TWO_OFFSET && opInfo_->outputAddr != nullptr && rank == root_) {
234 0 : HCCL_DEBUG(
235 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu], "
236 : "size[%llu] at userMemOut_",
237 : step, userRank_, lastStepOffset_, subSlice.size);
238 0 : DeviceMem src = inputMem_.range(cclSlice.offset, cclSlice.size);
239 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(opInfo_->outputAddr) + lastStepOffset_, subSlice.size);
240 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
241 0 : }
242 0 : return HCCL_SUCCESS;
243 : }
244 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_SCATTER_RING_DIRECT, ScatterRingDirect);
245 : } // namespace hccl
|
