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