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 
46  ffi::TypedFunction<ffi::Array<Any>(Schedule sch, const ffi::Array<Any>& inputs,
47  const ffi::Array<Any>& attrs, const Any& decision)>;
48 
58  ffi::TypedFunction<ffi::String(const ffi::Array<Any>& inputs, const ffi::Array<Any>& attrs,
59  const Any& decision, const ffi::Array<ffi::String>& outputs)>;
60 
67 using FInstructionAttrsAsJSON = ffi::TypedFunction<ObjectRef(ffi::Array<Any> attrs)>;
68 
75 using FInstructionAttrsFromJSON = ffi::TypedFunction<ffi::Array<Any>(ObjectRef json_attrs)>;
76 
89 class InstructionKindNode : public runtime::Object {
90  public:
92  ffi::String name;
99  bool is_pure{false};
114 
115  static void RegisterReflection() {
116  namespace refl = tvm::ffi::reflection;
117  refl::ObjectDef<InstructionKindNode>()
118  .def_ro("name", &InstructionKindNode::name)
119  .def_ro("_is_pure", &InstructionKindNode::is_pure);
120  }
121 
123  bool IsPostproc() const;
124  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.InstructionKind", InstructionKindNode, runtime::Object);
125 };
126 
131 class InstructionKind : public runtime::ObjectRef {
132  public:
138  static InstructionKind Get(const ffi::String& name);
141 };
142 
144 class InstructionNode : public runtime::Object {
145  public:
159  ffi::Array<Any> inputs;
165  ffi::Array<Any> attrs;
172  ffi::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  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Instruction", InstructionNode, runtime::Object);
183 };
184 
189 class Instruction : public runtime::ObjectRef {
190  public:
198  explicit Instruction(InstructionKind kind, ffi::Array<Any> inputs, ffi::Array<Any> attrs,
199  ffi::Array<Any> outputs);
200 
202 };
203 
209 #define TVM_INST_KIND_REGISTER_VAR_DEF \
210  static DMLC_ATTRIBUTE_UNUSED ::tvm::tir::InstructionKindRegEntry& __make_##InstructionKind
211 
229 #define TVM_REGISTER_INST_KIND(InstructionKindName) \
230  TVM_STR_CONCAT(TVM_INST_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
231  ::tvm::tir::InstructionKindRegEntry::RegisterOrGet(InstructionKindName).set_name()
232 
235  public:
236  static InstructionKindRegEntry& RegisterOrGet(const ffi::String& name);
237 
239  get_mutable()->name = this->name;
240  return *this;
241  }
242 
244  get_mutable()->is_pure = is_pure;
245  return *this;
246  }
247 
249  get_mutable()->f_apply_to_schedule = std::move(f_apply_to_schedule);
250  return *this;
251  }
252 
254  get_mutable()->f_as_python = std::move(f_as_python);
255  return *this;
256  }
257 
259  get_mutable()->f_attrs_as_json = std::move(f_attrs_as_json);
260  return *this;
261  }
262 
264  get_mutable()->f_attrs_from_json = std::move(f_attrs_from_json);
265  return *this;
266  }
267 
268  private:
270  explicit InstructionKindRegEntry(uint32_t reg_index);
272  InstructionKindNode* get_mutable() const {
273  return const_cast<InstructionKindNode*>(inst_kind_.get());
274  }
275 
277  ffi::String name;
279  InstructionKind inst_kind_;
280  template <typename, typename>
281  friend class ::tvm::AttrRegistry;
282  friend class InstructionKind;
283 };
284 
285 } // namespace tir
286 } // namespace tvm
287 
288 #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:89
ffi::String name
The name of a kind of instructions.
Definition: instruction.h:92
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.InstructionKind", InstructionKindNode, runtime::Object)
FInstructionApply f_apply_to_schedule
A functor that applies the instruction to a TensorIR schedule.
Definition: instruction.h:101
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:108
static void RegisterReflection()
Definition: instruction.h:115
FInstructionAttrsFromJSON f_attrs_from_json
A functor that deserialize its attributes from JSON.
Definition: instruction.h:113
bool is_pure
Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule state....
Definition: instruction.h:99
FInstructionAsPython f_as_python
A functor that converts the instruction to a statement in python syntax.
Definition: instruction.h:103
An entry in the registry of InstructionKind.
Definition: instruction.h:234
InstructionKindRegEntry & set_attrs_from_json(FInstructionAttrsFromJSON f_attrs_from_json)
Definition: instruction.h:263
InstructionKindRegEntry & set_attrs_as_json(FInstructionAttrsAsJSON f_attrs_as_json)
Definition: instruction.h:258
InstructionKindRegEntry & set_apply_to_schedule(FInstructionApply f_apply_to_schedule)
Definition: instruction.h:248
static InstructionKindRegEntry & RegisterOrGet(const ffi::String &name)
InstructionKindRegEntry & set_name()
Definition: instruction.h:238
InstructionKindRegEntry & set_as_python(FInstructionAsPython f_as_python)
Definition: instruction.h:253
InstructionKindRegEntry & set_is_pure(bool is_pure)
Definition: instruction.h:243
friend class InstructionKind
Definition: instruction.h:282
Managed reference to InstructionKindNode.
Definition: instruction.h:131
static InstructionKind Get(const ffi::String &name)
Retrieve an InstructionKind using its name.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(InstructionKind, runtime::ObjectRef, InstructionKindNode)
Schedule instructions each corresponds to a schedule primitive.
Definition: instruction.h:144
ffi::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
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Instruction", InstructionNode, runtime::Object)
ffi::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
ffi::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 void RegisterReflection()
Definition: instruction.h:174
Managed reference to InstructionNode.
Definition: instruction.h:189
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Instruction, runtime::ObjectRef, InstructionNode)
Instruction(InstructionKind kind, ffi::Array< Any > inputs, ffi::Array< Any > attrs, ffi::Array< Any > outputs)
Constructor.
Managed reference to ScheduleNode.
Definition: schedule.h:885
Definition: repr_printer.h:91
ffi::TypedFunction< ffi::String(const ffi::Array< Any > &inputs, const ffi::Array< Any > &attrs, const Any &decision, const ffi::Array< ffi::String > &outputs)> FInstructionAsPython
Type of the functor that converts the instruction to a statement in python syntax.
Definition: instruction.h:59
ffi::TypedFunction< ObjectRef(ffi::Array< Any > attrs)> FInstructionAttrsAsJSON
Type of the functor that serialize its attributes to JSON.
Definition: instruction.h:67
ffi::TypedFunction< ffi::Array< Any >(Schedule sch, const ffi::Array< Any > &inputs, const ffi::Array< Any > &attrs, const Any &decision)> FInstructionApply
Type of the functor that applies the instruction to a TensorIR schedule.
Definition: instruction.h:47
ffi::TypedFunction< ffi::Array< Any >(ObjectRef json_attrs)> FInstructionAttrsFromJSON
Type of the functor that deserialize its attributes from JSON.
Definition: instruction.h:75
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37