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 "all_gather_mesh.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AllGatherMesh::AllGatherMesh(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : AllGatherMesh::~AllGatherMesh() {}
18 :
19 0 : HcclResult AllGatherMesh::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 : u32 interRank, u32 interRankSize)
23 : {
24 0 : meshStreams_ = meshStreams;
25 0 : meshSignal_ = &meshSignal;
26 0 : meshSignalAux_ = &meshSignalAux;
27 0 : interRank_ = interRank;
28 0 : interRankSize_ = interRankSize;
29 0 : userRank_ = userRank;
30 0 : return HCCL_SUCCESS;
31 : }
32 :
33 0 : HcclResult AllGatherMesh::Tx(const LINK& link, const Slice& txSlice, const Slice& dstSlice, Stream stream)
34 : {
35 0 : DeviceMem srcMem = outputMem_.range(txSlice.offset, txSlice.size);
36 :
37 0 : HCCL_DEBUG(
38 : "rank[%u] tx srcMem[%p] output's offset[%llu] size[%llu] to dstrank's offset[%llu]", interRank_, srcMem.ptr(),
39 : txSlice.offset, txSlice.size, dstSlice.offset);
40 : HcclResult ret
41 0 : = link->TxAsync(UserMemType::OUTPUT_MEM, baseOffset_ + dstSlice.offset, srcMem.ptr(), txSlice.size, stream);
42 0 : CHK_PRT_RET(
43 : ret != HCCL_SUCCESS,
44 : HCCL_ERROR("[AllGatherMesh][Tx]rank[%u] tx srcMem[%p] tx_async Failed", interRank_, srcMem.ptr()), ret);
45 0 : return HCCL_SUCCESS;
46 0 : }
47 :
48 0 : HcclResult AllGatherMesh::Rx(const LINK& link, const Slice& srcSlice, const Slice& rxSlice, Stream stream)
49 : {
50 0 : DeviceMem rcvMem = outputMem_.range(rxSlice.offset, rxSlice.size);
51 :
52 0 : HCCL_DEBUG(
53 : "rank[%u] rx rcvMem[%p] output's offset[%llu] size[%llu] rcv data from srcrank's"
54 : "offset[%llu] ",
55 : interRank_, rcvMem.ptr(), rxSlice.offset, rxSlice.size, srcSlice.offset);
56 : HcclResult ret
57 0 : = link->RxAsync(UserMemType::OUTPUT_MEM, baseOffset_ + srcSlice.offset, rcvMem.ptr(), rxSlice.size, stream);
58 0 : CHK_PRT_RET(
59 : ret != HCCL_SUCCESS,
60 : HCCL_ERROR("[AllGatherMesh][Tx]rank[%u] rcvMem[%p] rx_async Failed", interRank_, rcvMem.ptr()), ret);
61 0 : return HCCL_SUCCESS;
62 0 : }
63 :
64 0 : HcclResult AllGatherMesh::RunAllGather(
65 : const std::vector<LINK>& links, const std::vector<Slice>& outputSlices, const std::vector<Slice>& inputSlices)
66 : {
67 0 : Stream subStream;
68 0 : HcclResult ret = HCCL_SUCCESS;
69 0 : for (u32 round = 1; round < interRankSize_; round++) {
70 0 : u32 dstRank = BackwardRank(interRank_, interRankSize_, round);
71 :
72 0 : subStream = (round == interRankSize_ - 1) ? stream_ : meshStreams_[round - 1];
73 :
74 0 : profilerInput_.streamID = subStream.id();
75 0 : profilerInput_.planeID = round - 1;
76 0 : profilerInput_.step = HCCL_EXEC_STEP_NOT_SET;
77 :
78 0 : CHK_SMART_PTR_NULL(links[dstRank]);
79 0 : HCCL_DEBUG("rank[%u] round[%u] tx ack to srcRank[%u] ", interRank_, round, dstRank);
80 0 : ret = links[dstRank]->TxAck(subStream);
81 0 : CHK_PRT_RET(
82 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]rank[%u] tx ack to rank[%u] Failed", interRank_, dstRank),
83 : ret);
84 0 : CHK_SMART_PTR_NULL(links[dstRank]);
85 :
86 0 : HCCL_DEBUG("rank[%u] round[%u] rx ack from Rank[%u] ", interRank_, round, dstRank);
87 0 : ret = links[dstRank]->RxAck(subStream);
88 0 : CHK_PRT_RET(
89 : ret != HCCL_SUCCESS,
90 : HCCL_ERROR("[Run][AllGather]rank[%u] rx_ack from rank[%u] Failed", interRank_, dstRank), ret);
91 :
92 0 : HCCL_DEBUG(
93 : "rank[%u] tx to rank[%u] inputslice offset[%llu] size[%llu] outputSlices "
94 : "offset[%llu] size[%llu] ",
95 : interRank_, dstRank, inputSlices[interRank_].offset, inputSlices[interRank_].size,
96 : outputSlices[interRank_].offset, outputSlices[interRank_].size);
97 :
98 0 : ret = Tx(links[dstRank], inputSlices[interRank_], outputSlices[interRank_], subStream);
99 0 : CHK_PRT_RET(
100 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]rank[%u] tx to rank[%u] failed", interRank_, dstRank),
101 : ret);
102 0 : HCCL_DEBUG(
103 : "rank[%u] will rcv from rank[%u] inputSlices offset[%llu] size[%llu] outputSlices"
104 : "offset[%llu] size[%llu]",
105 : interRank_, dstRank, inputSlices[dstRank].offset, inputSlices[dstRank].size, outputSlices[dstRank].offset,
106 : outputSlices[dstRank].size);
107 :
108 0 : ret = Rx(links[dstRank], inputSlices[dstRank], outputSlices[dstRank], subStream);
109 0 : CHK_PRT_RET(
110 : ret != HCCL_SUCCESS,
111 : HCCL_ERROR("[Run][AllGather]rank[%u] rx from srcRank[%u] run failed", interRank_, dstRank), ret);
112 0 : ret = ExecuteBarrier(links[dstRank], subStream);
113 0 : CHK_PRT_RET(
114 : ret != HCCL_SUCCESS,
115 : HCCL_ERROR(
116 : "[Run][AllGather]destrank[%u] AllGather mesh run tempAlg barrier "
117 : "failed",
118 : dstRank),
119 : ret);
120 0 : ret = links[dstRank]->RxWaitDone(subStream);
121 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]RxWaitDone failed"), ret);
122 0 : ret = links[dstRank]->TxWaitDone(subStream);
123 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]TxWaitDone failed"), ret);
124 : }
125 0 : return HCCL_SUCCESS;
126 0 : }
127 :
128 : // allgather的入口函数
129 0 : HcclResult AllGatherMesh::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
130 : {
131 0 : HcclResult ret = HCCL_SUCCESS;
132 0 : CHK_SMART_PTR_NULL(dispatcher_);
133 0 : CHK_PTR_NULL(stream_.ptr());
134 0 : HCCL_INFO(
135 : "AllGatherMesh run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
136 : inputMem_.ptr(), outputMem_.ptr(), count_);
137 :
138 0 : interRank_ = rank;
139 0 : interRankSize_ = rankSize;
140 :
141 0 : if (interRankSize_ == 1) {
142 0 : if (inputMem_ != outputMem_) {
143 0 : HCCL_DEBUG("rank[%u] mem copy async from input to output", rank);
144 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
145 : }
146 0 : return ret;
147 : }
148 :
149 0 : if (links.size() < rankSize) {
150 0 : HCCL_ERROR("[AllGatherMesh][RunAsync]rank[%u] linksize error", rank);
151 0 : return HCCL_E_INTERNAL;
152 : }
153 0 : u32 subStreamSize = interRankSize_ - 2; // 子流大小等于ranksize-2
154 0 : if (meshStreams_.size() < subStreamSize || (*meshSignal_).size() < subStreamSize
155 0 : || (*meshSignalAux_).size() < subStreamSize) {
156 0 : HCCL_ERROR(
157 : "[AllGatherMesh][RunAsync]AllGatherMesh stream size error: "
158 : "rank[%u] totalrank:%u substreamsize[%llu] signalsize[%llu], signal_aux size[%llu]",
159 : rank, rankSize, meshStreams_.size(), (*meshSignal_).size(), (*meshSignalAux_).size());
160 0 : return HCCL_E_PARA;
161 : }
162 :
163 0 : u32 unitSize = DataUnitSize(dataType_);
164 0 : if (unitSize == 0) {
165 0 : HCCL_ERROR("[AllGatherMesh][RunAsync]rank[%u] Unit Data Size is zero", rank);
166 0 : return HCCL_E_INTERNAL;
167 : }
168 :
169 0 : std::vector<Slice> inputSlices(slices_);
170 0 : if (slices_.size() == 0) {
171 0 : slices_.resize(interRankSize_);
172 0 : inputSlices.resize(rankSize);
173 :
174 : // 生成std::vector<Slice> slices_
175 0 : u64 sliceSize = count_ * unitSize;
176 :
177 0 : for (u32 i = 0; i < interRankSize_; i++) {
178 0 : slices_[i].size = sliceSize;
179 0 : slices_[i].offset = (i * sliceSize);
180 :
181 0 : inputSlices[i].size = sliceSize;
182 0 : inputSlices[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
183 0 : HCCL_DEBUG(
184 : "rank[%u], slices[%u].offset=%llu, slices[%u].size=%llu", rank, i, slices_[i].offset, i,
185 : slices_[i].size);
186 : }
187 : }
188 :
189 0 : for (u32 i = 0; i < interRankSize_; i++) {
190 0 : HCCL_DEBUG(
191 : "[AllGatherMesh][Outputslice]: size[%llu] offset[%llu] inputslice: size[%llu] offset[%llu]",
192 : slices_[i].size, slices_[i].offset, inputSlices[i].size, inputSlices[i].offset);
193 : }
194 :
195 0 : if (inputMem_ != outputMem_) {
196 0 : DeviceMem dst = outputMem_.range(slices_[rank].offset, slices_[rank].size);
197 0 : DeviceMem src = inputMem_.range(inputSlices[rank].offset, inputSlices[rank].size);
198 :
199 0 : HCCL_INFO(
200 : "inputMem != outputMem: rank[%u] copy src[%p] offset[%llu] size[%llu] to dst[%p] offset[%llu]"
201 : "size[%llu]",
202 : rank, src.ptr(), inputSlices[rank].offset, inputSlices[rank].size, dst.ptr(), slices_[rank].offset,
203 : slices_[rank].size);
204 :
205 : // 拷贝到自身rank的output_mem
206 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
207 0 : }
208 :
209 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize-2: stream num
210 0 : HCCL_DEBUG(
211 : "rank[%u] streamindex[%u] wait signalaux[%p]", rank, streamIndex, (*meshSignalAux_)[streamIndex]->ptr());
212 0 : CHK_RET(LocalNotify::Wait(
213 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
214 :
215 0 : HCCL_DEBUG(
216 : "rank[%u] siganl_aux index[%u] signal record signalaux[%p]", rank, streamIndex,
217 : (*meshSignalAux_)[streamIndex]->ptr());
218 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
219 : }
220 0 : CHK_RET(RunAllGather(links, slices_, inputSlices));
221 :
222 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize - 2 stream num
223 0 : HCCL_DEBUG("rank[%u] streamindex[%u] wait signal[%p] ", rank, streamIndex, (*meshSignal_)[streamIndex]->ptr());
224 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
225 :
226 0 : HCCL_DEBUG("rank[%u] streamindex[%u] record signal[%p]", rank, streamIndex, meshStreams_[streamIndex].ptr());
227 0 : CHK_RET(LocalNotify::Post(
228 : meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
229 : }
230 :
231 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
232 0 : HCCL_INFO("AllGatherMesh finished: rank[%u]", rank);
233 0 : return HCCL_SUCCESS;
234 0 : }
235 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_MESH, AllGatherMesh);
236 : } // namespace hccl
|
