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