1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | #ifndef TVM_TIR_SCHEDULE_INSTRUCTION_H_ |
20 | #define TVM_TIR_SCHEDULE_INSTRUCTION_H_ |
21 | |
22 | #include <tvm/node/reflection.h> |
23 | |
24 | #include <utility> |
25 | |
26 | namespace tvm { |
27 | |
28 | // Forward declaration |
29 | template <typename, typename> |
30 | class AttrRegistry; |
31 | |
32 | namespace tir { |
33 | |
34 | // Forward declaration |
35 | class Schedule; |
36 | |
37 | /*! |
38 | * \brief Type of the functor that applies the instruction to a TensorIR schedule |
39 | * \param sch The schedule to be applied on |
40 | * \param inputs The input random variables |
41 | * \param attrs Instruction attributes |
42 | * \param decision Decisions made on the instruction |
43 | * \return The functor returns an array of output random variables |
44 | */ |
45 | using FInstructionApply = runtime::TypedPackedFunc<Array<ObjectRef>( |
46 | Schedule sch, const Array<ObjectRef>& inputs, const Array<ObjectRef>& attrs, |
47 | const Optional<ObjectRef>& decision)>; |
48 | |
49 | /*! |
50 | * \brief Type of the functor that converts the instruction to a statement in python syntax |
51 | * \param inputs Names of the input random variables |
52 | * \param attrs Instruction attributes |
53 | * \param decisions Decisions made on the instruction |
54 | * \param outputs Names of the output random variables |
55 | * \return A string representing the python api call |
56 | */ |
57 | using FInstructionAsPython = runtime::TypedPackedFunc<String( |
58 | const Array<ObjectRef>& inputs, const Array<ObjectRef>& attrs, |
59 | const Optional<ObjectRef>& decision, const Array<String>& outputs)>; |
60 | |
61 | /*! |
62 | * \brief Type of the functor that serialize its attributes to JSON |
63 | * \param attrs The attributes to be serialized |
64 | * \return An array, serialized attributes |
65 | * \note This functor is nullable |
66 | */ |
67 | using FInstructionAttrsAsJSON = runtime::TypedPackedFunc<ObjectRef(Array<ObjectRef> attrs)>; |
68 | |
69 | /*! |
70 | * \brief Type of the functor that deserialize its attributes from JSON |
71 | * \param json_attrs The attributes to be serialized |
72 | * \return An array, deserialized attributes |
73 | * \note This functor is nullable |
74 | */ |
75 | using FInstructionAttrsFromJSON = runtime::TypedPackedFunc<Array<ObjectRef>(ObjectRef json_attrs)>; |
76 | |
77 | /*! |
78 | * \brief Kind of an instruction, e.g. Split, Reorder, etc. |
79 | * Besides the name, every kind of instruction has its own properties, including: |
80 | * 1) A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state |
81 | * 2) A functor that applies the instruction to a TensorIR schedule |
82 | * 3) A functor that converts the instruction to a statement in python syntax |
83 | * 4) A functor that serialize its attributes to JSON |
84 | * 5) A functor that deserialize its attributes from JSON |
85 | * |
86 | * Unlike `tvm::OpNode`, `InstructionKindNode` doesn't support unstructured properties, |
87 | * mainly because there is no such usecase yet to add any other property. |
88 | */ |
89 | class InstructionKindNode : public runtime::Object { |
90 | public: |
91 | /*! \brief The name of a kind of instructions */ |
92 | String name; |
93 | /*! |
94 | * \brief Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule |
95 | * state. For example, the instruction `GetBlock` is pure because it changes |
96 | * nothing, while `ComputeInline` is not because removing it leads to a different resulting |
97 | * schedule. |
98 | */ |
99 | bool is_pure{false}; |
100 | /*! \brief A functor that applies the instruction to a TensorIR schedule */ |
101 | FInstructionApply f_apply_to_schedule{nullptr}; |
102 | /*! \brief A functor that converts the instruction to a statement in python syntax */ |
103 | FInstructionAsPython f_as_python{nullptr}; |
104 | /*! |
105 | * \brief A functor that serialize its attributes to JSON |
106 | * \note If the functor is null, it means no conversion is needed |
107 | */ |
108 | FInstructionAttrsAsJSON f_attrs_as_json{nullptr}; |
109 | /*! |
110 | * \brief A functor that deserialize its attributes from JSON |
111 | * \note If the functor is null, it means no conversion is needed |
112 | */ |
113 | FInstructionAttrsFromJSON f_attrs_from_json{nullptr}; |
114 | |
115 | void VisitAttrs(tvm::AttrVisitor* v) { |
116 | v->Visit("name" , &name); |
117 | v->Visit("_is_pure" , &is_pure); |
118 | // not visited: f_apply_to_schedule |
119 | // not visited: f_as_python |
120 | // not visited: f_attrs_as_json |
121 | // not visited: f_attrs_from_json |
122 | } |
123 | |
124 | /*! \brief Checks if the instruction kind is EnterPostproc */ |
125 | bool IsPostproc() const; |
126 | |
127 | static constexpr const char* _type_key = "tir.InstructionKind" ; |
128 | TVM_DECLARE_FINAL_OBJECT_INFO(InstructionKindNode, runtime::Object); |
129 | }; |
130 | |
131 | /*! |
132 | * \brief Managed reference to InstructionKindNode |
133 | * \sa InstructionKindNode |
134 | */ |
135 | class InstructionKind : public runtime::ObjectRef { |
136 | public: |
137 | /*! |
138 | * \brief Retrieve an InstructionKind using its name |
139 | * \param name The registered name of the InstructionKind |
140 | * \return The InstructionKind retrieved |
141 | */ |
142 | static InstructionKind Get(const String& name); |
143 | TVM_DEFINE_OBJECT_REF_METHODS(InstructionKind, runtime::ObjectRef, InstructionKindNode); |
144 | }; |
145 | |
146 | /*! \brief Schedule instructions each corresponds to a schedule primitive */ |
147 | class InstructionNode : public runtime::Object { |
148 | public: |
149 | /*! \brief The kind of the instruction */ |
150 | InstructionKind kind; |
151 | /*! |
152 | * \brief The input random variables of the instruction, and the type of each element can be one |
153 | * of the following: |
154 | * - BlockRV |
155 | * - LoopRV |
156 | * - ExprRV |
157 | * - FloatImm |
158 | * - IntImm |
159 | * - String |
160 | * - null pointer |
161 | */ |
162 | Array<ObjectRef> inputs; |
163 | /*! |
164 | * \brief The attributes of the instruction. Similar to attributes of an operator, |
165 | * attributes of an instruction are arbitrary constant metadata required by the instructions. |
166 | * For example, the name of the block to be retrieved in `GetBlock`. |
167 | */ |
168 | Array<ObjectRef> attrs; |
169 | /*! \brief The output random variables of the instruction, and the type of each element can be one |
170 | * of the following: |
171 | * - BlockRV |
172 | * - LoopRV |
173 | * - ExprRV, atomic variables only, won't be constants or composite PrimExpr |
174 | */ |
175 | Array<ObjectRef> outputs; |
176 | |
177 | void VisitAttrs(tvm::AttrVisitor* v) { |
178 | v->Visit("kind" , &kind); |
179 | v->Visit("inputs" , &inputs); |
180 | v->Visit("attrs" , &attrs); |
181 | v->Visit("outputs" , &outputs); |
182 | } |
183 | |
184 | static constexpr const char* _type_key = "tir.Instruction" ; |
185 | TVM_DECLARE_FINAL_OBJECT_INFO(InstructionNode, runtime::Object); |
186 | }; |
187 | |
188 | /*! |
189 | * \brief Managed reference to InstructionNode |
190 | * \sa InstructionNode |
191 | */ |
192 | class Instruction : public runtime::ObjectRef { |
193 | public: |
194 | /*! |
195 | * \brief Constructor |
196 | * \param kind The kind of the instruction |
197 | * \param inputs The input random variables of the instruction |
198 | * \param attrs The attributes of the instruction |
199 | * \param outputs The output random variables of the instruction |
200 | */ |
201 | explicit Instruction(InstructionKind kind, Array<ObjectRef> inputs, Array<ObjectRef> attrs, |
202 | Array<ObjectRef> outputs); |
203 | |
204 | TVM_DEFINE_OBJECT_REF_METHODS(Instruction, runtime::ObjectRef, InstructionNode); |
205 | }; |
206 | |
207 | /*! |
208 | * \brief A helper macro to register InstructionKind, only used in `TVM_REGISTER_INST_KIND` |
209 | * \note This macro is not user-facing. |
210 | * \sa TVM_REGISTER_INST_KIND |
211 | */ |
212 | #define TVM_INST_KIND_REGISTER_VAR_DEF \ |
213 | static DMLC_ATTRIBUTE_UNUSED ::tvm::tir::InstructionKindRegEntry& __make_##InstructionKind |
214 | |
215 | /*! |
216 | * \brief Register an InstructionKind |
217 | * \param InstructionKindName The name of the InstructionKind |
218 | * |
219 | * Example: |
220 | * |
221 | * \code |
222 | * |
223 | * TVM_REGISTER_INST_KIND("ComputeInline") |
224 | * .set_is_pure(false) |
225 | * .set_apply_to_schedule(ApplyToSchedule) |
226 | * .set_attrs_as_json(AttrsAsJSON) |
227 | * .set_attrs_from_json(AttrsFromJSON) |
228 | * .set_as_python(AsPython); |
229 | * |
230 | * \endcode |
231 | */ |
232 | #define TVM_REGISTER_INST_KIND(InstructionKindName) \ |
233 | TVM_STR_CONCAT(TVM_INST_KIND_REGISTER_VAR_DEF, __COUNTER__) = \ |
234 | ::tvm::tir::InstructionKindRegEntry::RegisterOrGet(InstructionKindName).set_name() |
235 | |
236 | /*! \brief An entry in the registry of InstructionKind */ |
237 | class InstructionKindRegEntry { |
238 | public: |
239 | static InstructionKindRegEntry& RegisterOrGet(const String& name); |
240 | |
241 | InstructionKindRegEntry& set_name() { |
242 | get_mutable()->name = this->name; |
243 | return *this; |
244 | } |
245 | |
246 | InstructionKindRegEntry& set_is_pure(bool is_pure) { |
247 | get_mutable()->is_pure = is_pure; |
248 | return *this; |
249 | } |
250 | |
251 | InstructionKindRegEntry& set_apply_to_schedule(FInstructionApply f_apply_to_schedule) { |
252 | get_mutable()->f_apply_to_schedule = std::move(f_apply_to_schedule); |
253 | return *this; |
254 | } |
255 | |
256 | InstructionKindRegEntry& set_as_python(FInstructionAsPython f_as_python) { |
257 | get_mutable()->f_as_python = std::move(f_as_python); |
258 | return *this; |
259 | } |
260 | |
261 | InstructionKindRegEntry& set_attrs_as_json(FInstructionAttrsAsJSON f_attrs_as_json) { |
262 | get_mutable()->f_attrs_as_json = std::move(f_attrs_as_json); |
263 | return *this; |
264 | } |
265 | |
266 | InstructionKindRegEntry& set_attrs_from_json(FInstructionAttrsFromJSON f_attrs_from_json) { |
267 | get_mutable()->f_attrs_from_json = std::move(f_attrs_from_json); |
268 | return *this; |
269 | } |
270 | |
271 | private: |
272 | /*! \brief Private constructor, used only by AttrRegistry */ |
273 | explicit InstructionKindRegEntry(uint32_t reg_index); |
274 | /*! \brief Get the mutable reference to the internal InstructionKind */ |
275 | InstructionKindNode* get_mutable() const { |
276 | return const_cast<InstructionKindNode*>(inst_kind_.get()); |
277 | } |
278 | |
279 | /*! \brief The name of the registry entry */ |
280 | String name; |
281 | /*! \brief The instruction kind */ |
282 | InstructionKind inst_kind_; |
283 | template <typename, typename> |
284 | friend class ::tvm::AttrRegistry; |
285 | friend class InstructionKind; |
286 | }; |
287 | |
288 | } // namespace tir |
289 | } // namespace tvm |
290 | |
291 | #endif // TVM_TIR_SCHEDULE_INSTRUCTION_H_ |
292 | |