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