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 : #ifndef HCCLV2_CCU_INSTRUCTION_SCATTER_MESH_2D_H_
12 : #define HCCLV2_CCU_INSTRUCTION_SCATTER_MESH_2D_H_
13 :
14 : #include "template_utils.h"
15 : #include "instruction.h"
16 : #include "ins_queue.h"
17 : #include "ccu_context_utils.h"
18 : #include "ccu_ctx_signature.h"
19 : #include "ccu_ins.h"
20 : #include "ccu_rank_group.h"
21 :
22 : namespace Hccl {
23 :
24 : // 为ScatterMesh2D实现的CCUIns、CCUCtxArg与CCUTaskArg
25 : class CcuCtxArgScatterMesh2D : public CcuCtxArg {
26 : public:
27 0 : explicit CcuCtxArgScatterMesh2D(
28 : const std::vector<uint64_t>& dSize, uint32_t rankSize, uint32_t rId, uint32_t axisId, uint32_t root,
29 : const CollAlgOperator& op, const std::vector<std::vector<RankId>>& tempVTopo)
30 0 : : dimSize_(dSize),
31 0 : rankSize_(rankSize),
32 0 : rankId_(rId),
33 0 : axisId_(axisId),
34 0 : root_(root),
35 0 : op_(op),
36 0 : tempVTopo_(tempVTopo)
37 0 : {}
38 :
39 0 : CcuCtxSignature GetCtxSignature() const override
40 : {
41 0 : CcuCtxSignature signature;
42 0 : GenerateCcuCtxSignature(signature, CcuInstType::CCU_SCATTER_MESH_2D_DIRECT, op_, tempVTopo_);
43 0 : return signature;
44 0 : }
45 :
46 : std::vector<uint64_t> dimSize_;
47 : uint32_t rankSize_;
48 : uint32_t rankId_;
49 : uint32_t axisId_;
50 : uint32_t root_;
51 : CollAlgOperator op_;
52 : std::vector<std::vector<RankId>> tempVTopo_;
53 : };
54 :
55 : class CcuTaskArgScatterMesh2D : public CcuTaskArg {
56 : public:
57 0 : explicit CcuTaskArgScatterMesh2D(
58 : uint64_t inputAddr, uint64_t outputAddr, uint64_t scratchAddr, uint64_t token, uint64_t sliceSize,
59 : uint64_t stride, uint64_t xSliceSize, uint64_t ySliceSize)
60 0 : : inputAddr_(inputAddr),
61 0 : outputAddr_(outputAddr),
62 0 : scratchAddr_(scratchAddr),
63 0 : token_(token),
64 0 : sliceSize_(sliceSize),
65 0 : stride_(stride),
66 0 : xSliceSize_(xSliceSize),
67 0 : ySliceSize_(ySliceSize)
68 0 : {}
69 :
70 : uint64_t inputAddr_;
71 : uint64_t outputAddr_;
72 : uint64_t scratchAddr_;
73 : uint64_t token_;
74 : uint64_t sliceSize_;
75 : uint64_t stride_;
76 : uint64_t xSliceSize_;
77 : uint64_t ySliceSize_;
78 : };
79 :
80 : class CcuInstructionScatterMesh2D : public CcuInstruction {
81 : public:
82 0 : CcuInstructionScatterMesh2D() : CcuInstruction() {}
83 :
84 0 : void Init(
85 : uint32_t rankId, uint32_t rankSize, uint32_t axisId, uint64_t inputAddr, uint64_t outputAddr,
86 : uint64_t scratchAddr, uint64_t token, uint64_t sliceSize, uint64_t stride, uint64_t xSliceSize,
87 : uint64_t ySliceSize, CollAlgOperator& op, std::vector<std::vector<RankId>>& tempVTopo)
88 : {
89 0 : dimSize_.push_back(tempVTopo[0].size());
90 0 : dimSize_.push_back(tempVTopo[1].size());
91 0 : rankId_ = rankId;
92 0 : rankSize_ = rankSize;
93 0 : axisId_ = axisId;
94 0 : inputAddr_ = inputAddr;
95 0 : outputAddr_ = outputAddr;
96 0 : scratchAddr_ = scratchAddr;
97 0 : sliceSize_ = sliceSize;
98 0 : token_ = token;
99 0 : stride_ = stride;
100 0 : xSliceSize_ = xSliceSize;
101 0 : ySliceSize_ = ySliceSize;
102 0 : op_ = op;
103 0 : tempVTopo_ = tempVTopo;
104 :
105 0 : if (dimSize_[0] * dimSize_[1] != rankSize_) {
106 0 : THROW<InvalidParamsException>(StringFormat(
107 : "[CcuInstructionScatterMesh2D] each DimSize and rankSize is NOT match. dimSize[0][%llu], "
108 : "dimSize[1][%llu], rankSize[%llu]",
109 0 : dimSize_[0], dimSize_[1], rankSize_));
110 : }
111 :
112 0 : return;
113 : }
114 :
115 0 : CcuInstType GetInstType() const override
116 : {
117 0 : HCCL_INFO("CcuInstructionScatterMesh2D instype is CCU_SCATTER_MESH_2D_DIRECT.");
118 0 : return instType_;
119 : }
120 :
121 : void SetInstType(CcuInstType instType) { instType_ = instType; }
122 :
123 0 : std::string Describe() const override
124 : {
125 : return StringFormat(
126 0 : "CcuInstructionScatterMesh2D rankId [%u], instType[%s]", rankId_, instType_.Describe().c_str());
127 : }
128 :
129 0 : std::unique_ptr<CcuCtxArg> GetCtxArg() const override
130 : {
131 0 : return std::make_unique<CcuCtxArgScatterMesh2D>(
132 0 : dimSize_, rankSize_, rankId_, axisId_, op_.root, op_, tempVTopo_);
133 : }
134 :
135 0 : std::unique_ptr<CcuTaskArg> GetTaskArg() const override
136 : {
137 0 : return std::make_unique<CcuTaskArgScatterMesh2D>(
138 0 : inputAddr_, outputAddr_, scratchAddr_, token_, sliceSize_, stride_, xSliceSize_, ySliceSize_);
139 : }
140 :
141 : private:
142 : CcuInstType instType_ = CcuInstType::CCU_SCATTER_MESH_2D_DIRECT;
143 : std::vector<uint64_t> dimSize_;
144 : uint32_t rankId_{0};
145 : uint32_t rankSize_{0};
146 : uint32_t axisId_{0};
147 : uint64_t inputAddr_{0};
148 : uint64_t outputAddr_{0};
149 : uint64_t scratchAddr_{0};
150 : uint64_t token_{0};
151 : uint64_t stride_{0};
152 : uint64_t sliceSize_{0};
153 : uint64_t xSliceSize_{0};
154 : uint64_t ySliceSize_{0};
155 : CollAlgOperator op_;
156 : std::vector<std::vector<RankId>> tempVTopo_;
157 : };
158 :
159 : } // namespace Hccl
160 : #endif // HCCLV2_CCU_INSTRUCTION_SCATTER_MESH_2D_H_
|