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