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 "broadcast_nhr_oneshot.h"
12 : #include <cmath>
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 3 : BroadcastNHROneshot::BroadcastNHROneshot(const HcclDispatcher dispatcher)
17 : : NHRBase(dispatcher),
18 3 : localBaseOffset_(0),
19 3 : isForAllReduce_(false)
20 3 : {}
21 :
22 3 : BroadcastNHROneshot::~BroadcastNHROneshot() {}
23 :
24 3 : HcclResult BroadcastNHROneshot::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
25 : {
26 : // 基本的检查
27 3 : CHK_RET(SimpleCheck(rank, rankSize, links));
28 3 : HCCL_INFO(
29 : "[BroadcastNHROneshot][RunAsync] rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
30 : inputMem_.ptr(), outputMem_.ptr(), count_);
31 :
32 3 : u32 unitSize = DataUnitSize(dataType_);
33 3 : CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[BroadcastNHROneshot][RunAsync] unitSize is zero"), HCCL_E_INTERNAL);
34 :
35 3 : if (!isForAllReduce_) {
36 0 : localBaseOffset_ = baseOffset_; // broadcast的本地偏移量和baseOffset_一致
37 : }
38 :
39 : // 双buffer下, 先将input拷贝到output的合适位置
40 3 : if (inputMem_ != outputMem_ && rank == root_) {
41 3 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
42 3 : DeviceMem src = inputMem_.range(localBaseOffset_, totalSize);
43 3 : DeviceMem dst = outputMem_.range(localBaseOffset_, totalSize);
44 3 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
45 3 : }
46 :
47 : // 如果ranksize为1, 从input->output就结束
48 3 : if (rankSize == 1) {
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 : // 运行bcast， rst算法
53 3 : CHK_RET(RunBroadcastNHROneshot(rank, rankSize, links));
54 :
55 3 : HCCL_INFO("[BroadcastNHROneshot][RunAsync] finished: rank[%u] ranksize[%u]", rank, rankSize);
56 3 : return HCCL_SUCCESS;
57 : }
58 :
59 3 : HcclResult BroadcastNHROneshot::RunAsyncForAllReduce(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
60 : {
61 3 : isForAllReduce_ = true;
62 3 : return RunAsync(rank, rankSize, links);
63 : }
64 :
65 3 : HcclResult BroadcastNHROneshot::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
66 : {
67 : // 判断stream, dispatcher是否为空
68 3 : CHK_SMART_PTR_NULL(dispatcher_);
69 3 : CHK_PTR_NULL(stream_.ptr());
70 :
71 : // 检查memory
72 3 : CHK_PRT_RET(
73 : !outputMem_ || !inputMem_,
74 : HCCL_ERROR("[BroadcastNHROneshot][SimpleCheck] rank[%u] inputmem or outputmem is null", rank), HCCL_E_PTR);
75 :
76 : // 判断links数量是否正确
77 3 : CHK_PRT_RET(
78 : links.size() < rankSize,
79 : HCCL_ERROR(
80 : "[BroadcastNHROneshot][SimpleCheck] rank[%u] link size[%llu] is "
81 : "less than rank size[%u]",
82 : rank, links.size(), rankSize),
83 : HCCL_E_INTERNAL);
84 3 : return HCCL_SUCCESS;
85 : }
86 :
87 0 : HcclResult BroadcastNHROneshot::SdmaRx(
88 : LINK& linkLeft, LINK& linkRight, [[maybe_unused]] InterServerAlgoStep& stepInfo,
89 : [[maybe_unused]] const std::vector<LINK>& links)
90 : {
91 0 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
92 0 : DeviceMem srcMem = outputMem_.range(localBaseOffset_, totalSize);
93 :
94 0 : if (linkRight != nullptr) {
95 0 : CHK_RET(linkRight->TxAck(stream_));
96 : }
97 0 : if (linkLeft != nullptr) {
98 0 : CHK_RET(linkLeft->RxAck(stream_));
99 0 : void* srcMemPtr = nullptr;
100 0 : CHK_RET(linkLeft->GetRemoteMem(UserMemType::OUTPUT_MEM, &srcMemPtr));
101 0 : DeviceMem srcMemLeft(static_cast<s8*>(srcMemPtr) + baseOffset_, totalSize);
102 0 : CHK_RET(HcclD2DMemcpyAsync(
103 : dispatcher_, srcMem, srcMemLeft, stream_, linkLeft->GetRemoteRank(), // Memecpy
104 : linkLeft->GetLinkType()));
105 0 : CHK_RET(linkLeft->TxDataSignal(stream_));
106 0 : }
107 0 : if (linkRight != nullptr) {
108 0 : CHK_RET(linkRight->RxDataSignal(stream_));
109 : }
110 0 : return HCCL_SUCCESS;
111 0 : }
112 :
113 9 : HcclResult BroadcastNHROneshot::RdmaTxRx(
114 : LINK& linkLeft, LINK& linkRight, [[maybe_unused]] InterServerAlgoStep& stepInfo,
115 : [[maybe_unused]] const std::vector<LINK>& links)
116 : {
117 9 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
118 9 : DeviceMem srcMem = outputMem_.range(localBaseOffset_, totalSize);
119 :
120 9 : if (linkLeft != nullptr) {
121 0 : CHK_RET(linkLeft->TxAck(stream_));
122 : }
123 :
124 9 : if (linkRight != nullptr) {
125 9 : CHK_RET(linkRight->RxAck(stream_));
126 9 : CHK_RET(linkRight->TxAsync(UserMemType::OUTPUT_MEM, baseOffset_, srcMem.ptr(), srcMem.size(), stream_));
127 9 : CHK_RET(linkRight->WaitFinAck(stream_));
128 : }
129 :
130 9 : if (linkLeft != nullptr) {
131 0 : CHK_RET(linkLeft->RxAsync(UserMemType::OUTPUT_MEM, baseOffset_, srcMem.ptr(), srcMem.size(), stream_));
132 0 : CHK_RET(linkLeft->PostFinAck(stream_));
133 : }
134 9 : return HCCL_SUCCESS;
135 9 : }
136 :
137 3 : HcclResult BroadcastNHROneshot::RunBroadcastNHROneshot(u32 rank, u32 rankSize, const std::vector<LINK>& links)
138 : {
139 : // 计算通信步数
140 3 : u32 nSteps = GetStepNumInterServer(rankSize);
141 3 : HCCL_DEBUG(
142 : "[BroadcastNHROneshot][RunBroadcastNHROneshot] rank[%u] rankSize[%u] nSteps[%u]", rank, rankSize, nSteps);
143 :
144 : // 逐步编排任务
145 12 : for (u32 step = 0; step < nSteps; step++) {
146 9 : InterServerAlgoStep stepInfo;
147 9 : GetStepInfo(step, nSteps, rank, rankSize, stepInfo);
148 :
149 9 : HCCL_DEBUG(
150 : "[BroadcastNHROneshot][RunBroadcastNHROneshot] recvFrom[%u] sendTo[%u] step[%u]", stepInfo.fromRank,
151 : stepInfo.toRank, step);
152 :
153 9 : LINK linkLeft;
154 9 : LINK linkRight;
155 9 : if (stepInfo.txSliceIdxs.size() > 0) {
156 9 : linkRight = links[stepInfo.toRank];
157 9 : CHK_SMART_PTR_NULL(linkRight);
158 : }
159 9 : if (stepInfo.rxSliceIdxs.size() > 0) {
160 0 : linkLeft = links[stepInfo.fromRank];
161 0 : CHK_SMART_PTR_NULL(linkLeft);
162 : }
163 :
164 18 : if ((linkRight != nullptr && linkRight->IsSpInlineReduce())
165 18 : || (linkLeft != nullptr && linkLeft->IsSpInlineReduce())) {
166 0 : CHK_RET(SdmaRx(linkLeft, linkRight, stepInfo, links));
167 : } else {
168 9 : CHK_RET(RdmaTxRx(linkLeft, linkRight, stepInfo, links));
169 : }
170 9 : }
171 3 : return HCCL_SUCCESS;
172 : }
173 :
174 : // NHR每步的算法描述原理函数
175 9 : HcclResult BroadcastNHROneshot::GetStepInfo(u32 step, u32 nSteps, u32 rank, u32 rankSize, InterServerAlgoStep& stepInfo)
176 : {
177 9 : stepInfo.txSliceIdxs.clear();
178 9 : stepInfo.rxSliceIdxs.clear();
179 9 : stepInfo.nSlices = 1;
180 9 : stepInfo.toRank = rankSize;
181 9 : stepInfo.fromRank = rankSize;
182 9 : stepInfo.step = step;
183 9 : stepInfo.myRank = rank;
184 :
185 9 : u32 nRanks = (rankSize - 1 + (1 << (nSteps - 1 - step))) / (1 << (nSteps - step)); // 本步需要进行收/发的rank数
186 :
187 9 : u32 deltaRoot = (rank + rankSize - root_) % rankSize;
188 :
189 9 : u32 deltaRankPair = 1 << (nSteps - 1 - step);
190 9 : u32 deltaRankGroup = 1 << (nSteps - step);
191 :
192 9 : if (deltaRoot / deltaRankGroup < nRanks) {
193 9 : if (deltaRoot % deltaRankGroup == 0) {
194 9 : stepInfo.toRank = (rank + deltaRankPair) % rankSize;
195 9 : stepInfo.txSliceIdxs.push_back(0);
196 : }
197 :
198 9 : if ((deltaRoot + deltaRankPair) % deltaRankGroup == 0) {
199 0 : stepInfo.fromRank = (rank + rankSize - deltaRankPair) % rankSize;
200 0 : stepInfo.rxSliceIdxs.push_back(0);
201 : }
202 : }
203 9 : return HCCL_SUCCESS;
204 : }
205 :
206 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_BROADCAST_NHR_ONESHOT, BroadcastNHROneshot);
207 : } // namespace hccl
|
