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.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : BroadcastNHR::BroadcastNHR(const HcclDispatcher dispatcher) : NHRBase(dispatcher) {}
16 :
17 0 : BroadcastNHR::~BroadcastNHR() {}
18 :
19 : HcclResult
20 0 : BroadcastNHR::RunAsync(const u32 rank, const u32 rankSize, const std::vector<std::shared_ptr<Transport>>& links)
21 : {
22 0 : CHK_SMART_PTR_NULL(dispatcher_);
23 0 : CHK_PTR_NULL(stream_.ptr());
24 0 : HCCL_INFO("[BroadcastNHR][RunAsync] run: rank[%u] totalrank[%u] count[%llu]", rank, rankSize, count_);
25 :
26 0 : if (rankSize == 1) {
27 0 : return HCCL_SUCCESS;
28 : }
29 :
30 0 : CHK_PRT_RET(
31 : links.size() < rankSize,
32 : HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] linksize[%llu] is less than rank size", rank, links.size()),
33 : HCCL_E_INTERNAL);
34 :
35 0 : u32 unitSize = DataUnitSize(dataType_);
36 0 : CHK_PRT_RET(
37 : unitSize == 0, HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] unit data size is zero", rank), HCCL_E_INTERNAL);
38 :
39 0 : DeviceMem srcMem = inputMem_.range(baseOffset_, count_ * unitSize);
40 0 : DeviceMem dstMem = outputMem_.range(baseOffset_, count_ * unitSize);
41 :
42 0 : if (inputMem_ != outputMem_ && rank == root_) {
43 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, stream_));
44 : }
45 :
46 : // 准备slice
47 0 : PrepareSlice(rank, rankSize);
48 :
49 : // scatter
50 : std::unique_ptr<AlgTemplateBase> tempAlgScatter
51 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_SCATTER_NHR, dispatcher_);
52 0 : CHK_SMART_PTR_NULL(tempAlgScatter);
53 0 : CHK_RET(tempAlgScatter->Prepare(true));
54 0 : tempAlgScatter->CloseBarrier();
55 0 : CHK_RET(tempAlgScatter->Prepare(
56 : srcMem, srcMem, srcMem, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
57 :
58 0 : CHK_RET(
59 : tempAlgScatter->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
60 0 : CHK_RET(tempAlgScatter->RunAsync(rank, rankSize, links));
61 :
62 : // allgather
63 : std::unique_ptr<AlgTemplateBase> tempAlgAllgather
64 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_NHR, dispatcher_);
65 0 : CHK_SMART_PTR_NULL(tempAlgAllgather);
66 0 : CHK_RET(tempAlgAllgather->Prepare(true));
67 0 : CHK_RET(tempAlgAllgather->Prepare(
68 : srcMem, srcMem, srcMem, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
69 :
70 0 : CHK_RET(
71 : tempAlgAllgather->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
72 0 : CHK_RET(tempAlgAllgather->RunAsync(rank, rankSize, links));
73 :
74 0 : HCCL_INFO("[BroadcastNHR][RunAsync] finished: rank[%u] end count[%llu]", rank, count_);
75 0 : return HCCL_SUCCESS;
76 0 : }
77 :
78 0 : HcclResult BroadcastNHR::PrepareSlice(const u32 rank, const u32 rankSize)
79 : {
80 0 : u32 unitSize = DataUnitSize(dataType_);
81 :
82 : // 所有的数据平均到每个rank上
83 0 : u64 sizeAvg = ((count_ + rankSize - 1) / rankSize) * unitSize;
84 0 : u64 sizePerSlice = AlgTemplateBase::RoundUpWithDivisor(sizeAvg, HCCL_MIN_SLICE_ALIGN);
85 0 : HCCL_DEBUG(
86 : "[BroadcastNHR][PrepareSlice] bcast total count[%llu] sizeAverage[%llu] sizePerSlice after aligns[%llu]",
87 : count_, sizeAvg, sizePerSlice);
88 :
89 : // 准备slice
90 0 : slices_.resize(rankSize);
91 0 : u64 sizeResidue = count_ * unitSize;
92 0 : u64 sizePerRound = 0;
93 :
94 0 : for (u32 i = 0; i < rankSize; i++) {
95 0 : sizePerRound = (sizeResidue > sizePerSlice) ? sizePerSlice : sizeResidue;
96 0 : slices_[i].offset = count_ * unitSize - sizeResidue;
97 0 : slices_[i].size = sizePerRound;
98 :
99 0 : sizeResidue -= sizePerRound;
100 0 : HCCL_DEBUG(
101 : "[BroadcastNHR][PrepareSlice] rank[%u] default slice[%u]: offset: [%llu] size[%llu]", rank, i,
102 : slices_[i].offset, slices_[i].size);
103 : }
104 0 : return HCCL_SUCCESS;
105 : }
106 :
107 : HcclResult
108 0 : BroadcastNHR::GetNslbAdjInfo(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
109 : {
110 0 : if (rankSize == 1) {
111 0 : return HCCL_SUCCESS;
112 : }
113 0 : if (links.size() < rankSize) {
114 0 : return HCCL_SUCCESS;
115 : }
116 0 : u32 nSteps = 0;
117 0 : for (u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps) {
118 : }
119 :
120 0 : u32 deltaRoot = (rankSize - rank) % rankSize;
121 : // 先执行 scatter 流程
122 0 : for (u32 step = 0; step < nSteps; step++) {
123 0 : u32 deltaRankPair = 1 << step;
124 0 : u32 nRanks = 0;
125 0 : bool isPerfect = (rankSize & (rankSize - 1)) == 0;
126 0 : if (!isPerfect && step == nSteps - 1) {
127 0 : nRanks = rankSize - deltaRankPair;
128 : } else {
129 0 : nRanks = deltaRankPair;
130 : }
131 0 : if (deltaRoot >= nRanks) {
132 0 : continue;
133 : }
134 0 : u32 sendTo = (rank + rankSize - deltaRankPair) % rankSize;
135 0 : LINK linkRight = links[sendTo];
136 0 : CHK_SMART_PTR_NULL(linkRight);
137 :
138 0 : NslbDpAdjInfo adjInfoStep = {};
139 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
140 0 : adjInfoStep.phaseId = step + 1;
141 0 : adjInfoStep.rev = 0;
142 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
143 0 : }
144 0 : u32 begin = nSteps;
145 : // 后续执行AllGather的NHR流程
146 0 : for (u32 step = 0; step < nSteps; step++) {
147 0 : u32 deltaRank = 1 << (nSteps - 1 - step);
148 0 : u32 sendTo = (rank + deltaRank) % rankSize;
149 0 : LINK linkRight = links[sendTo];
150 0 : CHK_SMART_PTR_NULL(linkRight);
151 :
152 0 : NslbDpAdjInfo allGatherInfoStep = {};
153 0 : allGatherInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
154 0 : allGatherInfoStep.phaseId = step + begin + 1;
155 0 : allGatherInfoStep.rev = 0;
156 0 : nslbAdjInfo.nsAdjInfo.push_back(allGatherInfoStep);
157 0 : }
158 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
159 0 : return HCCL_SUCCESS;
160 : }
161 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_BROADCAST_NHR, BroadcastNHR);
162 : } // namespace hccl
|
