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_direct.h"
12 : #include "alg_template_register.h"
13 : // userin -> dmaout -> userout
14 : namespace hccl {
15 0 : AllgatherMeshDirect::AllgatherMeshDirect(const HcclDispatcher dispatcher)
16 0 : : AlgTemplateBase(dispatcher)
17 0 : {}
18 :
19 0 : AllgatherMeshDirect::~AllgatherMeshDirect() {}
20 :
21 0 : HcclResult AllgatherMeshDirect::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 : opInfo_ = opInfo;
29 0 : interRank_ = interRank;
30 0 : interRankSize_ = interRankSize;
31 0 : userRank_ = userRank;
32 0 : return HCCL_SUCCESS;
33 : }
34 :
35 0 : HcclResult AllgatherMeshDirect::MainRecordSub()
36 : {
37 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignalAux_).size(); signalIndex++) {
38 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex],
39 : profilerInput_.stage));
40 : }
41 0 : return HCCL_SUCCESS;
42 : }
43 :
44 0 : HcclResult AllgatherMeshDirect::SubWaitMain()
45 : {
46 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignalAux_).size(); streamIndex++) {
47 0 : CHK_RET(LocalNotify::Wait(meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex],
48 : profilerInput_.stage));
49 : }
50 0 : return HCCL_SUCCESS;
51 : }
52 :
53 0 : HcclResult AllgatherMeshDirect::MainWaitSub()
54 : {
55 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignal_).size(); signalIndex++) {
56 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
57 : }
58 0 : return HCCL_SUCCESS;
59 : }
60 :
61 0 : HcclResult AllgatherMeshDirect::SubRecordMain()
62 : {
63 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignal_).size(); streamIndex++) {
64 0 : CHK_RET(LocalNotify::Post(meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex],
65 : profilerInput_.stage));
66 : }
67 0 : return HCCL_SUCCESS;
68 : }
69 :
70 : // allgather的入口函数
71 0 : HcclResult AllgatherMeshDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
72 : {
73 0 : CHK_SMART_PTR_NULL(dispatcher_);
74 0 : CHK_PTR_NULL(stream_.ptr());
75 0 : HCCL_INFO("AllGatherMeshDirect run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]",
76 : rank, rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
77 :
78 0 : char* curUerMemInPtr = static_cast<char *>(opInfo_->inputAddr);
79 0 : char* curUerMemOutPtr = static_cast<char *>(opInfo_->outputAddr);
80 0 : char* curCommMemOutPtr = static_cast<char *>(outputMem_.ptr());
81 :
82 0 : u32 unitSize = DataUnitSize(dataType_);
83 0 : u64 curSize = count_ * unitSize; // 当前count
84 0 : u64 sliceSize = opInfo_->count * unitSize; // 总输入count
85 :
86 0 : if (rankSize == 1) {
87 0 : if (opInfo_->inputAddr != opInfo_->outputAddr) {
88 0 : HCCL_DEBUG("rank[%u] mem copy async from input to output", rank);
89 0 : DeviceMem userMemIn = DeviceMem::create(curUerMemInPtr, curSize);
90 0 : DeviceMem userMemOut = DeviceMem::create(curUerMemOutPtr, curSize);
91 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, userMemOut, userMemIn, stream_));
92 0 : }
93 0 : return HCCL_SUCCESS;
94 : }
95 :
96 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
97 :
98 0 : std::vector<Slice> inputSlices(slices_);
99 0 : if (slices_.size() == 0) {
100 : // slices_为空，临时构造等长slices
101 0 : slices_.resize(interRankSize_);
102 0 : inputSlices.resize(interRankSize_);
103 :
104 0 : for (u32 i = 0; i < interRankSize_; i++) {
105 0 : slices_[i].size = curSize;
106 0 : slices_[i].offset = (i * sliceSize);
107 :
108 0 : inputSlices[i].size = curSize;
109 0 : inputSlices[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
110 : }
111 : } else {
112 : // allgather_v场景下走else分支，每张卡的数据在CCLbuffer上偏移地址相同
113 0 : for(u32 i = 0; i < interRankSize_; i++) {
114 0 : inputSlices[i].offset = 0;
115 : }
116 : }
117 :
118 0 : for (u32 i = 0; i < interRankSize_; i++) {
119 0 : HCCL_DEBUG("[AllGatherMeshDirect][Slice]: rank[%u], outputslice: size[%llu] offset[%llu] "
120 : "inputslice: size[%llu] offset[%llu]",
121 : i, slices_[i].size, slices_[i].offset, inputSlices[i].size, inputSlices[i].offset);
122 : }
123 :
124 0 : DeviceMem src;
125 0 : DeviceMem dst;
126 0 : src = DeviceMem::create(curUerMemInPtr, inputSlices[rank].size);
127 0 : u64 localOffsetByte = inputSlices[rank].offset % HCCL_MIN_SLICE_ALIGN_910B;
128 0 : dst = DeviceMem::create(curCommMemOutPtr + localOffsetByte, inputSlices[rank].size);
129 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
130 :
131 0 : CHK_RET(MainRecordSub());
132 0 : CHK_RET(SubWaitMain());
133 :
134 0 : for (u32 round = 1; round < rankSize; round++) {
135 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
136 0 : Stream& subStream = meshStreams_[round - 1];
137 0 : CHK_RET(links[dstRank]->TxAck(subStream));
138 0 : CHK_RET(links[dstRank]->RxAck(subStream));
139 : }
140 :
141 0 : HCCL_DEBUG("[AllgatherMeshDirect]runAsync now rankSize is %u", rankSize);
142 0 : CHK_RET(SubRecordMain());
143 0 : CHK_RET(MainWaitSub());
144 :
145 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
146 :
147 0 : CHK_RET(SubWaitMain());
148 0 : CHK_RET(MainRecordSub());
149 :
150 0 : src = dst;
151 0 : dst = DeviceMem::create(curUerMemOutPtr + slices_[rank].offset, slices_[rank].size);
152 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
153 :
154 0 : for (u32 round = 1; round < rankSize; round++) {
155 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
156 0 : Stream& subStream = meshStreams_[round - 1];
157 : // 本rank要收数据
158 0 : void *remMemPtr = nullptr;
159 : // DMA消减场景，从对端的ccl out内存拿数据到本端的user out
160 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
161 0 : u64 remoteOffsetByte = inputSlices[dstRank].offset % HCCL_MIN_SLICE_ALIGN_910B;
162 0 : src = DeviceMem::create(static_cast<char *>(remMemPtr) + remoteOffsetByte, inputSlices[dstRank].size);
163 0 : dst = DeviceMem::create(curUerMemOutPtr + slices_[dstRank].offset, slices_[dstRank].size);
164 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStream,
165 : links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
166 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
167 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
168 : }
169 0 : CHK_RET(SubRecordMain());
170 0 : CHK_RET(MainWaitSub());
171 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem, emptyMem, stream_));
172 :
173 0 : HCCL_INFO("AllGatherMeshDirect finished: rank[%u]", rank);
174 0 : return HCCL_SUCCESS;
175 0 : }
176 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_MESH_DIRECT, AllgatherMeshDirect);
177 : } // namespace hccl
|
