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_mesh.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : GatherMesh::GatherMesh(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher), round_(0) {}
16 :
17 0 : GatherMesh::~GatherMesh() {}
18 :
19 0 : HcclResult GatherMesh::Prepare(
20 : std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
21 : std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 userRank, [[maybe_unused]] HcomCollOpInfo* opInfo,
22 : [[maybe_unused]] u32 interRank, [[maybe_unused]] u32 interRankSize)
23 : {
24 0 : meshStreams_ = meshStreams;
25 0 : meshSignal_ = &meshSignal;
26 0 : meshSignalAux_ = &meshSignalAux;
27 0 : userRank_ = userRank;
28 0 : return HCCL_SUCCESS;
29 : }
30 :
31 0 : void GatherMesh::PrepareSlicesData(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
32 : {
33 0 : slices_.resize(rankSize);
34 0 : u64 sliceSize = (totalCount / rankSize) * unitSize;
35 :
36 0 : for (u32 i = 0; i < rankSize; i++) {
37 0 : slices_[i].offset = i * sliceSize;
38 0 : slices_[i].size = sliceSize;
39 0 : HCCL_DEBUG("default slice[%u]: offset: [%llu] size[%llu]", i, i * sliceSize, sliceSize);
40 : }
41 0 : }
42 :
43 0 : HcclResult GatherMesh::ExecuteBarrierSrcRank(std::shared_ptr<Transport> link, Stream& stream) const
44 : {
45 0 : CHK_RET(link->RxAck(stream));
46 :
47 0 : CHK_RET(link->TxAck(stream));
48 :
49 0 : CHK_RET(link->RxDataSignal(stream));
50 :
51 0 : CHK_RET(link->TxDataSignal(stream));
52 :
53 0 : return HCCL_SUCCESS;
54 : }
55 :
56 : // root rank接收数据
57 0 : HcclResult GatherMesh::RunRecvGather(const u32 srcRank, const Slice& slice, const std::vector<LINK>& links)
58 : {
59 0 : DeviceMem src;
60 0 : DeviceMem dst;
61 0 : if (srcRank == root_) {
62 0 : if (inputMem_ != outputMem_) {
63 : // root rank给自身拷贝时候不需要同步信号，拷贝到outputmem的偏移不同
64 0 : Slice& rootSlice = slices_[root_];
65 0 : src = inputMem_.range(rootSlice.offset, rootSlice.size);
66 0 : dst = outputMem_.range(rootSlice.offset, rootSlice.size);
67 0 : HCCL_DEBUG(
68 : "root rank copy from input[%p] range[%llu] to output[%p] range[%llu], size[%llu]", src.ptr(),
69 : rootSlice.offset, dst.ptr(), rootSlice.offset, rootSlice.size);
70 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
71 : }
72 : } else {
73 : // 判断数据是否需要分片
74 0 : if (srcRank >= links.size()) {
75 0 : HCCL_ERROR("[Run][RecvGather]SrcRank[%u] is out of range, linkSize[%llu]", srcRank, links.size());
76 0 : return HCCL_E_INTERNAL;
77 : }
78 0 : const LINK& link = links[srcRank];
79 0 : dst = outputMem_.range(slice.offset, slice.size);
80 0 : HCCL_DEBUG(
81 : "rank[%u] will rcv with output's offset[%llu], size[%llu] dstmem[%p]", root_, slice.offset, slice.size,
82 : dst.ptr());
83 :
84 : // 向非root节点发送tx同步,rxmem可用
85 0 : Stream& curStream = (round_ == 0) ? stream_ : meshStreams_[round_ - 1];
86 0 : HcclResult ret = link->TxAck(curStream);
87 0 : CHK_PRT_RET(
88 : ret != HCCL_SUCCESS,
89 : HCCL_ERROR("[Run][RecvGather]root rank[%u] tx ack to srcrank[%u] failed", root_, srcRank), ret);
90 :
91 0 : ret = link->RxAsync(UserMemType::OUTPUT_MEM, slice.offset, dst.ptr(), slice.size, curStream);
92 0 : CHK_PRT_RET(
93 : ret != HCCL_SUCCESS,
94 : HCCL_ERROR("[Run][RecvGather]rank[%u] rx async with output's offset[%llu] failed", root_, slice.offset),
95 : ret);
96 :
97 0 : CHK_RET(link->TxDataSignal(curStream));
98 0 : CHK_RET(link->RxWaitDone(curStream));
99 : }
100 0 : return HCCL_SUCCESS;
101 0 : }
102 :
103 : // 非root rank发送数据
104 0 : HcclResult GatherMesh::RunSendGather(const u32 dstRank, const Slice& slice, const std::vector<LINK>& links)
105 : {
106 0 : DeviceMem src = inputMem_.range(slice.offset, slice.size);
107 0 : if (dstRank >= links.size()) {
108 0 : HCCL_ERROR("[Run][SendGather]DstRank[%u] is out of range, link Size[%llu]", dstRank, links.size());
109 0 : return HCCL_E_INTERNAL;
110 : }
111 0 : const LINK& link = links[dstRank];
112 :
113 0 : HCCL_DEBUG(
114 : "root rank[%u] tx input[%p] offset[%llu] to srcrank size[%llu] ", dstRank, src.ptr(), slice.offset, slice.size);
115 : // 接收目的rank的同步信号，便可进行下一轮发送
116 0 : CHK_RET(link->RxAck(stream_));
117 :
118 0 : HcclResult ret = link->TxAsync(UserMemType::OUTPUT_MEM, slice.offset, src.ptr(), slice.size, stream_);
119 0 : CHK_PRT_RET(
120 : ret != HCCL_SUCCESS,
121 : HCCL_ERROR("[Run][SendGather]srcrank[%u] tx async to root rank[%u] run failed", dstRank, dstRank), ret);
122 :
123 0 : CHK_RET(link->RxDataSignal(stream_));
124 0 : CHK_RET(link->TxWaitDone(stream_));
125 0 : return HCCL_SUCCESS;
126 0 : }
127 :
128 : // Gather的入口函数
129 : HcclResult
130 0 : GatherMesh::RunAsync(const u32 rank, const u32 rankSize, const std::vector<std::shared_ptr<Transport>>& links)
131 : {
132 0 : CHK_SMART_PTR_NULL(dispatcher_);
133 0 : CHK_PTR_NULL(stream_.ptr());
134 0 : HCCL_INFO(
135 : "GatherMesh run: rank[%u] rankSize[%u] inputMem[%p] to outputMem[%p] count[%llu]", rank, rankSize,
136 : inputMem_.ptr(), outputMem_.ptr(), count_);
137 : // ranksize ==1 的处理
138 0 : if (rankSize == 1) {
139 0 : if (inputMem_ != outputMem_) {
140 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
141 0 : CHK_PRT_RET(
142 : ret != HCCL_SUCCESS,
143 : HCCL_ERROR(
144 : "[GatherMesh][RunAsync]rank[%u] copy input[%p] to output[%p] "
145 : "failed",
146 : rank, inputMem_.ptr(), outputMem_.ptr()),
147 : ret);
148 : }
149 0 : return HCCL_SUCCESS;
150 : }
151 :
152 0 : if (links.size() < rankSize) {
153 0 : HCCL_ERROR(
154 : "[GatherMesh][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(), rankSize);
155 0 : return HCCL_E_INTERNAL;
156 : }
157 :
158 0 : u32 unitSize = SIZE_TABLE[dataType_];
159 0 : if (unitSize == 0) {
160 0 : HCCL_ERROR("[GatherMesh][RunAsync]rank[%u] unit data size is zero", rank);
161 0 : return HCCL_E_INTERNAL;
162 : }
163 0 : if (slices_.size() == 0) {
164 0 : PrepareSlicesData(unitSize, count_, rankSize);
165 : }
166 :
167 0 : if (rank == root_) {
168 0 : CHK_RET(AddMainSteamSubStreamSyncPre(rank, rankSize));
169 : // root rank接收其他rank发送的数据
170 0 : round_ = 0;
171 0 : for (u32 srcRank = 0; srcRank < rankSize; srcRank++) {
172 0 : HcclResult ret = RunRecvGather(srcRank, slices_[srcRank], links);
173 0 : CHK_PRT_RET(
174 : ret != HCCL_SUCCESS,
175 : HCCL_ERROR(
176 : "[GatherMesh][RunAsync]srcrank[%u] send gather to "
177 : "root rank[%u] run failed",
178 : srcRank, rank),
179 : ret);
180 0 : if (srcRank != root_) {
181 0 : round_++;
182 : }
183 : }
184 0 : CHK_RET(AddMainSteamSubStreamSyncPost(rank, rankSize));
185 :
186 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
187 : } else {
188 : // 非root rank向root rank发送数据
189 0 : CHK_RET(RunSendGather(root_, slices_[rank], links));
190 : }
191 0 : HCCL_INFO("GatherMesh finished: rank[%u], end", rank);
192 0 : return HCCL_SUCCESS;
193 : }
194 :
195 0 : HcclResult GatherMesh::AddMainSteamSubStreamSyncPre(u32 rank, u32 rankSize)
196 : {
197 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize-2: stream num
198 0 : HCCL_DEBUG(
199 : "rank[%u] streamindex[%u] wait signalaux[%p]", rank, streamIndex, (*meshSignalAux_)[streamIndex]->ptr());
200 0 : CHK_RET(LocalNotify::Wait(
201 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
202 :
203 0 : HCCL_DEBUG(
204 : "rank[%u] siganl_aux index[%u] signal record signalaux[%p] ", rank, streamIndex,
205 : (*meshSignalAux_)[streamIndex]->ptr());
206 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
207 : }
208 0 : return HCCL_SUCCESS;
209 : }
210 :
211 0 : HcclResult GatherMesh::AddMainSteamSubStreamSyncPost(u32 rank, u32 rankSize)
212 : {
213 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize - 2 stream num
214 0 : HCCL_DEBUG("rank[%u] streamindex[%u] wait signal[%p] ", rank, streamIndex, (*meshSignal_)[streamIndex]->ptr());
215 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
216 :
217 0 : HCCL_DEBUG("rank[%u] streamindex[%u] record signal[%p]", rank, streamIndex, (*meshSignal_)[streamIndex]->ptr());
218 0 : CHK_RET(LocalNotify::Post(
219 : meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
220 : }
221 0 : return HCCL_SUCCESS;
222 : }
223 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_GATHER_MESH, GatherMesh);
224 : } // namespace hccl
|
