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 "gather_star.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 : // Gather的入口函数
16 0 : GatherStar::GatherStar(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
17 :
18 0 : GatherStar::~GatherStar() {}
19 :
20 0 : HcclResult GatherStar::Prepare(u32 userRank)
21 : {
22 0 : userRank_ = userRank;
23 0 : return HCCL_SUCCESS;
24 : }
25 :
26 : HcclResult
27 0 : GatherStar::RunAsync(const u32 rank, const u32 rankSize, const std::vector<std::shared_ptr<Transport>>& links)
28 : {
29 0 : HCCL_INFO("GatherStar rank[%u] root[%u] linksize[%u]", rank, root_, links.size());
30 : // task下发接口
31 0 : CHK_SMART_PTR_NULL(dispatcher_);
32 : // ==1的处理
33 0 : if (rankSize == 1) {
34 0 : if (inputMem_ != outputMem_) {
35 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
36 0 : CHK_PRT_RET(
37 : ret != HCCL_SUCCESS,
38 : HCCL_ERROR(
39 : "[GatherStar][RunAsync]rank[%u] copy input[%p] to output[%p] "
40 : "failed",
41 : rank, inputMem_.ptr(), outputMem_.ptr()),
42 : ret);
43 : }
44 0 : return HCCL_SUCCESS;
45 : }
46 : // links本rank_id与通信域内其它rank的通信连接,rankSize本executor所在通信域的rank个数
47 0 : if (links.size() < rankSize) {
48 0 : HCCL_ERROR(
49 : "[GatherStar][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(), rankSize);
50 0 : return HCCL_E_INTERNAL;
51 : }
52 : // 计算offset
53 0 : u32 unitSize = DataUnitSize(dataType_);
54 0 : if (unitSize == 0) {
55 0 : HCCL_ERROR("[GatherStar][RunAsync]rank[%u] unit data size is zero", rank);
56 0 : return HCCL_E_INTERNAL;
57 : }
58 0 : if (slices_.size() == 0) {
59 0 : PrepareSlicesData(unitSize, count_, rankSize);
60 : }
61 :
62 0 : Slice sendSlice;
63 0 : sendSlice.offset = dataBytes_ * rank;
64 0 : sendSlice.size = dataBytes_;
65 0 : Slice recvSlice;
66 0 : recvSlice.offset = 0;
67 0 : recvSlice.size = dataBytes_;
68 0 : if (rank == root_) {
69 : // root 从对端其他rank收
70 0 : for (u32 srcRank = 0; srcRank < rankSize; srcRank++) {
71 0 : HcclResult ret = RunRecvGather(srcRank, recvSlice, links);
72 0 : CHK_PRT_RET(
73 : ret != HCCL_SUCCESS,
74 : HCCL_ERROR(
75 : "[GatherStar][RunAsync] srcrank[%u] recv broadcast from"
76 : "root rank[%u] run failed!",
77 : srcRank, root_),
78 : ret);
79 : }
80 : } else {
81 : // 非root 给root发
82 0 : CHK_RET(RunSendGather(root_, sendSlice, links));
83 : }
84 :
85 0 : return HCCL_SUCCESS;
86 : }
87 :
88 0 : void GatherStar::PrepareSlicesData(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
89 : {
90 0 : slices_.resize(rankSize);
91 0 : u64 sliceSize = totalCount * unitSize;
92 :
93 0 : for (u32 i = 0; i < rankSize; i++) {
94 0 : slices_[i].offset = i * sliceSize;
95 0 : slices_[i].size = sliceSize;
96 0 : HCCL_DEBUG("default slice[%u]: offset: [%llu] size[%llu]", i, i * sliceSize, sliceSize);
97 : }
98 0 : }
99 :
100 0 : HcclResult GatherStar::RunRecvGather(const u32 srcRank, const Slice& slice, const std::vector<LINK>& links)
101 : {
102 : // root 接受数据
103 0 : DeviceMem dst;
104 0 : if (slice.size > 0 && srcRank != root_) {
105 0 : if (srcRank >= links.size()) {
106 0 : HCCL_ERROR(
107 : "[Run][GatherStar][RecvGather]SrcRank[%u] is out of range, linkSize[%llu]", srcRank, links.size());
108 0 : return HCCL_E_INTERNAL;
109 : }
110 :
111 0 : dst = outputMem_.range(slice.size * srcRank, slice.size);
112 0 : HCCL_DEBUG(
113 : "rank[%u] will rcv with output's offset[%llu], size[%llu] dstmem[%p]", root_, slice.offset, slice.size,
114 : dst.ptr());
115 :
116 0 : if (links[srcRank]->IsTransportRoce()) {
117 0 : CHK_RET(links[srcRank]->TxEnv(dst.ptr(), slice.size, stream_));
118 : } else {
119 0 : CHK_RET(links[srcRank]->TxAck(stream_));
120 : }
121 :
122 0 : HcclResult ret = links[srcRank]->RxAsync(UserMemType::OUTPUT_MEM, slice.offset, dst.ptr(), slice.size, stream_);
123 0 : CHK_PRT_RET(
124 : ret != HCCL_SUCCESS,
125 : HCCL_ERROR(
126 : "[Run][GatherStar][RecvGather]rank[%u] rx async with output's offset[%llu] failed", root_,
127 : slice.offset),
128 : ret);
129 :
130 0 : if (!links[srcRank]->IsTransportRoce()) {
131 0 : ret = ExecuteBarrier(links[srcRank], stream_);
132 0 : CHK_PRT_RET(
133 : ret != HCCL_SUCCESS,
134 : HCCL_ERROR("[Run][GatherStar][SendGather]srcRank[%u] gather mesh run executor barrier failed", srcRank),
135 : ret);
136 :
137 0 : ret = links[srcRank]->RxWaitDone(stream_);
138 0 : CHK_PRT_RET(
139 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][GatherStar][GatherStar][SendGather]RxWaitDone failed"), ret);
140 : }
141 0 : } else if (srcRank == root_) {
142 0 : CHK_RET(HcclMemcpyAsync(
143 : dispatcher_, static_cast<u8*>(outputMem_.ptr()) + slice.offset, outputMem_.size() - slice.offset,
144 : inputMem_.ptr(), slice.size, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, stream_,
145 : INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP));
146 : }
147 :
148 0 : return HCCL_SUCCESS;
149 0 : }
150 :
151 0 : HcclResult GatherStar::RunSendGather(const u32 dstRank, const Slice& slice, const std::vector<LINK>& links)
152 : {
153 : // 非root发送数据
154 0 : if (slice.size > 0) {
155 0 : HCCL_DEBUG(
156 : "root rank[%u] tx input[%p] offset[%llu] to srcrank size[%llu] ", dstRank, inputMem_.ptr(), slice.offset,
157 : slice.size);
158 0 : if (dstRank >= links.size()) {
159 0 : HCCL_ERROR(
160 : "[Run][GatherStar][SendGather]DstRank[%u] is out of range, link Size[%llu]", dstRank, links.size());
161 0 : return HCCL_E_INTERNAL;
162 : }
163 :
164 0 : if (links[dstRank]->IsTransportRoce()) {
165 0 : CHK_RET(links[dstRank]->RxEnv(stream_));
166 : } else {
167 0 : CHK_RET(links[dstRank]->RxAck(stream_));
168 : }
169 :
170 : HcclResult ret
171 0 : = links[dstRank]->TxAsync(UserMemType::OUTPUT_MEM, slice.offset, inputMem_.ptr(), slice.size, stream_);
172 0 : CHK_PRT_RET(
173 : ret != HCCL_SUCCESS,
174 : HCCL_ERROR(
175 : "[Run][GatherStar][SendGather]srcrank[%u] tx async to root"
176 : "rank[%u] run failed",
177 : dstRank, root_),
178 : ret);
179 :
180 0 : if (!links[dstRank]->IsTransportRoce()) {
181 0 : ret = ExecuteBarrierSrcRank(links[dstRank], stream_);
182 0 : CHK_PRT_RET(
183 : ret != HCCL_SUCCESS,
184 : HCCL_ERROR(
185 : "[Run][GatherStar][SendGather]dstRank[%u] gather mesh run"
186 : "executor barrier failed",
187 : dstRank),
188 : ret);
189 0 : ret = links[dstRank]->TxWaitDone(stream_);
190 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][GatherStar][SendGather]TxWaitDone failed"), ret);
191 : }
192 : }
193 0 : return HCCL_SUCCESS;
194 : }
195 :
196 0 : HcclResult GatherStar::ExecuteBarrierSrcRank(std::shared_ptr<Transport> link, Stream& stream) const
197 : {
198 0 : CHK_RET(link->RxAck(stream));
199 :
200 0 : CHK_RET(link->TxAck(stream));
201 :
202 0 : CHK_RET(link->RxDataSignal(stream));
203 :
204 0 : CHK_RET(link->TxDataSignal(stream));
205 :
206 0 : return HCCL_SUCCESS;
207 : }
208 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_GATHER_STAR, GatherStar);
209 : } // namespace hccl
|
