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 "scatter_mesh.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : ScatterMesh::ScatterMesh(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : ScatterMesh::~ScatterMesh() {}
18 :
19 0 : HcclResult ScatterMesh::Prepare(u32 interRank, u32 interRankSize)
20 : {
21 0 : interRank_ = interRank;
22 0 : interRankSize_ = interRankSize;
23 0 : return HCCL_SUCCESS;
24 : }
25 :
26 0 : void ScatterMesh::PrepareSlicesData(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
27 : {
28 0 : slices_.resize(rankSize);
29 0 : u64 sliceSize = (totalCount / rankSize) * unitSize;
30 :
31 0 : for (u32 i = 0; i < rankSize; i++) {
32 0 : slices_[i].offset = i * sliceSize;
33 0 : slices_[i].size = sliceSize;
34 0 : HCCL_DEBUG(" default slice[%u]: offset: [%llu] size[%llu]", i, i * sliceSize, sliceSize);
35 : }
36 0 : }
37 : // scatter的入口函数
38 0 : HcclResult ScatterMesh::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
39 : {
40 0 : CHK_SMART_PTR_NULL(dispatcher_);
41 0 : CHK_PTR_NULL(stream_.ptr());
42 0 : HCCL_INFO(
43 : "ScatterMesh run: rank[%u] rankSize[%u] inputMem[%p] to outputMem[%p] count[%llu]", rank, rankSize,
44 : inputMem_.ptr(), outputMem_.ptr(), count_);
45 : // ranksize ==1 的处理
46 0 : if (rankSize == 1) {
47 0 : if (inputMem_ != outputMem_) {
48 0 : HcclResult ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
49 0 : CHK_PRT_RET(
50 : ret != HCCL_SUCCESS,
51 : HCCL_ERROR(
52 : "[ScatterMesh][RunAsync]rank[%u] copy input[%p] to output[%p] "
53 : "failed",
54 : rank, inputMem_.ptr(), outputMem_.ptr()),
55 : ret);
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : if (links.size() < rankSize) {
61 0 : HCCL_ERROR(
62 : "[ScatterMesh][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(), rankSize);
63 0 : return HCCL_E_INTERNAL;
64 : }
65 :
66 0 : u32 unitSize = DataUnitSize(dataType_);
67 0 : if (unitSize == 0) {
68 0 : HCCL_ERROR("[ScatterMesh][RunAsync]rank[%u] unit data size is zero", rank);
69 0 : return HCCL_E_INTERNAL;
70 : }
71 0 : if (slices_.size() == 0) {
72 0 : PrepareSlicesData(unitSize, count_, rankSize);
73 : }
74 : // rank存放scatter 结果的偏移
75 0 : u64 scatterOffset = slices_[rank].offset;
76 0 : HCCL_DEBUG("rank[%u] scatter_offset is [%llu] reslutsize[%llu]", rank, scatterOffset, slices_[rank].size);
77 :
78 : // root rank向其他rank发送数据
79 0 : if (rank == root_) {
80 0 : for (u32 dstRank = 0; dstRank < rankSize; dstRank++) {
81 0 : HcclResult ret = RunSendScatter(dstRank, slices_[dstRank], links);
82 0 : CHK_PRT_RET(
83 : ret != HCCL_SUCCESS,
84 : HCCL_ERROR(
85 : "[ScatterMesh][RunAsync]root rank[%u] send scatter to "
86 : "dstrank[%u] run failed",
87 : rank, dstRank),
88 : ret);
89 : }
90 : } else { // 非rootrank 从rootrank接收数据
91 0 : CHK_RET(RunRecvScatter(root_, slices_[rank], links));
92 : }
93 :
94 : // Barrier确保数据收发完成
95 0 : if (barrierSwitchOn_) {
96 0 : for (u32 dstRank = 0; dstRank < rankSize; dstRank++) {
97 0 : if (dstRank != interRank_) {
98 0 : CHK_RET(ExecuteBarrier(links[dstRank], stream_));
99 : }
100 : }
101 : }
102 :
103 0 : HCCL_INFO("ScatterMesh finished: rank[%u]", rank);
104 0 : return HCCL_SUCCESS;
105 : }
106 :
107 0 : HcclResult ScatterMesh::RunSendScatter(const u32 dstRank, const Slice& slice, const std::vector<LINK>& links)
108 : {
109 0 : DeviceMem dst;
110 0 : DeviceMem src;
111 : // 本rank是root rank，直接进行数据的拷贝，从input拷贝到output
112 0 : if (dstRank == interRank_) {
113 0 : if (inputMem_ != outputMem_) {
114 : // root rank给自身拷贝时候不需要同步信号，拷贝到outputmem的偏移不同
115 0 : src = inputMem_.range(slices_[interRank_].offset, slices_[interRank_].size);
116 0 : dst = outputMem_.range(slices_[interRank_].offset, slices_[interRank_].size);
117 0 : HCCL_DEBUG(
118 : "root rank copy from input[%p] range[%llu] to output[%p] range[%llu], size[%llu]", src.ptr(),
119 : slices_[interRank_].offset, dst.ptr(), slices_[interRank_].offset, slices_[interRank_].size);
120 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_));
121 : }
122 : } else { // root rank给其他rank进行数据发送
123 0 : src = inputMem_.range(slice.offset, slice.size);
124 0 : if (dstRank >= links.size()) {
125 0 : HCCL_ERROR("[Run][SendScatter]DstRank[%u] is out of range, link Size[%llu]", dstRank, links.size());
126 0 : return HCCL_E_INTERNAL;
127 : }
128 :
129 0 : HCCL_DEBUG(
130 : "root rank tx input[%p] offset[%llu] to dstrank[%u] size[%llu] ", src.ptr(), slice.offset, dstRank,
131 : slice.size);
132 : // 接收目的rank的同步信号，便可进行下一轮发送
133 0 : CHK_RET(links[dstRank]->RxAck(stream_));
134 :
135 0 : HcclResult ret = links[dstRank]->TxAsync(
136 0 : UserMemType::OUTPUT_MEM, slice.offset + baseOffset_, src.ptr(), slice.size, stream_);
137 0 : CHK_PRT_RET(
138 : ret != HCCL_SUCCESS,
139 : HCCL_ERROR("[Run][SendScatter]root rank[%u] tx async to dstrank[%u] run failed", interRank_, dstRank), ret);
140 0 : ret = links[dstRank]->TxWaitDone(stream_);
141 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][SendScatter]TxWaitDone failed"), ret);
142 : }
143 0 : return HCCL_SUCCESS;
144 0 : }
145 : // 只有非root节点才接收数据
146 0 : HcclResult ScatterMesh::RunRecvScatter(const u32 srcRank, const Slice& slice, const std::vector<LINK>& links)
147 : {
148 0 : DeviceMem dst;
149 : // 判断数据是否需要分片
150 0 : if (srcRank >= links.size()) {
151 0 : HCCL_ERROR("[Run][RecvScatter]SrcRank[%u] is out of range, linkSize[%llu]", srcRank, links.size());
152 0 : return HCCL_E_INTERNAL;
153 : }
154 0 : dst = outputMem_.range(slice.offset, slice.size);
155 0 : HCCL_DEBUG("rank[%u] will rcv with output's offset[%llu], size[%llu] ", interRank_, slice.offset, slice.size);
156 :
157 : // 向root节点发送tx同步,rxmem可用
158 0 : HcclResult ret = links[srcRank]->TxAck(stream_);
159 0 : CHK_PRT_RET(
160 : ret != HCCL_SUCCESS, HCCL_ERROR("[Run][RecvScatter]rank[%u] tx ack to srcrank[%u] failed", interRank_, srcRank),
161 : ret);
162 :
163 0 : ret = links[srcRank]->RxAsync(UserMemType::INPUT_MEM, slice.offset + baseOffset_, dst.ptr(), slice.size, stream_);
164 0 : CHK_PRT_RET(
165 : ret != HCCL_SUCCESS,
166 : HCCL_ERROR("[Run][RecvScatter]rank[%u] rx async with output's offset[%llu] failed", interRank_, slice.offset),
167 : ret);
168 0 : ret = links[srcRank]->RxWaitDone(stream_);
169 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][SendScatter]RxWaitDone failed"), ret);
170 0 : return HCCL_SUCCESS;
171 0 : }
172 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_SCATTER_MESH, ScatterMesh);
173 : } // namespace hccl
|
