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 uint32_t _type_child_slots = 2;
71 };
72 
77 class IterMapExpr : public PrimExpr {
78  public:
80 };
81 
88 class IterMarkNode : public Object {
89  public:
99 
100  static void RegisterReflection() {
101  namespace refl = tvm::ffi::reflection;
102  refl::ObjectDef<IterMarkNode>()
103  .def_ro("source", &IterMarkNode::source)
104  .def_ro("extent", &IterMarkNode::extent);
105  }
106 
107  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
109 };
110 
115 class IterMark : public ObjectRef {
116  public:
122  TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
123 
126 };
127 
134  public:
143 
144  static void RegisterReflection() {
145  namespace refl = tvm::ffi::reflection;
146  refl::ObjectDef<IterSplitExprNode>()
147  .def_ro("source", &IterSplitExprNode::source)
148  .def_ro("lower_factor", &IterSplitExprNode::lower_factor)
149  .def_ro("extent", &IterSplitExprNode::extent)
150  .def_ro("scale", &IterSplitExprNode::scale);
151  }
152 
153  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
155 };
156 
161 class IterSplitExpr : public IterMapExpr {
162  public:
167  TVM_DLL explicit IterSplitExpr(IterMark source);
173  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
181  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
182  PrimExpr scale);
183 
186 };
187 
194  public:
196  ffi::Array<IterSplitExpr> args;
199 
200  static void RegisterReflection() {
201  namespace refl = tvm::ffi::reflection;
202  refl::ObjectDef<IterSumExprNode>()
203  .def_ro("args", &IterSumExprNode::args)
204  .def_ro("base", &IterSumExprNode::base);
205  }
206 
207  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
209 };
210 
215 class IterSumExpr : public IterMapExpr {
216  public:
222  TVM_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
223 
226 };
227 
230  // Require the mapping to be bijective.
232  // Require the mapping to be surjective.
234  // No mapping safety check.
235  NoCheck = 3
236 };
237 
241 class IterMapResultNode : public Object {
242  public:
243  // The detected pattern if a match exists.
244  ffi::Array<IterSumExpr> indices;
245 
246  // Any errors that occurred while converting the input indices. If
247  // the array is empty, the conversion was successful.
248  ffi::Array<ffi::String> errors;
249 
259 
260  static void RegisterReflection() {
261  namespace refl = tvm::ffi::reflection;
262  refl::ObjectDef<IterMapResultNode>()
263  .def_ro("indices", &IterMapResultNode::indices)
264  .def_ro("errors", &IterMapResultNode::errors)
265  .def_ro("padding_predicate", &IterMapResultNode::padding_predicate);
266  }
267  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, Object);
268 };
269 
274 class IterMapResult : public ObjectRef {
275  public:
276  // constructor
277  IterMapResult() { data_ = ffi::make_object<IterMapResultNode>(); }
278 
280  IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
281 };
282 
306 IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
307  const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
308  IterMapLevel check_level, arith::Analyzer* analyzer,
309  bool simplify_trivial_iterators = true);
310 
322 ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
323  const ffi::Map<Var, Range>& input_iters,
324  const PrimExpr& input_pred, IterMapLevel check_level,
325  arith::Analyzer* analyzer,
326  bool simplify_trivial_iterators = true);
327 
348 ffi::Map<Var, PrimExpr> InverseAffineIterMap(const ffi::Array<IterSumExpr>& iter_map,
349  const ffi::Array<PrimExpr> outputs);
350 
378 ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
379  const ffi::Map<Var, Range>& input_iters,
380  const ffi::Array<Var>& sub_iters,
381  const PrimExpr& predicate, IterMapLevel check_level,
382  arith::Analyzer* analyzer,
383  bool simplify_trivial_iterators = true);
384 
391 
408 IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
409  arith::Analyzer* analyzer);
410 
411 } // namespace arith
412 } // namespace tvm
413 #endif // TVM_ARITH_ITER_AFFINE_MAP_H_
Algebra expression simplifications.
Base node of all primitive expressions.
Definition: expr.h:91
Reference to PrimExprNode.
Definition: expr.h:124
Analyzer that contains bunch of sub-analyzers.
Definition: analyzer.h:634
Base class of all iter map expressions.
Definition: iter_affine_map.h:67
static constexpr const uint32_t _type_child_slots
Definition: iter_affine_map.h:69
TVM_FFI_DECLARE_OBJECT_INFO("arith.IterMapExpr", IterMapExprNode, PrimExprNode)
Managed reference to IterMapExprNode.
Definition: iter_affine_map.h:77
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMapExpr, PrimExpr, IterMapExprNode)
Result of DetectIterMap.
Definition: iter_affine_map.h:241
ffi::Array< ffi::String > errors
Definition: iter_affine_map.h:248
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, Object)
ffi::Array< IterSumExpr > indices
Definition: iter_affine_map.h:244
static void RegisterReflection()
Definition: iter_affine_map.h:260
PrimExpr padding_predicate
Boolean expression indicating if a specific value w.
Definition: iter_affine_map.h:258
Managed reference to IterMapResultNode.
Definition: iter_affine_map.h:274
IterMapResult()
Definition: iter_affine_map.h:277
IterMapResultNode * operator->() const
Definition: iter_affine_map.h:280
Mark the source as an iterator in [0, extent).
Definition: iter_affine_map.h:88
PrimExpr extent
The extent of the iteration.
Definition: iter_affine_map.h:98
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:107
PrimExpr source
The source expression, can either be a IterSumExpr or a Var.
Definition: iter_affine_map.h:94
static void RegisterReflection()
Definition: iter_affine_map.h:100
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMark", IterMarkNode, Object)
Managed reference to IterMarkExprNode.
Definition: iter_affine_map.h:115
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode)
IterMark(PrimExpr source, PrimExpr extent)
constructor.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMark, ObjectRef, IterMarkNode)
Split of an iterator.
Definition: iter_affine_map.h:133
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSplitExpr", IterSplitExprNode, IterMapExprNode)
PrimExpr lower_factor
The lower factor to split the source.
Definition: iter_affine_map.h:138
IterMark source
The source marked iterator.
Definition: iter_affine_map.h:136
PrimExpr scale
Additional scale.
Definition: iter_affine_map.h:142
static void RegisterReflection()
Definition: iter_affine_map.h:144
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:153
PrimExpr extent
The extent of the split.
Definition: iter_affine_map.h:140
Managed reference to IterSplitExprNode.
Definition: iter_affine_map.h:161
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode)
IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent, PrimExpr scale)
constructor
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSplitExpr, IterMapExpr, IterSplitExprNode)
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:193
ffi::Array< IterSplitExpr > args
The args to the sum.
Definition: iter_affine_map.h:196
static void RegisterReflection()
Definition: iter_affine_map.h:200
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: iter_affine_map.h:207
PrimExpr base
The base offset.
Definition: iter_affine_map.h:198
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSumExpr", IterSumExprNode, IterMapExprNode)
Managed reference to IterSumExprNode.
Definition: iter_affine_map.h:215
IterSumExpr(ffi::Array< IterSplitExpr > args, PrimExpr base)
constructor.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSumExpr, IterMapExpr, IterSumExprNode)
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode)
A new diagnostic interface for TVM error reporting.
Base expr nodes in TVM.
ffi::Array< ffi::Array< IterMark > > SubspaceDivide(const ffi::Array< PrimExpr > &bindings, const ffi::Map< Var, Range > &input_iters, const ffi::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:229
@ NoCheck
Definition: iter_affine_map.h:235
@ Bijective
Definition: iter_affine_map.h:231
@ Surjective
Definition: iter_affine_map.h:233
IterMapResult DetectIterMap(const ffi::Array< PrimExpr > &indices, const ffi::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].
ffi::Array< PrimExpr > IterMapSimplify(const ffi::Array< PrimExpr > &indices, const ffi::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.
PrimExpr NormalizeIterMapToExpr(const PrimExpr &expr)
Given an expression that may contain IterMapExpr, transform it to normal PrimExpr.
ffi::Map< Var, PrimExpr > InverseAffineIterMap(const ffi::Array< IterSumExpr > &iter_map, const ffi::Array< PrimExpr > outputs)
Apply the inverse of the affine transformation to the outputs.
IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map< Var, Range > &input_iters, arith::Analyzer *analyzer)
Rewrite index as IterSumExpr.
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Variables in the TIR.