tvm
index_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 
26 #ifndef TVM_TIR_INDEX_MAP_H_
27 #define TVM_TIR_INDEX_MAP_H_
28 
29 #include <tvm/ffi/container/array.h>
30 #include <tvm/ir/expr.h>
31 #include <tvm/runtime/object.h>
32 #include <tvm/tir/var.h>
33 
34 #include <utility>
35 
36 namespace tvm {
37 namespace arith {
38 class Analyzer;
39 }
40 } // namespace tvm
41 
42 namespace tvm {
43 namespace tir {
44 
52 class IndexMapNode : public Object {
53  public:
59  Array<Var> initial_indices;
60 
71  Array<PrimExpr> final_indices;
72 
83  Optional<ObjectRef> inverse_index_map;
84 
92 
105  Array<PrimExpr> MapIndices(const Array<PrimExpr>& indices, arith::Analyzer* analyzer) const;
106 
123  Array<Range> MapRanges(const Array<Range>& ranges, arith::Analyzer* analyzer) const;
124 
136  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape, arith::Analyzer* analyzer) const;
137 
138  /* \brief Map an NDArray according to this index map
139  *
140  * \param arr_src The NDArray whose layout is transformed by this index map.
141  *
142  * \returns The transformed NDArray.
143  */
145 
152  const std::function<Optional<String>(const Var& var)>& f_name_map = nullptr) const;
153 
154  static void RegisterReflection() {
155  namespace refl = tvm::ffi::reflection;
156  refl::ObjectDef<IndexMapNode>()
157  .def_ro("initial_indices", &IndexMapNode::initial_indices,
158  refl::AttachFieldFlag::SEqHashDef())
159  .def_ro("final_indices", &IndexMapNode::final_indices)
160  .def_ro("inverse_index_map", &IndexMapNode::inverse_index_map,
161  refl::AttachFieldFlag::SEqHashIgnore());
162  }
163 
164  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
165  static constexpr const char* _type_key = "tir.IndexMap";
167 };
168 
169 class IndexMap : public ObjectRef {
170  public:
177  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices,
178  Optional<IndexMap> inverse_index_map = std::nullopt);
179 
187  static IndexMap FromFunc(int ndim, ffi::TypedFunction<Array<PrimExpr>(Array<Var>)> func,
188  Optional<IndexMap> inverse_index_map = std::nullopt);
189 
198  IndexMap Inverse(Array<Range> initial_ranges, arith::Analyzer* analyzer) const;
199 
209  const std::function<Optional<String>(const Var& var)>& f_name_map = nullptr) const;
210 
220  std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(Array<Range> initial_ranges,
221  arith::Analyzer* analyzer) const;
222 
224 };
225 
231 IndexMap Substitute(const IndexMap& index_map,
232  std::function<Optional<PrimExpr>(const Var& var)> f_subst);
233 
234 } // namespace tir
235 } // namespace tvm
236 
237 #endif // TVM_TIR_INDEX_MAP_H_
Analyzer that contains bunch of sub-analyzers.
Definition: analyzer.h:636
Managed NDArray. The array is backed by reference counted blocks.
Definition: ndarray.h:53
Defines a mapping between two representations of indices into a buffer.
Definition: index_map.h:52
Array< PrimExpr > MapIndices(const Array< PrimExpr > &indices, arith::Analyzer *analyzer) const
Map indices to the output space.
static constexpr const char * _type_key
Definition: index_map.h:165
Array< PrimExpr > final_indices
Expressions defining the indices after remapping.
Definition: index_map.h:71
Optional< ObjectRef > inverse_index_map
The inverse index map.
Definition: index_map.h:83
Array< PrimExpr > MapShape(const Array< PrimExpr > &shape, arith::Analyzer *analyzer) const
Map a buffer shape to the output space.
Array< Var > initial_indices
Variables representing the indices prior to remapping.
Definition: index_map.h:59
static void RegisterReflection()
Definition: index_map.h:154
IndexMapNode()
Default constructor.
Definition: index_map.h:91
Array< Range > MapRanges(const Array< Range > &ranges, arith::Analyzer *analyzer) const
Map a memory range to the output space.
TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object)
String ToPythonString(const std::function< Optional< String >(const Var &var)> &f_name_map=nullptr) const
Convert to string representation in Python.
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: index_map.h:164
runtime::NDArray MapNDArray(runtime::NDArray arr_src) const
Definition: index_map.h:169
TVM_DEFINE_OBJECT_REF_METHODS(IndexMap, ObjectRef, IndexMapNode)
IndexMap Inverse(Array< Range > initial_ranges, arith::Analyzer *analyzer) const
Generate the inverse mapping.
IndexMap RenameVariables(const std::function< Optional< String >(const Var &var)> &f_name_map=nullptr) const
Rename the variables in the index map and ensure the names are unique.
std::pair< IndexMap, PrimExpr > NonSurjectiveInverse(Array< Range > initial_ranges, arith::Analyzer *analyzer) const
Generate the inverse mapping.
IndexMap(Array< Var > initial_indices, Array< PrimExpr > final_indices, Optional< IndexMap > inverse_index_map=std::nullopt)
The constructor.
static IndexMap FromFunc(int ndim, ffi::TypedFunction< Array< PrimExpr >(Array< Var >)> func, Optional< IndexMap > inverse_index_map=std::nullopt)
Create an index map from a packed function.
a named variable in TIR
Definition: var.h:78
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.
IndexMap Substitute(const IndexMap &index_map, std::function< Optional< PrimExpr >(const Var &var)> f_subst)
Substitute variables in an index map.
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1945
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
Variables in the TIR.