tvm
instruction.h
Go to the documentation of this file.
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/ffi/reflection/registry.h>
23 
24 #include <utility>
25 
26 namespace tvm {
27 
28 // Forward declaration
29 template <typename, typename>
31 
32 namespace tir {
33 
34 // Forward declaration
35 class Schedule;
36 
45 using FInstructionApply = ffi::TypedFunction<Array<Any>(
46  Schedule sch, const Array<Any>& inputs, const Array<Any>& attrs, const Any& decision)>;
47 
57  ffi::TypedFunction<String(const Array<Any>& inputs, const Array<Any>& attrs,
58  const Any& decision, const Array<String>& outputs)>;
59 
66 using FInstructionAttrsAsJSON = ffi::TypedFunction<ObjectRef(Array<Any> attrs)>;
67 
74 using FInstructionAttrsFromJSON = ffi::TypedFunction<Array<Any>(ObjectRef json_attrs)>;
75 
88 class InstructionKindNode : public runtime::Object {
89  public:
91  String name;
98  bool is_pure{false};
113 
114  static void RegisterReflection() {
115  namespace refl = tvm::ffi::reflection;
116  refl::ObjectDef<InstructionKindNode>()
117  .def_ro("name", &InstructionKindNode::name)
118  .def_ro("_is_pure", &InstructionKindNode::is_pure);
119  }
120 
122  bool IsPostproc() const;
123 
124  static constexpr const char* _type_key = "tir.InstructionKind";
126 };
127 
132 class InstructionKind : public runtime::ObjectRef {
133  public:
139  static InstructionKind Get(const String& name);
141 };
142 
144 class InstructionNode : public runtime::Object {
145  public:
159  Array<Any> inputs;
165  Array<Any> attrs;
172  Array<Any> outputs;
173 
174  static void RegisterReflection() {
175  namespace refl = tvm::ffi::reflection;
176  refl::ObjectDef<InstructionNode>()
177  .def_ro("kind", &InstructionNode::kind)
178  .def_ro("inputs", &InstructionNode::inputs)
179  .def_ro("attrs", &InstructionNode::attrs)
180  .def_ro("outputs", &InstructionNode::outputs);
181  }
182 
183  static constexpr const char* _type_key = "tir.Instruction";
185 };
186 
191 class Instruction : public runtime::ObjectRef {
192  public:
200  explicit Instruction(InstructionKind kind, Array<Any> inputs, Array<Any> attrs,
201  Array<Any> outputs);
202 
204 };
205 
211 #define TVM_INST_KIND_REGISTER_VAR_DEF \
212  static DMLC_ATTRIBUTE_UNUSED ::tvm::tir::InstructionKindRegEntry& __make_##InstructionKind
213 
231 #define TVM_REGISTER_INST_KIND(InstructionKindName) \
232  TVM_STR_CONCAT(TVM_INST_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
233  ::tvm::tir::InstructionKindRegEntry::RegisterOrGet(InstructionKindName).set_name()
234 
237  public:
238  static InstructionKindRegEntry& RegisterOrGet(const String& name);
239 
241  get_mutable()->name = this->name;
242  return *this;
243  }
244 
246  get_mutable()->is_pure = is_pure;
247  return *this;
248  }
249 
251  get_mutable()->f_apply_to_schedule = std::move(f_apply_to_schedule);
252  return *this;
253  }
254 
256  get_mutable()->f_as_python = std::move(f_as_python);
257  return *this;
258  }
259 
261  get_mutable()->f_attrs_as_json = std::move(f_attrs_as_json);
262  return *this;
263  }
264 
266  get_mutable()->f_attrs_from_json = std::move(f_attrs_from_json);
267  return *this;
268  }
269 
270  private:
272  explicit InstructionKindRegEntry(uint32_t reg_index);
274  InstructionKindNode* get_mutable() const {
275  return const_cast<InstructionKindNode*>(inst_kind_.get());
276  }
277 
279  String name;
281  InstructionKind inst_kind_;
282  template <typename, typename>
283  friend class ::tvm::AttrRegistry;
284  friend class InstructionKind;
285 };
286 
287 } // namespace tir
288 } // namespace tvm
289 
290 #endif // TVM_TIR_SCHEDULE_INSTRUCTION_H_
Definition: instruction.h:30
Kind of an instruction, e.g. Split, Reorder, etc. Besides the name, every kind of instruction has its...
Definition: instruction.h:88
String name
The name of a kind of instructions.
Definition: instruction.h:91
TVM_DECLARE_FINAL_OBJECT_INFO(InstructionKindNode, runtime::Object)
FInstructionApply f_apply_to_schedule
A functor that applies the instruction to a TensorIR schedule.
Definition: instruction.h:100
bool IsPostproc() const
Checks if the instruction kind is EnterPostproc.
FInstructionAttrsAsJSON f_attrs_as_json
A functor that serialize its attributes to JSON.
Definition: instruction.h:107
static constexpr const char * _type_key
Definition: instruction.h:124
static void RegisterReflection()
Definition: instruction.h:114
FInstructionAttrsFromJSON f_attrs_from_json
A functor that deserialize its attributes from JSON.
Definition: instruction.h:112
bool is_pure
Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule state....
Definition: instruction.h:98
FInstructionAsPython f_as_python
A functor that converts the instruction to a statement in python syntax.
Definition: instruction.h:102
An entry in the registry of InstructionKind.
Definition: instruction.h:236
InstructionKindRegEntry & set_attrs_from_json(FInstructionAttrsFromJSON f_attrs_from_json)
Definition: instruction.h:265
InstructionKindRegEntry & set_attrs_as_json(FInstructionAttrsAsJSON f_attrs_as_json)
Definition: instruction.h:260
InstructionKindRegEntry & set_apply_to_schedule(FInstructionApply f_apply_to_schedule)
Definition: instruction.h:250
static InstructionKindRegEntry & RegisterOrGet(const String &name)
InstructionKindRegEntry & set_name()
Definition: instruction.h:240
InstructionKindRegEntry & set_as_python(FInstructionAsPython f_as_python)
Definition: instruction.h:255
InstructionKindRegEntry & set_is_pure(bool is_pure)
Definition: instruction.h:245
friend class InstructionKind
Definition: instruction.h:284
Managed reference to InstructionKindNode.
Definition: instruction.h:132
static InstructionKind Get(const String &name)
Retrieve an InstructionKind using its name.
TVM_DEFINE_OBJECT_REF_METHODS(InstructionKind, runtime::ObjectRef, InstructionKindNode)
Schedule instructions each corresponds to a schedule primitive.
Definition: instruction.h:144
Array< Any > inputs
The input random variables of the instruction, and the type of each element can be one of the followi...
Definition: instruction.h:159
Array< Any > attrs
The attributes of the instruction. Similar to attributes of an operator, attributes of an instruction...
Definition: instruction.h:165
InstructionKind kind
The kind of the instruction.
Definition: instruction.h:147
static constexpr const char * _type_key
Definition: instruction.h:183
TVM_DECLARE_FINAL_OBJECT_INFO(InstructionNode, runtime::Object)
Array< Any > outputs
The output random variables of the instruction, and the type of each element can be one of the follow...
Definition: instruction.h:172
static void RegisterReflection()
Definition: instruction.h:174
Managed reference to InstructionNode.
Definition: instruction.h:191
Instruction(InstructionKind kind, Array< Any > inputs, Array< Any > attrs, Array< Any > outputs)
Constructor.
TVM_DEFINE_OBJECT_REF_METHODS(Instruction, runtime::ObjectRef, InstructionNode)
Managed reference to ScheduleNode.
Definition: schedule.h:880
Definition: repr_printer.h:91
ffi::TypedFunction< Array< Any >(ObjectRef json_attrs)> FInstructionAttrsFromJSON
Type of the functor that deserialize its attributes from JSON.
Definition: instruction.h:74
ffi::TypedFunction< ObjectRef(Array< Any > attrs)> FInstructionAttrsAsJSON
Type of the functor that serialize its attributes to JSON.
Definition: instruction.h:66
ffi::TypedFunction< Array< Any >(Schedule sch, const Array< Any > &inputs, const Array< Any > &attrs, const Any &decision)> FInstructionApply
Type of the functor that applies the instruction to a TensorIR schedule.
Definition: instruction.h:46
ffi::TypedFunction< String(const Array< Any > &inputs, const Array< Any > &attrs, const Any &decision, const Array< String > &outputs)> FInstructionAsPython
Type of the functor that converts the instruction to a statement in python syntax.
Definition: instruction.h:58
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37