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 "all_reduce_nhr_oneshot.h"
12 : #include "alg_template_register.h"
13 : #include "reduce_nhr_oneshot.h"
14 : #include "broadcast_nhr_oneshot.h"
15 :
16 : namespace hccl {
17 10 : AllReduceNHROneshot::AllReduceNHROneshot(const HcclDispatcher dispatcher) : NHRBase(dispatcher) {}
18 :
19 20 : AllReduceNHROneshot::~AllReduceNHROneshot() {}
20 :
21 10 : HcclResult AllReduceNHROneshot::Prepare(u64 reduceAttrBitMap, [[maybe_unused]] HcomCollOpInfo* opInfo)
22 : {
23 10 : reduceAttr_ = reduceAttrBitMap;
24 10 : return HCCL_SUCCESS;
25 : }
26 :
27 3 : HcclResult AllReduceNHROneshot::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
28 : {
29 : // 基本的检查
30 3 : CHK_RET(SimpleCheck(rank, rankSize, links));
31 3 : HCCL_INFO(
32 : "[AllReduceNHROneshot][RunAsync] run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank,
33 : rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
34 :
35 3 : HcclResult ret = HCCL_SUCCESS;
36 : // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
37 3 : if (rankSize == 1) {
38 0 : if (inputMem_ != outputMem_) {
39 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
40 0 : CHK_PRT_RET(
41 : ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceNHROneshot][RunAsync] rank[%u] memcpy async failed", rank),
42 : ret);
43 : }
44 :
45 0 : return ret;
46 : }
47 :
48 : // 先执行1-reduce
49 3 : ret = RunReduceOneshot(rank, rankSize, links);
50 3 : CHK_PRT_RET(
51 : ret != HCCL_SUCCESS,
52 : HCCL_ERROR(
53 : "[AllReduceNHROneshot][RunAsync] rank[%u] count[%llu] failed in "
54 : "1-reduce step",
55 : rank, count_),
56 : ret);
57 :
58 : // 再执行1-bcast
59 3 : ret = RunBroadcastOneshot(rank, rankSize, links);
60 3 : CHK_PRT_RET(
61 : ret != HCCL_SUCCESS,
62 : HCCL_ERROR(
63 : "[AllReduceNHROneshot][RunAsync] rank[%u] count[%llu] failed in "
64 : "1-bcast step",
65 : rank, count_),
66 : ret);
67 :
68 3 : HCCL_INFO("[AllReduceNHROneshot][RunAsync] finished: rank[%u] ranksize[%u]", rank, rankSize);
69 3 : return HCCL_SUCCESS;
70 : }
71 :
72 3 : HcclResult AllReduceNHROneshot::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
73 : {
74 : // 判断stream, dispatcher是否为空
75 3 : CHK_SMART_PTR_NULL(dispatcher_);
76 3 : CHK_PTR_NULL(stream_.ptr());
77 :
78 : // 检查memory
79 3 : CHK_PRT_RET(
80 : !outputMem_ || !inputMem_,
81 : HCCL_ERROR("[AllReduceNHROneshot][SimpleCheck] rank[%u] inputmem or outputmem is null", rank), HCCL_E_PTR);
82 :
83 : // 判断links数量是否正确
84 3 : CHK_PRT_RET(
85 : links.size() < rankSize,
86 : HCCL_ERROR(
87 : "[AllReduceNHROneshot][SimpleCheck] rank[%u] link size[%llu] is "
88 : "less than rank size[%u]",
89 : rank, links.size(), rankSize),
90 : HCCL_E_INTERNAL);
91 :
92 3 : return HCCL_SUCCESS;
93 : }
94 :
95 3 : HcclResult AllReduceNHROneshot::RunReduceOneshot(u32 rank, u32 rankSize, const std::vector<LINK>& links)
96 : {
97 3 : std::unique_ptr<AlgTemplateBase> tempAlg;
98 3 : tempAlg = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCE_NHR_ONE_SHOT, dispatcher_);
99 3 : CHK_SMART_PTR_NULL(tempAlg);
100 3 : CHK_RET(tempAlg->Prepare(reduceAttr_));
101 3 : if (!barrierSwitchOn_) {
102 0 : tempAlg->CloseBarrier();
103 : }
104 3 : HCCL_INFO(
105 : "[AllReduceNHROneshot][RunReduceOneshot] 1-reduce tempAlg rank[%u] inputMem[%p] outputMem[%p] "
106 : "mem_size[%llu] count[%llu] planeID:[%d]",
107 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
108 9 : CHK_RET(tempAlg->Prepare(
109 : inputMem_, inputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
110 :
111 3 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
112 :
113 3 : return tempAlg->RunAsync(rank, rankSize, links);
114 3 : }
115 :
116 3 : HcclResult AllReduceNHROneshot::RunBroadcastOneshot(u32 rank, u32 rankSize, const std::vector<LINK>& links)
117 : {
118 3 : BroadcastNHROneshot tempAlg(dispatcher_);
119 3 : HCCL_INFO(
120 : "[AllReduceNHROneshot][RunBroadcastOneshot] 1-broadcast tempAlg rank[%u] inputMem[%p] outputMem[%p] "
121 : "mem_size[%llu] count[%llu] planeID:[%d]",
122 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
123 9 : CHK_RET(tempAlg.Prepare(
124 : inputMem_, outputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
125 :
126 3 : CHK_RET(tempAlg.RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
127 :
128 3 : return tempAlg.RunAsyncForAllReduce(rank, rankSize, links);
129 3 : }
130 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_NHR_ONESHOT, AllReduceNHROneshot);
131 : } // namespace hccl
|
