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
26namespace tvm {
27
28// Forward declaration
29template <typename, typename>
30class AttrRegistry;
31
32namespace tir {
33
34// Forward declaration
35class 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 */
45using 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 */
57using 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 */
67using 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 */
75using 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 */
89class 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 */
135class 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 */
147class 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 */
192class 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 */
237class 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