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 "reduce_scatter_mesh_mix.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 : using namespace std;
16 :
17 1 : ReduceScatterMeshMix::ReduceScatterMeshMix(const HcclDispatcher dispatcher)
18 1 : : AlgTemplateBase(dispatcher)
19 1 : {}
20 :
21 2 : ReduceScatterMeshMix::~ReduceScatterMeshMix() {}
22 :
23 1 : HcclResult ReduceScatterMeshMix::Prepare(DeviceMem &inputMem, DeviceMem &outputMem, DeviceMem &scratchMem,
24 : const u64 count, const HcclDataType dataType, const Stream &stream,
25 : const HcclReduceOp reductionOp, const u32 root,
26 : const std::vector<Slice> &slices, const u64 baseOffset,
27 : const u64 reduceAttrBitMap, std::vector<Stream> &meshStreams,
28 : const std::vector<std::shared_ptr<LocalNotify>> &meshSignal,
29 : const std::vector<std::shared_ptr<LocalNotify>> &meshSignalAux,
30 : u32 interRank, u32 interRankSize, HcomCollOpInfo *opInfo)
31 : {
32 1 : reduceAttr_ = reduceAttrBitMap;
33 1 : meshStreams_ = meshStreams;
34 1 : meshSignalPtr_ = &meshSignal;
35 1 : meshSignalAuxPtr_ = &meshSignalAux;
36 1 : interRank_ = interRank;
37 1 : interRankSize_ = interRankSize;
38 1 : opInfo_ = opInfo;
39 4 : return AlgTemplateBase::Prepare(inputMem, outputMem, scratchMem, count, dataType,
40 2 : stream, reductionOp, root, slices, baseOffset);
41 : }
42 :
43 0 : HcclResult ReduceScatterMeshMix::MainRecordSub()
44 : {
45 0 : for (u32 signalIndex = 0; signalIndex < meshSignalAuxPtr_->size(); signalIndex++) {
46 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[signalIndex],
47 : profilerInput_.stage));
48 : }
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 0 : HcclResult ReduceScatterMeshMix::SubWaitMain()
53 : {
54 0 : for (u32 streamIndex = 0; streamIndex < meshSignalAuxPtr_->size(); streamIndex++) {
55 0 : CHK_RET(LocalNotify::Wait(meshStreams_[streamIndex], dispatcher_,
56 : (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
57 : }
58 0 : return HCCL_SUCCESS;
59 : }
60 :
61 0 : HcclResult ReduceScatterMeshMix::MainWaitSub()
62 : {
63 0 : for (u32 signalIndex = 0; signalIndex < meshSignalPtr_->size(); signalIndex++) {
64 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
65 : }
66 0 : return HCCL_SUCCESS;
67 : }
68 :
69 0 : HcclResult ReduceScatterMeshMix::SubRecordMain()
70 : {
71 0 : for (u32 streamIndex = 0; streamIndex < meshSignalPtr_->size(); streamIndex++) {
72 0 : CHK_RET(LocalNotify::Post(meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex],
73 : profilerInput_.stage));
74 : }
75 0 : return HCCL_SUCCESS;
76 : }
77 :
78 0 : HcclResult ReduceScatterMeshMix::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
79 : {
80 0 : HCCL_INFO("ReduceScatterMeshMix run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank,
81 : rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
82 :
83 : // 数据准备
84 0 : u32 unitSize = SIZE_TABLE[dataType_];
85 0 : u64 totalSize = opInfo_->count * unitSize;
86 0 : u64 sliceSize = count_ * unitSize;
87 :
88 0 : u8* curUerMemInPtr = static_cast<u8 *>(opInfo_->inputAddr);
89 0 : u8* curCommMemOutPtr = static_cast<u8 *>(outputMem_.ptr());
90 :
91 0 : DeviceMem userMemIn = DeviceMem::create(curUerMemInPtr + totalSize * rank, sliceSize);
92 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
93 :
94 0 : DeviceMem src;
95 0 : DeviceMem dst;
96 :
97 0 : for (u32 i = 0; i < interRankSize_; i++) {
98 0 : src = DeviceMem::create(curUerMemInPtr + (i * rankSize + rank) * totalSize, sliceSize);
99 0 : dst = DeviceMem::create(curCommMemOutPtr + (i * rankSize + rank) * sliceSize, sliceSize);
100 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
101 : }
102 :
103 0 : CHK_RET(MainRecordSub());
104 0 : CHK_RET(SubWaitMain());
105 :
106 : // 每个stream只负责一个对端的交互
107 0 : for (u32 round = 1; round < rankSize; round++) {
108 0 : u32 dstRank = (round + rank) % rankSize;
109 0 : const LINK &dstLink = links[dstRank];
110 0 : Stream &subStream = meshStreams_[round - 1];
111 0 : CHK_RET(dstLink->TxAck(subStream));
112 0 : CHK_RET(dstLink->RxAck(subStream));
113 : }
114 0 : CHK_RET(SubRecordMain());
115 0 : CHK_RET(MainWaitSub());
116 : // 为子图增加一个从stream到主stream的附着点
117 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
118 :
119 0 : CHK_RET(SubWaitMain());
120 0 : CHK_RET(MainRecordSub());
121 :
122 : // inline执行notice reduce
123 0 : for (u32 round = 1; round < rankSize; round++) {
124 0 : u32 dstRank = (round + rank) % rankSize;
125 0 : const LINK &dstLink = links[dstRank];
126 0 : Stream &subStream = meshStreams_[round - 1];
127 : // 本rank要发数据
128 0 : void *remMemPtr = nullptr;
129 : // 获取远端的commoutMem
130 0 : CHK_RET(dstLink->GetRemoteMem(UserMemType::INPUT_MEM, &remMemPtr));
131 :
132 0 : for (u32 i = 0; i < interRankSize_; i++) {
133 0 : src = DeviceMem::create(curUerMemInPtr + (i * rankSize + dstRank) * totalSize, sliceSize);
134 0 : dst = DeviceMem::create(static_cast<u8 *>(remMemPtr) + (i * rankSize + dstRank) * sliceSize, sliceSize);
135 0 : CHK_RET(HcclReduceAsync(dispatcher_, static_cast<void *>(src.ptr()), count_, dataType_, reductionOp_,
136 : subStream, static_cast<void *>(dst.ptr()), dstLink->GetRemoteRank(), dstLink->GetLinkType(),
137 : INLINE_REDUCE_BIT));
138 : }
139 :
140 0 : CHK_RET(dstLink->TxDataSignal(subStream));
141 0 : CHK_RET(dstLink->RxDataSignal(subStream));
142 : }
143 :
144 0 : CHK_RET(SubRecordMain());
145 0 : CHK_RET(MainWaitSub());
146 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
147 :
148 0 : HCCL_INFO("ReduceScatterMeshMix finished: rank[%u]", rank);
149 0 : return HCCL_SUCCESS;
150 0 : }
151 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_MESH_MIX, ReduceScatterMeshMix);
152 : }
|
