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