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  ffi::Array<Var> initial_indices;
60 
71  ffi::Array<PrimExpr> final_indices;
72 
83  ffi::Optional<ObjectRef> inverse_index_map;
84 
92 
105  ffi::Array<PrimExpr> MapIndices(const ffi::Array<PrimExpr>& indices,
106  arith::Analyzer* analyzer) const;
107 
124  ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges, arith::Analyzer* analyzer) const;
125 
137  ffi::Array<PrimExpr> MapShape(const ffi::Array<PrimExpr>& shape, arith::Analyzer* analyzer) const;
138 
139  /* \brief Map an Tensor according to this index map
140  *
141  * \param arr_src The Tensor whose layout is transformed by this index map.
142  *
143  * \returns The transformed Tensor.
144  */
146 
152  ffi::String ToPythonString(
153  const std::function<ffi::Optional<ffi::String>(const Var& var)>& f_name_map = nullptr) const;
154 
155  static void RegisterReflection() {
156  namespace refl = tvm::ffi::reflection;
157  refl::ObjectDef<IndexMapNode>()
158  .def_ro("initial_indices", &IndexMapNode::initial_indices,
159  refl::AttachFieldFlag::SEqHashDef())
160  .def_ro("final_indices", &IndexMapNode::final_indices)
161  .def_ro("inverse_index_map", &IndexMapNode::inverse_index_map,
162  refl::AttachFieldFlag::SEqHashIgnore());
163  }
164 
165  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
167 };
168 
169 class IndexMap : public ObjectRef {
170  public:
177  IndexMap(ffi::Array<Var> initial_indices, ffi::Array<PrimExpr> final_indices,
178  ffi::Optional<IndexMap> inverse_index_map = std::nullopt);
179 
187  static IndexMap FromFunc(int ndim, ffi::TypedFunction<ffi::Array<PrimExpr>(ffi::Array<Var>)> func,
188  ffi::Optional<IndexMap> inverse_index_map = std::nullopt);
189 
198  IndexMap Inverse(ffi::Array<Range> initial_ranges, arith::Analyzer* analyzer) const;
199 
209  const std::function<ffi::Optional<ffi::String>(const Var& var)>& f_name_map = nullptr) const;
210 
220  std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(ffi::Array<Range> initial_ranges,
221  arith::Analyzer* analyzer) const;
222 
224 };
225 
231 IndexMap Substitute(const IndexMap& index_map,
232  std::function<ffi::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:634
Managed Tensor. The array is backed by reference counted blocks.
Definition: tensor.h:53
Defines a mapping between two representations of indices into a buffer.
Definition: index_map.h:52
runtime::Tensor MapTensor(runtime::Tensor arr_src) const
ffi::Array< PrimExpr > MapShape(const ffi::Array< PrimExpr > &shape, arith::Analyzer *analyzer) const
Map a buffer shape to the output space.
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.IndexMap", IndexMapNode, Object)
ffi::Array< Var > initial_indices
Variables representing the indices prior to remapping.
Definition: index_map.h:59
ffi::Optional< ObjectRef > inverse_index_map
The inverse index map.
Definition: index_map.h:83
static void RegisterReflection()
Definition: index_map.h:155
ffi::Array< PrimExpr > MapIndices(const ffi::Array< PrimExpr > &indices, arith::Analyzer *analyzer) const
Map indices to the output space.
IndexMapNode()
Default constructor.
Definition: index_map.h:91
ffi::Array< PrimExpr > final_indices
Expressions defining the indices after remapping.
Definition: index_map.h:71
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: index_map.h:165
ffi::Array< Range > MapRanges(const ffi::Array< Range > &ranges, arith::Analyzer *analyzer) const
Map a memory range to the output space.
ffi::String ToPythonString(const std::function< ffi::Optional< ffi::String >(const Var &var)> &f_name_map=nullptr) const
Convert to string representation in Python.
Definition: index_map.h:169
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IndexMap, ObjectRef, IndexMapNode)
IndexMap RenameVariables(const std::function< ffi::Optional< ffi::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(ffi::Array< Range > initial_ranges, arith::Analyzer *analyzer) const
Generate the inverse mapping.
static IndexMap FromFunc(int ndim, ffi::TypedFunction< ffi::Array< PrimExpr >(ffi::Array< Var >)> func, ffi::Optional< IndexMap > inverse_index_map=std::nullopt)
Create an index map from a packed function.
IndexMap(ffi::Array< Var > initial_indices, ffi::Array< PrimExpr > final_indices, ffi::Optional< IndexMap > inverse_index_map=std::nullopt)
The constructor.
IndexMap Inverse(ffi::Array< Range > initial_ranges, arith::Analyzer *analyzer) const
Generate the inverse mapping.
a named variable in TIR
Definition: var.h:77
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< ffi::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:1960
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
Variables in the TIR.