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)
17 0 : : AlgTemplateBase(dispatcher)
18 : {
19 0 : }
20 :
21 0 : GatherStar::~GatherStar()
22 : {
23 0 : }
24 :
25 0 : HcclResult GatherStar::Prepare(u32 userRank)
26 : {
27 0 : userRank_ = userRank;
28 0 : return HCCL_SUCCESS;
29 : }
30 :
31 0 : HcclResult GatherStar::RunAsync(const u32 rank, const u32 rankSize,
32 : const std::vector<std::shared_ptr<Transport>> &links)
33 : {
34 0 : HCCL_INFO("GatherStar rank[%u] root[%u] linksize[%u]", rank, root_, links.size());
35 : // task下发接口
36 0 : CHK_SMART_PTR_NULL(dispatcher_);
37 : // ==1的处理
38 0 : if (rankSize == 1) {
39 0 : if (inputMem_ != outputMem_) {
40 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
41 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[GatherStar][RunAsync]rank[%u] copy input[%p] to output[%p] "\
42 : "failed", rank, inputMem_.ptr(), outputMem_.ptr()), 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("[GatherStar][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]",
49 : 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(ret != HCCL_SUCCESS, HCCL_ERROR("[GatherStar][RunAsync] srcrank[%u] recv broadcast from"\
73 : "root rank[%u] run failed!", srcRank, root_), ret);
74 : }
75 : } else {
76 : // 非root 给root发
77 0 : CHK_RET(RunSendGather(root_, sendSlice, links));
78 : }
79 :
80 0 : return HCCL_SUCCESS;
81 : }
82 :
83 0 : void GatherStar::PrepareSlicesData(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
84 : {
85 0 : slices_.resize(rankSize);
86 0 : u64 sliceSize = totalCount * unitSize;
87 :
88 0 : for (u32 i = 0; i < rankSize; i++) {
89 0 : slices_[i].offset = i * sliceSize;
90 0 : slices_[i].size = sliceSize;
91 0 : HCCL_DEBUG("default slice[%u]: offset: [%llu] size[%llu]", i, i * sliceSize, sliceSize);
92 : }
93 0 : }
94 :
95 0 : HcclResult GatherStar::RunRecvGather(const u32 srcRank, const Slice &slice, const std::vector<LINK> &links)
96 : {
97 : // root 接受数据
98 0 : DeviceMem dst;
99 0 : if (slice.size > 0 && srcRank != root_) {
100 0 : if (srcRank >= links.size()) {
101 0 : HCCL_ERROR("[Run][GatherStar][RecvGather]SrcRank[%u] is out of range, linkSize[%llu]", \
102 : srcRank, links.size());
103 0 : return HCCL_E_INTERNAL;
104 : }
105 :
106 0 : dst = outputMem_.range(slice.size * srcRank, slice.size);
107 0 : HCCL_DEBUG("rank[%u] will rcv with output's offset[%llu], size[%llu] dstmem[%p]", \
108 : root_, slice.offset, slice.size, dst.ptr());
109 :
110 0 : if (links[srcRank]->IsTransportRoce()) {
111 0 : CHK_RET(links[srcRank]->TxEnv(dst.ptr(), slice.size, stream_));
112 : } else {
113 0 : CHK_RET(links[srcRank]->TxAck(stream_));
114 : }
115 :
116 0 : HcclResult ret = links[srcRank]->RxAsync(UserMemType::OUTPUT_MEM, slice.offset, dst.ptr(), slice.size, stream_);
117 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
118 : HCCL_ERROR("[Run][GatherStar][RecvGather]rank[%u] rx async with output's offset[%llu] failed", \
119 : root_, slice.offset), ret);
120 :
121 0 : if (!links[srcRank]->IsTransportRoce()) {
122 0 : ret = ExecuteBarrier(links[srcRank], stream_);
123 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
124 : HCCL_ERROR("[Run][GatherStar][SendGather]srcRank[%u] gather mesh run executor barrier failed", \
125 : srcRank), ret);
126 :
127 0 : ret = links[srcRank]->RxWaitDone(stream_);
128 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
129 : HCCL_ERROR("[Run][GatherStar][GatherStar][SendGather]RxWaitDone failed"), ret);
130 : }
131 0 : } else if (srcRank == root_) {
132 0 : CHK_RET(HcclMemcpyAsync(dispatcher_, static_cast<u8 *>(outputMem_.ptr()) + slice.offset,
133 : outputMem_.size() - slice.offset, inputMem_.ptr(), slice.size,
134 : HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, stream_,
135 : INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP));
136 : }
137 :
138 0 : return HCCL_SUCCESS;
139 0 : }
140 :
141 0 : HcclResult GatherStar::RunSendGather(const u32 dstRank, const Slice &slice,
142 : const std::vector<LINK> &links)
143 : {
144 : // 非root发送数据
145 0 : if (slice.size > 0) {
146 0 : HCCL_DEBUG("root rank[%u] tx input[%p] offset[%llu] to srcrank size[%llu] ", \
147 : dstRank, inputMem_.ptr(), slice.offset, slice.size);
148 0 : if (dstRank >= links.size()) {
149 0 : HCCL_ERROR("[Run][GatherStar][SendGather]DstRank[%u] is out of range, link Size[%llu]", \
150 : dstRank, links.size());
151 0 : return HCCL_E_INTERNAL;
152 : }
153 :
154 0 : if (links[dstRank]->IsTransportRoce()) {
155 0 : CHK_RET(links[dstRank]->RxEnv(stream_));
156 : } else {
157 0 : CHK_RET(links[dstRank]->RxAck(stream_));
158 : }
159 :
160 0 : HcclResult ret = links[dstRank]->TxAsync(UserMemType::OUTPUT_MEM, slice.offset, inputMem_.ptr(),
161 0 : slice.size, stream_);
162 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][GatherStar][SendGather]srcrank[%u] tx async to root" \
163 : "rank[%u] run failed", dstRank, root_), ret);
164 :
165 0 : if (!links[dstRank]->IsTransportRoce()) {
166 0 : ret = ExecuteBarrierSrcRank(links[dstRank], stream_);
167 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][GatherStar][SendGather]dstRank[%u] gather mesh run" \
168 : "executor barrier failed", dstRank), ret);
169 0 : ret = links[dstRank]->TxWaitDone(stream_);
170 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][GatherStar][SendGather]TxWaitDone failed"), ret);
171 : }
172 : }
173 0 : return HCCL_SUCCESS;
174 : }
175 :
176 0 : HcclResult GatherStar::ExecuteBarrierSrcRank(std::shared_ptr<Transport> link, Stream &stream) const
177 : {
178 0 : CHK_RET(link->RxAck(stream));
179 :
180 0 : CHK_RET(link->TxAck(stream));
181 :
182 0 : CHK_RET(link->RxDataSignal(stream));
183 :
184 0 : CHK_RET(link->TxDataSignal(stream));
185 :
186 0 : return HCCL_SUCCESS;
187 : }
188 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_GATHER_STAR, GatherStar);
189 : }
|
