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_NHR_1D_H_
12 : #define HCCLV2_CCU_INSTRUCTION_SCATTER_NHR_1D_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 : using NHRStepInfo = struct NHRStepInfoDef {
24 : u32 step = 0;
25 : u32 myRank = 0;
26 : u32 nSlices;
27 : u32 toRank = 0;
28 : u32 fromRank = 0;
29 : std::vector<u32> txSliceIdxs;
30 : std::vector<u32> rxSliceIdxs;
31 0 : NHRStepInfoDef() : nSlices(0)
32 : {
33 0 : }
34 : };
35 :
36 : // 为ScatterNHR1D实现的CCUIns、CCUCtxArg与CCUTaskArg
37 : class CcuCtxArgScatterNHR1D : public CcuCtxArg {
38 : public:
39 0 : explicit CcuCtxArgScatterNHR1D(const std::vector<uint64_t> &dimSize, uint32_t rankId, uint32_t rootId,
40 : uint32_t axisId, uint32_t axisSize, const std::vector<NHRStepInfo> stepInfoVector,
41 : const std::map<u32, u32> indexMap, const CollAlgOperator &op,
42 : const std::vector<std::vector<RankId>> &tempVTopo)
43 0 : : dimSize_(dimSize), rankId_(rankId), rootId_(rootId), axisId_(axisId), axisSize_(axisSize),
44 0 : stepInfoVector_(stepInfoVector), indexMap_(indexMap), op_(op), tempVTopo_(tempVTopo)
45 : {
46 0 : }
47 0 : CcuCtxSignature GetCtxSignature() const override
48 : {
49 0 : CcuCtxSignature signature;
50 0 : GenerateCcuCtxSignature(signature, CcuInstType::CCU_SCATTER_NHR_1D_MEM2MEM, op_, tempVTopo_);
51 0 : return signature;
52 0 : }
53 : std::vector<uint64_t> dimSize_;
54 : uint64_t rankId_;
55 : uint64_t rootId_;
56 : uint64_t axisId_;
57 : uint64_t axisSize_;
58 : std::vector<NHRStepInfo> stepInfoVector_;
59 : std::map<u32, u32> indexMap_;
60 : CollAlgOperator op_;
61 : std::vector<std::vector<RankId>> tempVTopo_;
62 : };
63 :
64 : class CcuTaskArgScatterNHR1D : public CcuTaskArg {
65 : public:
66 0 : explicit CcuTaskArgScatterNHR1D(uint64_t inputAddr, uint64_t outputAddr, uint64_t scratchAddr, uint64_t token,
67 : uint64_t sliceSize, uint64_t die0Size, uint64_t die1Size, uint64_t inputSliceStride,
68 : uint64_t outputSliceStride, uint64_t inputRepeatStride, uint64_t outputRepeatStride,
69 : uint64_t repeatNum, uint64_t repeatNumVar, uint64_t isOutputScratch)
70 0 : : inputAddr_(inputAddr), outputAddr_(outputAddr), scratchAddr_(scratchAddr), token_(token),
71 0 : sliceSize_(sliceSize), die0Size_(die0Size), die1Size_(die1Size), inputSliceStride_(inputSliceStride),
72 0 : outputSliceStride_(outputSliceStride), inputRepeatStride_(inputRepeatStride),
73 0 : outputRepeatStride_(outputRepeatStride), repeatNum_(repeatNum), repeatNumVar_(repeatNumVar),
74 0 : isOutputScratch_(isOutputScratch)
75 : {
76 0 : }
77 :
78 : uint64_t inputAddr_;
79 : uint64_t outputAddr_;
80 : uint64_t scratchAddr_;
81 : uint64_t token_;
82 : uint64_t sliceSize_;
83 : uint64_t die0Size_;
84 : uint64_t die1Size_;
85 : uint64_t inputSliceStride_;
86 : uint64_t outputSliceStride_;
87 : uint64_t inputRepeatStride_;
88 : uint64_t outputRepeatStride_;
89 : uint64_t repeatNum_;
90 : uint64_t repeatNumVar_;
91 : uint64_t isOutputScratch_;
92 : };
93 :
94 : class CcuInstructionScatterNHR1D : public CcuInstruction {
95 : public:
96 0 : CcuInstructionScatterNHR1D() : CcuInstruction()
97 : {
98 0 : }
99 :
100 0 : void Init(uint32_t rankId, uint32_t rootId, uint64_t inputAddr, uint64_t outputAddr, uint64_t scratchAddr,
101 : uint64_t axisId, uint64_t axisSize, uint64_t sliceSize, uint64_t die0Size, uint64_t die1Size,
102 : uint64_t inputSliceStride, uint64_t outputSliceStride, uint64_t inputRepeatStride,
103 : uint64_t outputRepeatStride, uint32_t repeatNum, uint64_t repeatNumVar,
104 : std::vector<NHRStepInfo> stepInfoVector, std::map<u32, u32> indexMap, uint64_t token, CollAlgOperator &op,
105 : std::vector<std::vector<RankId>> &tempVTopo, uint64_t isOutputScratch)
106 : {
107 0 : u32 maxDimNum = 1;
108 0 : if (tempVTopo.size() != maxDimNum) {
109 0 : THROW<InvalidParamsException>(
110 0 : StringFormat("[CcuInstructionScatterNHR1D] tempVTopo size is not 1, size is [%zu].", tempVTopo.size()));
111 : }
112 0 : dimSize_.push_back(tempVTopo[0].size());
113 0 : rankId_ = rankId;
114 0 : rootId_ = rootId;
115 0 : token_ = token;
116 0 : inputAddr_ = inputAddr;
117 0 : outputAddr_ = outputAddr;
118 0 : scratchAddr_ = scratchAddr;
119 0 : axisId_ = axisId;
120 0 : axisSize_ = axisSize;
121 0 : sliceSize_ = sliceSize;
122 0 : die0Size_ = die0Size;
123 0 : die1Size_ = die1Size;
124 0 : inputSliceStride_ = inputSliceStride;
125 0 : outputSliceStride_ = outputSliceStride;
126 0 : inputRepeatStride_ = inputRepeatStride;
127 0 : outputRepeatStride_ = outputRepeatStride;
128 0 : stepInfoVector_ = stepInfoVector;
129 0 : indexMap_ = indexMap;
130 0 : token_ = token;
131 0 : op_ = op;
132 0 : tempVTopo_ = tempVTopo;
133 0 : repeatNum_ = repeatNum;
134 0 : repeatNumVar_ = repeatNumVar;
135 0 : isOutputScratch_ = isOutputScratch;
136 0 : return;
137 : }
138 :
139 0 : CcuInstType GetInstType() const override
140 : {
141 0 : HCCL_INFO("CcuInstructionScatterNHR1D instype is CCU_SCATTER_NHR_1D_MEM2MEM.");
142 0 : return instType_;
143 : }
144 :
145 0 : std::string Describe() const override
146 : {
147 0 : return StringFormat("CcuInstructionScatterNHR1D rankId [%u], instType[%s]", rankId_,
148 0 : instType_.Describe().c_str());
149 : }
150 :
151 : void SetInstType(CcuInstType instType)
152 : {
153 : instType_ = instType;
154 : }
155 :
156 0 : std::unique_ptr<CcuCtxArg> GetCtxArg() const override
157 : {
158 0 : return std::make_unique<CcuCtxArgScatterNHR1D>(dimSize_, rankId_, rootId_, axisId_, axisSize_, stepInfoVector_,
159 0 : indexMap_, op_, tempVTopo_);
160 : }
161 :
162 0 : std::unique_ptr<CcuTaskArg> GetTaskArg() const override
163 : {
164 0 : return std::make_unique<CcuTaskArgScatterNHR1D>(
165 0 : inputAddr_, outputAddr_, scratchAddr_, token_, sliceSize_, die0Size_, die1Size_, inputSliceStride_,
166 0 : outputSliceStride_, inputRepeatStride_, outputRepeatStride_, repeatNum_, repeatNumVar_,isOutputScratch_);
167 : }
168 :
169 : private:
170 : CcuInstType instType_ = CcuInstType::CCU_SCATTER_NHR_1D_MEM2MEM;
171 : std::vector<uint64_t> dimSize_;
172 : CollAlgOperator op_;
173 : uint32_t rankId_{0};
174 : uint32_t rootId_{0};
175 : uint32_t axisId_{0};
176 : uint32_t axisSize_{0};
177 : uint32_t sliceSize_{0};
178 : uint64_t inputAddr_{0};
179 : uint64_t outputAddr_{0};
180 : uint64_t scratchAddr_{0};
181 : uint64_t die0Size_{0};
182 : uint64_t die1Size_{0};
183 : uint64_t inputSliceStride_{0};
184 : uint64_t outputSliceStride_{0};
185 : uint64_t inputRepeatStride_{0};
186 : uint64_t outputRepeatStride_{0};
187 : uint32_t repeatNum_{0};
188 : uint64_t repeatNumVar_{0};
189 : uint64_t isOutputScratch_{0};
190 : std::vector<NHRStepInfo> stepInfoVector_;
191 : std::map<u32, u32> indexMap_;
192 : uint64_t token_{0};
193 : std::vector<std::vector<RankId>> tempVTopo_;
194 : };
195 : } // namespace Hccl
196 :
197 : #endif // HCCLV2_CCU_INSTRUCTION_SCATTER_NHR_1D_H_
|