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 "reduce_scatter_mesh_atomic.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 : using namespace std;
16 :
17 1 : ReduceScatterMeshAtomic::ReduceScatterMeshAtomic(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
18 :
19 2 : ReduceScatterMeshAtomic::~ReduceScatterMeshAtomic() {}
20 :
21 1 : HcclResult ReduceScatterMeshAtomic::Prepare(
22 : DeviceMem& inputMem, DeviceMem& outputMem, DeviceMem& scratchMem, const u64 count, const HcclDataType dataType,
23 : const Stream& stream, const HcclReduceOp reductionOp, const u32 root, const std::vector<Slice>& slices,
24 : const u64 baseOffset, const u64 reduceAttrBitMap, std::vector<Stream>& meshStreams,
25 : std::vector<std::shared_ptr<LocalNotify>>& meshSignal, std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux,
26 : u32 userRank, const HcomCollOpInfo* opInfo)
27 : {
28 1 : reduceAttr_ = reduceAttrBitMap;
29 1 : userRank_ = userRank;
30 1 : meshStreams_ = meshStreams;
31 1 : meshSignalPtr_ = &meshSignal;
32 1 : meshSignalAuxPtr_ = &meshSignalAux;
33 : (void)opInfo;
34 4 : return AlgTemplateBase::Prepare(
35 2 : inputMem, outputMem, scratchMem, count, dataType, stream, reductionOp, root, slices, baseOffset);
36 : }
37 :
38 0 : HcclResult ReduceScatterMeshAtomic::RunReduceScatter(const std::vector<LINK>& links)
39 : {
40 : // 拼接所有stream
41 0 : vector<Stream> streamVct;
42 0 : streamVct.reserve(localRankSize_ - 1); // 有ranksize-1个对端，每个对端对应一条stream
43 0 : streamVct.push_back(stream_); // 增加主stream
44 0 : streamVct.insert(streamVct.end(), meshStreams_.begin(), meshStreams_.end()); // 增加从stream
45 :
46 : // 每个stream只负责一个对端的交互
47 0 : for (u32 streamIndex = 0; streamIndex < localRankSize_ - 1; streamIndex++) {
48 0 : u32 remoteRank = (streamIndex + localRank_ + 1) % localRankSize_;
49 0 : const LINK& dstLink = links[remoteRank];
50 0 : Stream& stream = streamVct[streamIndex];
51 :
52 0 : CHK_RET(dstLink->TxAck(stream));
53 0 : CHK_RET(dstLink->RxAck(stream));
54 : }
55 :
56 0 : for (u32 streamIndex = 0; streamIndex < localRankSize_ - 1; streamIndex++) {
57 0 : u32 remoteRank = (streamIndex + localRank_ + 1) % localRankSize_;
58 0 : const LINK& dstLink = links[remoteRank];
59 0 : Stream& stream = streamVct[streamIndex];
60 0 : profilerInput_.streamID = stream.id();
61 0 : profilerInput_.planeID = streamIndex - 1;
62 0 : profilerInput_.step = HCCL_EXEC_STEP_NOT_SET;
63 :
64 0 : Slice& rxSlice = slices_[localRank_];
65 0 : if (streamIndex == 0) {
66 0 : for (u32 signalIndex = 0; signalIndex < localRankSize_ - 2; signalIndex++) { // rankSize-2: stream num
67 0 : CHK_RET(LocalNotify::Wait(stream, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
68 : }
69 0 : for (u32 signalIndex = 0; signalIndex < localRankSize_ - 2; signalIndex++) { // rankSize-2: stream num
70 0 : CHK_RET(
71 : LocalNotify::Post(stream, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
72 : }
73 : } else {
74 0 : u32 signalIndex = streamIndex - 1;
75 0 : CHK_RET(LocalNotify::Post(stream, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
76 0 : CHK_RET(LocalNotify::Wait(stream, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
77 : }
78 :
79 0 : HCCL_DEBUG(
80 : "ReduceScatterMeshAtomic RX rank[%u] inputMemSize[%llu], rxSlice.offset[%llu], rxSlice.size[%llu] "
81 : "baseOffset_[%llu]",
82 : localRank_, inputMem_.size(), rxSlice.offset, rxSlice.size, baseOffset_);
83 0 : DeviceMem dstMem = inputMem_.range(rxSlice.offset, rxSlice.size);
84 0 : void* remoteMem = nullptr;
85 0 : CHK_RET(dstLink->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMem));
86 0 : CHK_RET(HcclReduceAsync(
87 : dispatcher_, static_cast<s8*>(remoteMem) + baseOffset_ + rxSlice.offset,
88 : dstMem.size() / SIZE_TABLE[dataType_], dataType_, reductionOp_, stream, dstMem.ptr(),
89 : dstLink->GetRemoteRank(), dstLink->GetLinkType(), INLINE_REDUCE_BIT));
90 :
91 0 : CHK_RET(dstLink->TxDataSignal(stream));
92 0 : CHK_RET(dstLink->RxDataSignal(stream));
93 0 : }
94 : // 添加空task,保证执行时不乱序
95 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, streamVct[0], dispatcher_));
96 0 : return HCCL_SUCCESS;
97 0 : }
98 :
99 0 : HcclResult ReduceScatterMeshAtomic::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
100 : {
101 0 : CHK_SMART_PTR_NULL(dispatcher_);
102 0 : CHK_PTR_NULL(stream_.ptr());
103 0 : HCCL_INFO(
104 : "ReduceScatterMeshAtomic run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
105 : inputMem_.ptr(), outputMem_.ptr(), count_);
106 :
107 0 : localRank_ = rank;
108 0 : localRankSize_ = rankSize;
109 :
110 0 : if (localRankSize_ == 1) {
111 0 : if (inputMem_ != outputMem_) {
112 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
113 : }
114 0 : return HCCL_SUCCESS;
115 : }
116 :
117 0 : if (links.size() < rankSize) {
118 0 : HCCL_ERROR("[ReduceScatterMeshAtomic][RunAsync]rank[%u] linksize[%zu] error", rank, links.size());
119 0 : return HCCL_E_INTERNAL;
120 : }
121 0 : CHK_RET(MemSlice());
122 :
123 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize-2: stream num
124 0 : HCCL_DEBUG(
125 : "rank[%u] streamindex[%u] wait signalaux[%p]", rank, streamIndex, (*meshSignalAuxPtr_)[streamIndex]->ptr());
126 0 : CHK_RET(LocalNotify::Wait(
127 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
128 : }
129 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize-2: stream num
130 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
131 : }
132 :
133 0 : CHK_RET(RunReduceScatter(links));
134 :
135 0 : for (u32 streamIndex = 0; streamIndex < rankSize - 2; streamIndex++) { // rankSize - 2 stream num
136 0 : HCCL_DEBUG(
137 : "rank[%u] streamindex[%u] wait signal[%p] ", rank, streamIndex, (*meshSignalPtr_)[streamIndex]->ptr());
138 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
139 0 : CHK_RET(LocalNotify::Post(
140 : meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
141 : }
142 :
143 : // 添加空task,保证执行时不乱序
144 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
145 :
146 0 : if (inputMem_ != outputMem_) {
147 0 : DeviceMem src = inputMem_.range(slices_[localRank_].offset, slices_[localRank_].size);
148 0 : HCCL_DEBUG(
149 : "rank[%u] copy result from to output[%p] offset[%llu] size[%llu] ", localRank_, outputMem_.ptr(),
150 : slices_[localRank_].offset, slices_[localRank_].size);
151 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, src, stream_);
152 0 : CHK_PRT_RET(
153 : ret != HCCL_SUCCESS,
154 : HCCL_ERROR(
155 : "[ReduceScatterMeshAtomic][RunAsync]rank[%u] memcpy async from mem[%p] "
156 : "to ouputmem[%p] failed",
157 : rank, src.ptr(), outputMem_.ptr()),
158 : ret);
159 0 : }
160 :
161 0 : HCCL_INFO("ReduceScatterMeshAtomic finished: rank[%u]", rank);
162 0 : return HCCL_SUCCESS;
163 : }
164 :
165 0 : HcclResult ReduceScatterMeshAtomic::MemSlice()
166 : {
167 0 : u32 unitSize = SIZE_TABLE[dataType_];
168 0 : if (unitSize == 0) {
169 0 : HCCL_ERROR("[ReduceScatterMeshAtomic][RunAsync]rank[%u] unit data size is zero", localRank_);
170 0 : return HCCL_E_INTERNAL;
171 : }
172 :
173 0 : if (HcclCheckLogLevel(DLOG_DEBUG)) {
174 0 : for (size_t i = 0; i < slices_.size(); i++) {
175 0 : HCCL_DEBUG(
176 : "[ReduceScatterMeshAtomic] rank[%u] index[%zu] size[%llu] offset[%llu]", localRank_, i, slices_[i].size,
177 : slices_[i].offset);
178 : }
179 0 : HCCL_DEBUG("[ReduceScatterMeshAtomic] localRankSize[%u]", localRankSize_);
180 : }
181 :
182 0 : if (slices_.size() == 0) {
183 0 : slices_.resize(localRankSize_);
184 0 : u64 sliceSize = count_ * unitSize;
185 0 : for (u32 i = 0; i < localRankSize_; i++) {
186 0 : slices_[i].size = sliceSize;
187 0 : slices_[i].offset = (i * sliceSize);
188 : }
189 : }
190 :
191 0 : scratchSlices_.resize(localRankSize_);
192 0 : for (u32 i = 0; i < localRankSize_; i++) {
193 0 : scratchSlices_[i].size = slices_[i].size;
194 0 : scratchSlices_[i].offset = (scratchMem_.size() < inputMem_.size()) ? 0 : slices_[i].offset;
195 : }
196 :
197 0 : if (HcclCheckLogLevel(DLOG_DEBUG)) {
198 0 : for (size_t i = 0; i < slices_.size(); i++) {
199 0 : HCCL_DEBUG(
200 : "ReduceScatterMeshAtomic rank[%u] index[%zu] "
201 : "size[%llu] offset[%llu] scratch_size[%llu] scratch_offset[%llu]",
202 : localRank_, i, slices_[i].size, slices_[i].offset, scratchSlices_[i].size, scratchSlices_[i].offset);
203 : }
204 : }
205 :
206 0 : return HCCL_SUCCESS;
207 : }
208 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_MESH_ATOMIC, ReduceScatterMeshAtomic);
209 : } // namespace hccl
|
