tvm
var.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 
24 #ifndef TVM_TIR_VAR_H_
25 #define TVM_TIR_VAR_H_
26 
27 #include <tvm/ir/expr.h>
28 #include <tvm/node/node.h>
29 #include <tvm/runtime/data_type.h>
30 
31 #include <functional>
32 #include <string>
33 
34 namespace tvm {
35 namespace tir {
36 
48 class VarNode : public PrimExprNode {
49  public:
54  String name_hint;
63 
64  static void RegisterReflection() {
65  namespace refl = tvm::ffi::reflection;
66  refl::ObjectDef<VarNode>()
67  .def_ro("name", &VarNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
68  .def_ro("type_annotation", &VarNode::type_annotation);
69  }
70 
71  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
72  static constexpr const char* _type_key = "tir.Var";
73  static constexpr const uint32_t _type_child_slots = 1;
75 };
76 
78 class Var : public PrimExpr {
79  public:
80  explicit Var(ObjectPtr<Object> n) : PrimExpr(n) {}
87  TVM_DLL explicit Var(String name_hint = "v", DataType dtype = DataType::Int(32),
88  Span span = Span());
95  TVM_DLL explicit Var(String name_hint, Type type_annotation, Span span = Span());
101  TVM_DLL Var copy_with_name(const String& name) const;
107  TVM_DLL Var copy_with_suffix(const String& suffix) const;
114 
119  const VarNode* operator->() const { return get(); }
124  const VarNode* get() const { return static_cast<const VarNode*>(data_.get()); }
127 };
128 
133 class SizeVarNode : public VarNode {
134  public:
135  static void RegisterReflection() {
136  namespace refl = tvm::ffi::reflection;
137  refl::ObjectDef<SizeVarNode>();
138  }
139  static constexpr const char* _type_key = "tir.SizeVar";
141 };
142 
144 class SizeVar : public Var {
145  public:
146  explicit SizeVar(ObjectPtr<Object> n) : Var(n) {}
153  TVM_DLL explicit SizeVar(String name_hint = "s", DataType t = DataType::Int(32),
154  Span span = Span());
161  TVM_DLL explicit SizeVar(String name_hint, Type type_annotation, Span span = Span());
166  const SizeVarNode* operator->() const { return get(); }
171  const SizeVarNode* get() const { return static_cast<const SizeVarNode*>(data_.get()); }
174 };
175 
176 using Region = Array<Range>;
177 
185 enum IterVarType : int {
194  kDataPar = 0,
217  kOrdered = 3,
227  kOpaque = 4,
228  // The following are possible additional
229  // types that are provided during schedule
245  kTensorized = 8
246 };
247 
255  public:
269  String thread_tag;
274  mutable Span span;
275 
276  PrimExpr ToPrimExpr() const final { return var; }
277 
278  static void RegisterReflection() {
279  namespace refl = tvm::ffi::reflection;
280  refl::ObjectDef<IterVarNode>()
281  .def_ro("dom", &IterVarNode::dom)
282  .def_ro("var", &IterVarNode::var, refl::AttachFieldFlag::SEqHashDef())
283  .def_ro("iter_type", &IterVarNode::iter_type)
284  .def_ro("thread_tag", &IterVarNode::thread_tag);
285  }
286 
287  static constexpr const char* _type_key = "tir.IterVar";
288  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
290 };
291 
298 class IterVar : public PrimExprConvertible {
299  public:
300  TVM_DLL IterVar(Range dom, Var var, IterVarType iter_type, String thread_tag = "",
301  Span span = Span());
305  inline operator PrimExpr() const;
306 
309 };
310 
311 // inline implementations
312 inline IterVar::operator PrimExpr() const { return (*this)->var; }
313 
314 inline const char* IterVarType2String(IterVarType t) {
315  switch (t) {
316  case kDataPar:
317  return "DataPar";
318  case kThreadIndex:
319  return "ThreadIndex";
320  case kCommReduce:
321  return "CommReduce";
322  case kOrdered:
323  return "Ordered";
324  case kOpaque:
325  return "Opaque";
326  case kUnrolled:
327  return "Unrolled";
328  case kVectorized:
329  return "Vectorized";
330  case kParallelized:
331  return "Parallelized";
332  case kTensorized:
333  return "Tensorized";
334  }
335  return "Unknown";
336 }
337 } // namespace tir
338 } // namespace tvm
339 
340 /* \brief Allow tir.Var as key in STL tables
341  *
342  * For most TIR expressions, it would be ambiguous whether the
343  * expression should follow reference equality or structural equality.
344  * This is not the case for variables, which do not contain nested
345  * internal structure, and are frequently used as keys in lookup
346  * tables.
347  *
348  * Providing `std::hash` and `std::equal_to` specializations for
349  * `tir::Var` allows it to be used as a key in STL tables. For
350  * `PrimExpr`, the user must specify the type of equality used
351  * (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>` or
352  * `std::unordered_set<T, ObjectPtrHash, ObjectPtrEqual>`).
353  */
354 template <>
355 struct std::hash<tvm::tir::Var> {
356  std::size_t operator()(const tvm::tir::Var& var) const {
357  return tvm::runtime::ObjectPtrHash()(var);
358  }
359 };
360 
361 template <>
362 struct std::equal_to<tvm::tir::Var> {
363  bool operator()(const tvm::tir::Var& var_a, const tvm::tir::Var& var_b) const {
364  return tvm::runtime::ObjectPtrEqual()(var_a, var_b);
365  }
366 };
367 
368 #endif // TVM_TIR_VAR_H_
Base class for other IR constructs that can be converted to PrimExpr. This is useful for the FFI to c...
Definition: expr.h:159
Managed reference to PrimExprConvertibleNode.
Definition: expr.h:172
Base node of all primitive expressions.
Definition: expr.h:95
Reference to PrimExprNode.
Definition: expr.h:129
DataType dtype() const
Definition: expr.h:143
Range container
Definition: expr.h:698
Definition: source_map.h:113
Managed reference to TypeNode.
Definition: type.h:101
Runtime primitive data type.
Definition: data_type.h:47
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
An iteration variable representing an iteration over a one dimensional interval.
Definition: var.h:254
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: var.h:288
Var var
The looping variable.
Definition: var.h:262
String thread_tag
additional tag on the iteration variable, set this if this is bound already to a known thread tag.
Definition: var.h:269
static constexpr const char * _type_key
Definition: var.h:287
Span span
Span that points to the original source code. Reserved debug information.
Definition: var.h:274
static void RegisterReflection()
Definition: var.h:278
PrimExpr ToPrimExpr() const final
Definition: var.h:276
Range dom
the domain of iteration, if known, can be None For the intermediate schedule node,...
Definition: var.h:260
TVM_DECLARE_FINAL_OBJECT_INFO(IterVarNode, PrimExprConvertibleNode)
IterVarType iter_type
The type of the IterVar.
Definition: var.h:264
Iteration Variable, represents an iteration over an integer interval.
Definition: var.h:298
IterVar(Range dom, Var var, IterVarType iter_type, String thread_tag="", Span span=Span())
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterVarNode)
TVM_DEFINE_OBJECT_REF_METHODS(IterVar, PrimExprConvertible, IterVarNode)
A variable node represent a tensor index size, whose value must be non-negative.
Definition: var.h:133
static void RegisterReflection()
Definition: var.h:135
TVM_DECLARE_FINAL_OBJECT_INFO(SizeVarNode, VarNode)
static constexpr const char * _type_key
Definition: var.h:139
a named variable represents a tensor index size
Definition: var.h:144
SizeVar(String name_hint="s", DataType t=DataType::Int(32), Span span=Span())
constructor
const SizeVarNode * get() const
Get pointer to the internal value.
Definition: var.h:171
SizeVar(String name_hint, Type type_annotation, Span span=Span())
Constructor which provides a more detailed type annotation.
SizeVar(ObjectPtr< Object > n)
Definition: var.h:146
const SizeVarNode * operator->() const
Get pointer to the internal value.
Definition: var.h:166
A variable node in the IR.
Definition: var.h:48
static void RegisterReflection()
Definition: var.h:64
static constexpr const char * _type_key
Definition: var.h:72
TVM_DECLARE_BASE_OBJECT_INFO(VarNode, PrimExprNode)
Type type_annotation
type annotation of the variable.
Definition: var.h:62
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: var.h:71
static constexpr const uint32_t _type_child_slots
Definition: var.h:73
String name_hint
The hint to the variable name.
Definition: var.h:54
a named variable in TIR
Definition: var.h:78
const VarNode * get() const
Get pointer to the internal value.
Definition: var.h:124
Var(ObjectPtr< Object > n)
Definition: var.h:80
Var(String name_hint, Type type_annotation, Span span=Span())
Constructor which provides a more detailed type annotation.
const VarNode * operator->() const
Get pointer to the internal value.
Definition: var.h:119
Var copy_with_name(const String &name) const
Make a new copy of var with same type, but a different nam.
Var(String name_hint="v", DataType dtype=DataType::Int(32), Span span=Span())
Constructor.
Var copy_with_suffix(const String &suffix) const
Make a new copy of var with same type, append suffix.
Var copy_with_dtype(DataType dtype) const
Make a new copy of the variable with specified dtype.
Base expr nodes in TVM.
Definition: repr_printer.h:91
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
Array< Range > Region
Definition: var.h:176
IterVarType
Type of iteration variable. Each IterVar have a specific type.
Definition: var.h:185
@ kVectorized
The loop is vectorized.
Definition: var.h:237
@ kThreadIndex
The IterVar itself is a thread-index of a fixed thread launching group. Note that this is already ass...
Definition: var.h:202
@ kUnrolled
The execution is unrolled.
Definition: var.h:233
@ kTensorized
Marks boundary of tensorization intrinsic.
Definition: var.h:245
@ kDataPar
Data parallel iteration. This normally corresponds to axis of Tensor. Allow all IterVar manipulations...
Definition: var.h:194
@ kOrdered
Serial loops with loop carry dependency, the iteration must execute in order. Cannot be re-ordered.
Definition: var.h:217
@ kCommReduce
Communicative reduction. Cannot be directly parallelized.
Definition: var.h:209
@ kParallelized
The loop is parallelized.
Definition: var.h:241
@ kOpaque
IterVar is opaque,.
Definition: var.h:227
const char * IterVarType2String(IterVarType t)
Definition: var.h:314
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Definitions and helper macros for IR/AST nodes.