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) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : AllgatherMeshDirect::~AllgatherMeshDirect() {}
18 :
19 0 : HcclResult AllgatherMeshDirect::Prepare(
20 : std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
21 : std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 userRank, HcomCollOpInfo* opInfo, u32 interRank,
22 : u32 interRankSize)
23 : {
24 0 : meshStreams_ = meshStreams;
25 0 : meshSignal_ = &meshSignal;
26 0 : meshSignalAux_ = &meshSignalAux;
27 0 : opInfo_ = opInfo;
28 0 : interRank_ = interRank;
29 0 : interRankSize_ = interRankSize;
30 0 : userRank_ = userRank;
31 0 : return HCCL_SUCCESS;
32 : }
33 :
34 0 : HcclResult AllgatherMeshDirect::MainRecordSub()
35 : {
36 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignalAux_).size(); signalIndex++) {
37 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex], profilerInput_.stage));
38 : }
39 0 : return HCCL_SUCCESS;
40 : }
41 :
42 0 : HcclResult AllgatherMeshDirect::SubWaitMain()
43 : {
44 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignalAux_).size(); streamIndex++) {
45 0 : CHK_RET(LocalNotify::Wait(
46 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
47 : }
48 0 : return HCCL_SUCCESS;
49 : }
50 :
51 0 : HcclResult AllgatherMeshDirect::MainWaitSub()
52 : {
53 0 : for (u32 signalIndex = 0; signalIndex < (*meshSignal_).size(); signalIndex++) {
54 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
55 : }
56 0 : return HCCL_SUCCESS;
57 : }
58 :
59 0 : HcclResult AllgatherMeshDirect::SubRecordMain()
60 : {
61 0 : for (u32 streamIndex = 0; streamIndex < (*meshSignal_).size(); streamIndex++) {
62 0 : CHK_RET(LocalNotify::Post(
63 : meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
64 : }
65 0 : return HCCL_SUCCESS;
66 : }
67 :
68 : // allgather的入口函数
69 0 : HcclResult AllgatherMeshDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
70 : {
71 0 : CHK_SMART_PTR_NULL(dispatcher_);
72 0 : CHK_PTR_NULL(stream_.ptr());
73 0 : HCCL_INFO(
74 : "AllGatherMeshDirect run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
75 : inputMem_.ptr(), outputMem_.ptr(), count_);
76 :
77 0 : char* curUerMemInPtr = static_cast<char*>(opInfo_->inputAddr);
78 0 : char* curUerMemOutPtr = static_cast<char*>(opInfo_->outputAddr);
79 0 : char* curCommMemOutPtr = static_cast<char*>(outputMem_.ptr());
80 :
81 0 : u32 unitSize = DataUnitSize(dataType_);
82 0 : u64 curSize = count_ * unitSize; // 当前count
83 0 : u64 sliceSize = opInfo_->count * unitSize; // 总输入count
84 :
85 0 : if (rankSize == 1) {
86 0 : if (opInfo_->inputAddr != opInfo_->outputAddr) {
87 0 : HCCL_DEBUG("rank[%u] mem copy async from input to output", rank);
88 0 : DeviceMem userMemIn = DeviceMem::create(curUerMemInPtr, curSize);
89 0 : DeviceMem userMemOut = DeviceMem::create(curUerMemOutPtr, curSize);
90 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, userMemOut, userMemIn, stream_));
91 0 : }
92 0 : return HCCL_SUCCESS;
93 : }
94 :
95 0 : DeviceMem emptyMem = outputMem_.range(0, 0);
96 :
97 0 : std::vector<Slice> inputSlices(slices_);
98 0 : if (slices_.size() == 0) {
99 : // slices_为空，临时构造等长slices
100 0 : slices_.resize(interRankSize_);
101 0 : inputSlices.resize(interRankSize_);
102 :
103 0 : for (u32 i = 0; i < interRankSize_; i++) {
104 0 : slices_[i].size = curSize;
105 0 : slices_[i].offset = (i * sliceSize);
106 :
107 0 : inputSlices[i].size = curSize;
108 0 : inputSlices[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
109 : }
110 : } else {
111 : // allgather_v场景下走else分支，每张卡的数据在CCLbuffer上偏移地址相同
112 0 : for (u32 i = 0; i < interRankSize_; i++) {
113 0 : inputSlices[i].offset = 0;
114 : }
115 : }
116 :
117 0 : for (u32 i = 0; i < interRankSize_; i++) {
118 0 : HCCL_DEBUG(
119 : "[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(
165 : dispatcher_, dst, src, subStream, 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
|
