tvm
iter_affine_map.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 
48 #ifndef TVM_ARITH_ITER_AFFINE_MAP_H_
49 #define TVM_ARITH_ITER_AFFINE_MAP_H_
50 
51 #include <tvm/arith/analyzer.h>
52 #include <tvm/ffi/reflection/registry.h>
53 #include <tvm/ir/diagnostic.h>
54 #include <tvm/ir/expr.h>
55 #include <tvm/tir/var.h>
56 
57 namespace tvm {
58 namespace arith {
59 
67 class IterMapExprNode : public PrimExprNode {
68  public:
69  static constexpr const char* _type_key = "arith.IterMapExpr";
70  static constexpr const uint32_t _type_child_slots = 2;
72 };
73 
78 class IterMapExpr : public PrimExpr {
79  public:
81 };
82 
89 class IterMarkNode : public Object {
90  public:
100 
101  static void RegisterReflection() {
102  namespace refl = tvm::ffi::reflection;
103  refl::ObjectDef<IterMarkNode>()
104  .def_ro("source", &IterMarkNode::source)
105  .def_ro("extent", &IterMarkNode::extent);
106  }
107 
108  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
109 
110  static constexpr const char* _type_key = "arith.IterMark";
112 };
113 
118 class IterMark : public ObjectRef {
119  public:
125  TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
126 
129 };
130 
137  public:
146 
147  static void RegisterReflection() {
148  namespace refl = tvm::ffi::reflection;
149  refl::ObjectDef<IterSplitExprNode>()
150  .def_ro("source", &IterSplitExprNode::source)
151  .def_ro("lower_factor", &IterSplitExprNode::lower_factor)
152  .def_ro("extent", &IterSplitExprNode::extent)
153  .def_ro("scale", &IterSplitExprNode::scale);
154  }
155 
156  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
157  static constexpr const char* _type_key = "arith.IterSplitExpr";
159 };
160 
165 class IterSplitExpr : public IterMapExpr {
166  public:
171  TVM_DLL explicit IterSplitExpr(IterMark source);
177  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
185  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
186  PrimExpr scale);
187 
190 };
191 
198  public:
200  Array<IterSplitExpr> args;
203 
204  static void RegisterReflection() {
205  namespace refl = tvm::ffi::reflection;
206  refl::ObjectDef<IterSumExprNode>()
207  .def_ro("args", &IterSumExprNode::args)
208  .def_ro("base", &IterSumExprNode::base);
209  }
210 
211  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
212  static constexpr const char* _type_key = "arith.IterSumExpr";
214 };
215 
220 class IterSumExpr : public IterMapExpr {
221  public:
227  TVM_DLL IterSumExpr(Array<IterSplitExpr> args, PrimExpr base);
228 
231 };
232 
235  // Require the mapping to be bijective.
237  // Require the mapping to be surjective.
239  // No mapping safety check.
240  NoCheck = 3
241 };
242 
246 class IterMapResultNode : public Object {
247  public:
248  // The detected pattern if a match exists.
249  Array<IterSumExpr> indices;
250 
251  // Any errors that occurred while converting the input indices. If
252  // the array is empty, the conversion was successful.
253  Array<String> errors;
254 
264 
265  static void RegisterReflection() {
266  namespace refl = tvm::ffi::reflection;
267  refl::ObjectDef<IterMapResultNode>()
268  .def_ro("indices", &IterMapResultNode::indices)
269  .def_ro("errors", &IterMapResultNode::errors)
270  .def_ro("padding_predicate", &IterMapResultNode::padding_predicate);
271  }
272 
273  static constexpr const char* _type_key = "arith.IterMapResult";
275 };
276 
281 class IterMapResult : public ObjectRef {
282  public:
283  // constructor
284  IterMapResult() { data_ = make_object<IterMapResultNode>(); }
285 
287  IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
288 };
289 
313 IterMapResult DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
314  const PrimExpr& predicate, IterMapLevel check_level,
315  arith::Analyzer* analyzer, bool simplify_trivial_iterators = true);
316 
328 Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
329  const PrimExpr& input_pred, IterMapLevel check_level,
330  arith::Analyzer* analyzer, bool simplify_trivial_iterators = true);
331 
352 Map<Var, PrimExpr> InverseAffineIterMap(const Array<IterSumExpr>& iter_map,
353  const Array<PrimExpr> outputs);
354 
382 Array<Array<IterMark>> SubspaceDivide(const Array<PrimExpr>& bindings,
383  const Map<Var, Range>& input_iters,
384  const Array<Var>& sub_iters, const PrimExpr& predicate,
385  IterMapLevel check_level, arith::Analyzer* analyzer,
386  bool simplify_trivial_iterators = true);
387 
394 
411 IterSumExpr NormalizeToIterSum(PrimExpr index, const Map<Var, Range>& input_iters,
412  arith::Analyzer* analyzer);
413 
414 } // namespace arith
415 } // namespace tvm
416 #endif // TVM_ARITH_ITER_AFFINE_MAP_H_
Algebra expression simplifications.
Base node of all primitive expressions.
Definition: expr.h:95
Reference to PrimExprNode.
Definition: expr.h:129
Analyzer that contains bunch of sub-analyzers.
Definition: analyzer.h:636
Base class of all iter map expressions.
Definition: iter_affine_map.h:67
TVM_DECLARE_BASE_OBJECT_INFO(IterMapExprNode, PrimExprNode)
static constexpr const char * _type_key
Definition: iter_affine_map.h:69
static constexpr const uint32_t _type_child_slots
Definition: iter_affine_map.h:70
Managed reference to IterMapExprNode.
Definition: iter_affine_map.h:78
TVM_DEFINE_OBJECT_REF_METHODS(IterMapExpr, PrimExpr, IterMapExprNode)
Result of DetectIterMap.
Definition: iter_affine_map.h:246
Array< IterSumExpr > indices
Definition: iter_affine_map.h:249
static constexpr const char * _type_key
Definition: iter_affine_map.h:273
TVM_DECLARE_FINAL_OBJECT_INFO(IterMapResultNode, Object)
Array< String > errors
Definition: iter_affine_map.h:253
static void RegisterReflection()
Definition: iter_affine_map.h:265
PrimExpr padding_predicate
Boolean expression indicating if a specific value w.
Definition: iter_affine_map.h:263
Managed reference to IterMapResultNode.
Definition: iter_affine_map.h:281
IterMapResult()
Definition: iter_affine_map.h:284
IterMapResultNode * operator->() const
Definition: iter_affine_map.h:287
Mark the source as an iterator in [0, extent).
Definition: iter_affine_map.h:89
TVM_DECLARE_FINAL_OBJECT_INFO(IterMarkNode, Object)
PrimExpr extent
The extent of the iteration.
Definition: iter_affine_map.h:99
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:108
PrimExpr source
The source expression, can either be a IterSumExpr or a Var.
Definition: iter_affine_map.h:95
static void RegisterReflection()
Definition: iter_affine_map.h:101
static constexpr const char * _type_key
Definition: iter_affine_map.h:110
Managed reference to IterMarkExprNode.
Definition: iter_affine_map.h:118
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode)
IterMark(PrimExpr source, PrimExpr extent)
constructor.
TVM_DEFINE_OBJECT_REF_METHODS(IterMark, ObjectRef, IterMarkNode)
Split of an iterator.
Definition: iter_affine_map.h:136
static constexpr const char * _type_key
Definition: iter_affine_map.h:157
PrimExpr lower_factor
The lower factor to split the source.
Definition: iter_affine_map.h:141
TVM_DECLARE_FINAL_OBJECT_INFO(IterSplitExprNode, IterMapExprNode)
IterMark source
The source marked iterator.
Definition: iter_affine_map.h:139
PrimExpr scale
Additional scale.
Definition: iter_affine_map.h:145
static void RegisterReflection()
Definition: iter_affine_map.h:147
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:156
PrimExpr extent
The extent of the split.
Definition: iter_affine_map.h:143
Managed reference to IterSplitExprNode.
Definition: iter_affine_map.h:165
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode)
TVM_DEFINE_OBJECT_REF_METHODS(IterSplitExpr, IterMapExpr, IterSplitExprNode)
IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent, PrimExpr scale)
constructor
IterSplitExpr(IterMark source)
constructor from just source.
IterSplitExpr(IterMark source, PrimExpr scale)
constructor from just source.
Fuse multiple iterators by summing them with scaling.
Definition: iter_affine_map.h:197
static void RegisterReflection()
Definition: iter_affine_map.h:204
static constexpr const char * _type_key
Definition: iter_affine_map.h:212
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:211
Array< IterSplitExpr > args
The args to the sum.
Definition: iter_affine_map.h:200
TVM_DECLARE_FINAL_OBJECT_INFO(IterSumExprNode, IterMapExprNode)
PrimExpr base
The base offset.
Definition: iter_affine_map.h:202
Managed reference to IterSumExprNode.
Definition: iter_affine_map.h:220
IterSumExpr(Array< IterSplitExpr > args, PrimExpr base)
constructor.
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode)
TVM_DEFINE_OBJECT_REF_METHODS(IterSumExpr, IterMapExpr, IterSumExprNode)
A new diagnostic interface for TVM error reporting.
Base expr nodes in TVM.
Array< Array< IterMark > > SubspaceDivide(const Array< PrimExpr > &bindings, const Map< Var, Range > &input_iters, const Array< Var > &sub_iters, const PrimExpr &predicate, IterMapLevel check_level, arith::Analyzer *analyzer, bool simplify_trivial_iterators=true)
Detect if bindings can be written as [a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n].
IterMapLevel
Mapping level for iterators.
Definition: iter_affine_map.h:234
@ NoCheck
Definition: iter_affine_map.h:240
@ Bijective
Definition: iter_affine_map.h:236
@ Surjective
Definition: iter_affine_map.h:238
PrimExpr NormalizeIterMapToExpr(const PrimExpr &expr)
Given an expression that may contain IterMapExpr, transform it to normal PrimExpr.
Array< PrimExpr > IterMapSimplify(const Array< PrimExpr > &indices, const Map< Var, Range > &input_iters, const PrimExpr &input_pred, IterMapLevel check_level, arith::Analyzer *analyzer, bool simplify_trivial_iterators=true)
Use IterVarMap detector to rewrite and simplify the indices.
IterMapResult DetectIterMap(const Array< PrimExpr > &indices, const Map< Var, Range > &input_iters, const PrimExpr &predicate, IterMapLevel check_level, arith::Analyzer *analyzer, bool simplify_trivial_iterators=true)
Detect if indices can be written as [y_0 + c_0, y_1 + c_1, ..., y_n + c_n].
IterSumExpr NormalizeToIterSum(PrimExpr index, const Map< Var, Range > &input_iters, arith::Analyzer *analyzer)
Rewrite index as IterSumExpr.
Map< Var, PrimExpr > InverseAffineIterMap(const Array< IterSumExpr > &iter_map, const Array< PrimExpr > outputs)
Apply the inverse of the affine transformation to the outputs.
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Variables in the TIR.