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