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