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_gather_mesh_mix.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AllgatherMeshMix::AllgatherMeshMix(const HcclDispatcher dispatcher)
16 0 : : AlgTemplateBase(dispatcher)
17 0 : {}
18 :
19 0 : AllgatherMeshMix::~AllgatherMeshMix() {}
20 :
21 0 : HcclResult AllgatherMeshMix::Prepare(std::vector<Stream> &meshStreams,
22 : std::vector<std::shared_ptr<LocalNotify>> &meshSignal, std::vector<std::shared_ptr<LocalNotify>> &meshSignalAux,
23 : u32 userRank, HcomCollOpInfo *opInfo, u32 interRank, u32 interRankSize)
24 : {
25 0 : meshStreams_ = meshStreams;
26 0 : meshSignal_ = &meshSignal;
27 0 : meshSignalAux_ = &meshSignalAux;
28 0 : interRank_ = interRank;
29 0 : interRankSize_ = interRankSize;
30 0 : opInfo_ = opInfo;
31 0 : return HCCL_SUCCESS;
32 : }
33 :
34 0 : HcclResult AllgatherMeshMix::MainRecordSub()
35 : {
36 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignalAux_).size(); signalIndex++) {
37 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex],
38 : profilerInput_.stage));
39 : }
40 0 : return HCCL_SUCCESS;
41 : }
42 :
43 0 : HcclResult AllgatherMeshMix::SubWaitMain()
44 : {
45 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignalAux_).size(); streamIndex++) {
46 0 : CHK_RET(LocalNotify::Wait(meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex],
47 : profilerInput_.stage));
48 : }
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 0 : HcclResult AllgatherMeshMix::MainWaitSub()
53 : {
54 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignal_).size(); signalIndex++) {
55 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : HcclResult AllgatherMeshMix::SubRecordMain()
61 : {
62 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignal_).size(); streamIndex++) {
63 0 : CHK_RET(LocalNotify::Post(meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex],
64 : profilerInput_.stage));
65 : }
66 0 : return HCCL_SUCCESS;
67 : }
68 :
69 : // allgather的入口函数
70 0 : HcclResult AllgatherMeshMix::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
71 : {
72 0 : HCCL_INFO("AllGatherMesh run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]",
73 : rank, rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
74 0 : u32 unitSize = DataUnitSize(dataType_);
75 0 : u64 sliceSize = count_ * unitSize; // 当前count
76 0 : u64 totalSize = opInfo_->count * unitSize; // 总输入count
77 :
78 0 : u8* curUerMemOutPtr = static_cast<u8 *>(opInfo_->outputAddr);
79 0 : u8* curCommMemOutPtr = static_cast<u8 *>(outputMem_.ptr());
80 :
81 0 : CHK_RET(MainRecordSub());
82 0 : CHK_RET(SubWaitMain());
83 :
84 0 : for (u32 round = 1; round < rankSize; round++) {
85 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
86 0 : Stream& subStream = meshStreams_[round - 1];
87 0 : CHK_RET(links[dstRank]->TxAck(subStream));
88 0 : CHK_RET(links[dstRank]->RxAck(subStream));
89 : }
90 :
91 0 : CHK_RET(SubRecordMain());
92 0 : CHK_RET(MainWaitSub());
93 :
94 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
95 :
96 0 : CHK_RET(SubWaitMain());
97 0 : CHK_RET(MainRecordSub());
98 :
99 0 : DeviceMem src;
100 0 : DeviceMem dst;
101 0 : for (u32 i = 0; i < interRankSize_; i++) {
102 0 : src = DeviceMem::create(curCommMemOutPtr + (i * rankSize + rank) * sliceSize, sliceSize);
103 0 : dst = DeviceMem::create(curUerMemOutPtr + (i * rankSize + rank) * totalSize, sliceSize);
104 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
105 : }
106 :
107 0 : for (u32 round = 1; round < rankSize; round++) {
108 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
109 0 : Stream& subStream = meshStreams_[round - 1];
110 : // 本rank要收数据
111 0 : void *remMemPtr = nullptr;
112 : // 从对端的input内存拿数据，input==output也没有关系
113 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
114 :
115 0 : for (u32 i = 0; i < interRankSize_; i++) {
116 0 : src = DeviceMem::create(static_cast<u8 *>(remMemPtr) + (i * rankSize + dstRank) * sliceSize, sliceSize);
117 0 : dst = DeviceMem::create(curUerMemOutPtr + (i * rankSize + dstRank) * totalSize, sliceSize);
118 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStream,
119 : links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
120 : }
121 :
122 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
123 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
124 : }
125 0 : CHK_RET(SubRecordMain());
126 0 : CHK_RET(MainWaitSub());
127 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
128 :
129 0 : HCCL_INFO("AllGatherMesh finished: rank[%u]", rank);
130 0 : return HCCL_SUCCESS;
131 0 : }
132 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_MESH_MIX, AllgatherMeshMix);
133 : } // namespace hccl
|
