1 //===- AffineMap.h - MLIR Affine Map Class ----------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Affine maps are mathematical functions which map a list of dimension
10 // identifiers and symbols, to multidimensional affine expressions.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #ifndef MLIR_IR_AFFINE_MAP_H
15 #define MLIR_IR_AFFINE_MAP_H
16
17 #include "mlir/IR/AffineExpr.h"
18 #include "mlir/Support/LLVM.h"
19 #include "llvm/ADT/ArrayRef.h"
20 #include "llvm/ADT/DenseMapInfo.h"
21
22 namespace mlir {
23
24 namespace detail {
25 struct AffineMapStorage;
26 } // end namespace detail
27
28 class Attribute;
29 struct LogicalResult;
30 class MLIRContext;
31
32 /// A multi-dimensional affine map
33 /// Affine map's are immutable like Type's, and they are uniqued.
34 /// Eg: (d0, d1) -> (d0/128, d0 mod 128, d1)
35 /// The names used (d0, d1) don't matter - it's the mathematical function that
36 /// is unique to this affine map.
37 class AffineMap {
38 public:
39 using ImplType = detail::AffineMapStorage;
40
AffineMap()41 constexpr AffineMap() : map(nullptr) {}
AffineMap(ImplType * map)42 explicit AffineMap(ImplType *map) : map(map) {}
43
44 /// Returns a zero result affine map with no dimensions or symbols: () -> ().
45 static AffineMap get(MLIRContext *context);
46
47 /// Returns a zero result affine map with `dimCount` dimensions and
48 /// `symbolCount` symbols, e.g.: `(...) -> ()`.
49 static AffineMap get(unsigned dimCount, unsigned symbolCount,
50 MLIRContext *context);
51
52 /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
53 /// to a single output dimension
54 static AffineMap get(unsigned dimCount, unsigned symbolCount,
55 AffineExpr result);
56
57 /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
58 /// to the given results.
59 static AffineMap get(unsigned dimCount, unsigned symbolCount,
60 ArrayRef<AffineExpr> results, MLIRContext *context);
61
62 /// Returns a single constant result affine map.
63 static AffineMap getConstantMap(int64_t val, MLIRContext *context);
64
65 /// Returns an AffineMap with 'numDims' identity result dim exprs.
66 static AffineMap getMultiDimIdentityMap(unsigned numDims,
67 MLIRContext *context);
68
69 /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
70 /// minor dimensions.
71 static AffineMap getMinorIdentityMap(unsigned dims, unsigned results,
72 MLIRContext *context);
73
74 /// Returns an AffineMap representing a permutation.
75 /// The permutation is expressed as a non-empty vector of integers.
76 /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with
77 /// `permutation = [1,2,0]`. All values in `permutation` must be
78 /// integers, in the range 0..`permutation.size()-1` without duplications
79 /// (i.e. `[1,1,2]` is an invalid permutation).
80 static AffineMap getPermutationMap(ArrayRef<unsigned> permutation,
81 MLIRContext *context);
82
83 /// Returns a vector of AffineMaps; each with as many results as
84 /// `exprs.size()`, as many dims as the largest dim in `exprs` and as many
85 /// symbols as the largest symbol in `exprs`.
86 static SmallVector<AffineMap, 4>
87 inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList);
88 static SmallVector<AffineMap, 4>
89 inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList);
90
91 MLIRContext *getContext() const;
92
93 explicit operator bool() { return map != nullptr; }
94 bool operator==(AffineMap other) const { return other.map == map; }
95 bool operator!=(AffineMap other) const { return !(other.map == map); }
96
97 /// Returns true if this affine map is an identity affine map.
98 /// An identity affine map corresponds to an identity affine function on the
99 /// dimensional identifiers.
100 bool isIdentity() const;
101
102 /// Returns true if this affine map is a minor identity, i.e. an identity
103 /// affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
104 bool isMinorIdentity() const;
105
106 /// Returns true if this affine map is an empty map, i.e., () -> ().
107 bool isEmpty() const;
108
109 /// Returns true if this affine map is a single result constant function.
110 bool isSingleConstant() const;
111
112 /// Returns the constant result of this map. This methods asserts that the map
113 /// has a single constant result.
114 int64_t getSingleConstantResult() const;
115
116 // Prints affine map to 'os'.
117 void print(raw_ostream &os) const;
118 void dump() const;
119
120 unsigned getNumDims() const;
121 unsigned getNumSymbols() const;
122 unsigned getNumResults() const;
123 unsigned getNumInputs() const;
124
125 ArrayRef<AffineExpr> getResults() const;
126 AffineExpr getResult(unsigned idx) const;
127
128 /// Extracts the position of the dimensional expression at the given result,
129 /// when the caller knows it is safe to do so.
130 unsigned getDimPosition(unsigned idx) const;
131
132 /// Walk all of the AffineExpr's in this mapping. Each node in an expression
133 /// tree is visited in postorder.
134 void walkExprs(std::function<void(AffineExpr)> callback) const;
135
136 /// This method substitutes any uses of dimensions and symbols (e.g.
137 /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
138 /// expression mapping. Because this can be used to eliminate dims and
139 /// symbols, the client needs to specify the number of dims and symbols in
140 /// the result. The returned map always has the same number of results.
141 AffineMap replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
142 ArrayRef<AffineExpr> symReplacements,
143 unsigned numResultDims,
144 unsigned numResultSyms) const;
145
146 /// Folds the results of the application of an affine map on the provided
147 /// operands to a constant if possible.
148 LogicalResult constantFold(ArrayRef<Attribute> operandConstants,
149 SmallVectorImpl<Attribute> &results) const;
150
151 /// Propagates the constant operands into this affine map. Operands are
152 /// allowed to be null, at which point they are treated as non-constant. This
153 /// does not change the number of symbols and dimensions. Returns a new map,
154 /// which may be equal to the old map if no folding happened. If `results` is
155 /// provided and if all expressions in the map were folded to constants,
156 /// `results` will contain the values of these constants.
157 AffineMap
158 partialConstantFold(ArrayRef<Attribute> operandConstants,
159 SmallVectorImpl<int64_t> *results = nullptr) const;
160
161 /// Returns the AffineMap resulting from composing `this` with `map`.
162 /// The resulting AffineMap has as many AffineDimExpr as `map` and as many
163 /// AffineSymbolExpr as the concatenation of `this` and `map` (in which case
164 /// the symbols of `this` map come first).
165 ///
166 /// Prerequisites:
167 /// The maps are composable, i.e. that the number of AffineDimExpr of `this`
168 /// matches the number of results of `map`.
169 ///
170 /// Example:
171 /// map1: `(d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0)`
172 /// map2: `(d0)[s0] -> (d0 + s0, d0 - s0)`
173 /// map1.compose(map2):
174 /// `(d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)`
175 AffineMap compose(AffineMap map);
176
177 /// Applies composition by the dims of `this` to the integer `values` and
178 /// returns the resulting values. `this` must be symbol-less.
179 SmallVector<int64_t, 4> compose(ArrayRef<int64_t> values);
180
181 /// Returns true if the AffineMap represents a subset (i.e. a projection) of a
182 /// symbol-less permutation map.
183 bool isProjectedPermutation();
184
185 /// Returns true if the AffineMap represents a symbol-less permutation map.
186 bool isPermutation();
187
188 /// Returns the map consisting of the `resultPos` subset.
189 AffineMap getSubMap(ArrayRef<unsigned> resultPos);
190
191 /// Returns the map consisting of the most major `numResults` results.
192 /// Returns the null AffineMap if `numResults` == 0.
193 /// Returns `*this` if `numResults` >= `this->getNumResults()`.
194 AffineMap getMajorSubMap(unsigned numResults);
195
196 /// Returns the map consisting of the most minor `numResults` results.
197 /// Returns the null AffineMap if `numResults` == 0.
198 /// Returns `*this` if `numResults` >= `this->getNumResults()`.
199 AffineMap getMinorSubMap(unsigned numResults);
200
201 friend ::llvm::hash_code hash_value(AffineMap arg);
202
203 /// Methods supporting C API.
getAsOpaquePointer()204 const void *getAsOpaquePointer() const {
205 return static_cast<const void *>(map);
206 }
getFromOpaquePointer(const void * pointer)207 static AffineMap getFromOpaquePointer(const void *pointer) {
208 return AffineMap(reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
209 }
210
211 private:
212 ImplType *map;
213
214 static AffineMap getImpl(unsigned dimCount, unsigned symbolCount,
215 ArrayRef<AffineExpr> results, MLIRContext *context);
216 };
217
218 // Make AffineExpr hashable.
hash_value(AffineMap arg)219 inline ::llvm::hash_code hash_value(AffineMap arg) {
220 return ::llvm::hash_value(arg.map);
221 }
222
223 /// A mutable affine map. Its affine expressions are however unique.
224 struct MutableAffineMap {
225 public:
MutableAffineMapMutableAffineMap226 MutableAffineMap() {}
227 MutableAffineMap(AffineMap map);
228
getResultsMutableAffineMap229 ArrayRef<AffineExpr> getResults() const { return results; }
getResultMutableAffineMap230 AffineExpr getResult(unsigned idx) const { return results[idx]; }
setResultMutableAffineMap231 void setResult(unsigned idx, AffineExpr result) { results[idx] = result; }
getNumResultsMutableAffineMap232 unsigned getNumResults() const { return results.size(); }
getNumDimsMutableAffineMap233 unsigned getNumDims() const { return numDims; }
setNumDimsMutableAffineMap234 void setNumDims(unsigned d) { numDims = d; }
getNumSymbolsMutableAffineMap235 unsigned getNumSymbols() const { return numSymbols; }
setNumSymbolsMutableAffineMap236 void setNumSymbols(unsigned d) { numSymbols = d; }
getContextMutableAffineMap237 MLIRContext *getContext() const { return context; }
238
239 /// Returns true if the idx'th result expression is a multiple of factor.
240 bool isMultipleOf(unsigned idx, int64_t factor) const;
241
242 /// Resets this MutableAffineMap with 'map'.
243 void reset(AffineMap map);
244
245 /// Simplify the (result) expressions in this map using analysis (used by
246 //-simplify-affine-expr pass).
247 void simplify();
248 /// Get the AffineMap corresponding to this MutableAffineMap. Note that an
249 /// AffineMap will be uniqued and stored in context, while a mutable one
250 /// isn't.
251 AffineMap getAffineMap() const;
252
253 private:
254 // Same meaning as AffineMap's fields.
255 SmallVector<AffineExpr, 8> results;
256 unsigned numDims;
257 unsigned numSymbols;
258 /// A pointer to the IR's context to store all newly created
259 /// AffineExprStorage's.
260 MLIRContext *context;
261 };
262
263 /// Simplifies an affine map by simplifying its underlying AffineExpr results.
264 AffineMap simplifyAffineMap(AffineMap map);
265
266 /// Returns a map with the same dimension and symbol count as `map`, but whose
267 /// results are the unique affine expressions of `map`.
268 AffineMap removeDuplicateExprs(AffineMap map);
269
270 /// Returns a map of codomain to domain dimensions such that the first codomain
271 /// dimension for a particular domain dimension is selected.
272 /// Returns an empty map if the input map is empty.
273 /// Returns null map (not empty map) if `map` is not invertible (i.e. `map` does
274 /// not contain a subset that is a permutation of full domain rank).
275 ///
276 /// Prerequisites:
277 /// 1. `map` has no symbols.
278 ///
279 /// Example 1:
280 ///
281 /// ```mlir
282 /// (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0)
283 /// 0 2 3
284 /// ```
285 ///
286 /// returns:
287 ///
288 /// ```mlir
289 /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3)
290 /// ```
291 ///
292 /// Example 2:
293 ///
294 /// ```mlir
295 /// (d0, d1, d2) -> (d1, d0 + d1, d0, d2, d1, d2, d1, d0)
296 /// 0 2 3
297 /// ```
298 ///
299 /// returns:
300 ///
301 /// ```mlir
302 /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3)
303 /// ```
304 AffineMap inversePermutation(AffineMap map);
305
306 /// Concatenates a list of `maps` into a single AffineMap, stepping over
307 /// potentially empty maps. Assumes each of the underlying map has 0 symbols.
308 /// The resulting map has a number of dims equal to the max of `maps`' dims and
309 /// the concatenated results as its results.
310 /// Returns an empty map if all input `maps` are empty.
311 ///
312 /// Example:
313 /// When applied to the following list of 3 affine maps,
314 ///
315 /// ```mlir
316 /// {
317 /// (i, j, k) -> (i, k),
318 /// (i, j, k) -> (k, j),
319 /// (i, j, k) -> (i, j)
320 /// }
321 /// ```
322 ///
323 /// Returns the map:
324 ///
325 /// ```mlir
326 /// (i, j, k) -> (i, k, k, j, i, j)
327 /// ```
328 AffineMap concatAffineMaps(ArrayRef<AffineMap> maps);
329
330 AffineMap getProjectedMap(AffineMap map,
331 ArrayRef<unsigned> projectedDimensions);
332
333 inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
334 map.print(os);
335 return os;
336 }
337 } // end namespace mlir
338
339 namespace llvm {
340
341 // AffineExpr hash just like pointers
342 template <> struct DenseMapInfo<mlir::AffineMap> {
343 static mlir::AffineMap getEmptyKey() {
344 auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
345 return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
346 }
347 static mlir::AffineMap getTombstoneKey() {
348 auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
349 return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
350 }
351 static unsigned getHashValue(mlir::AffineMap val) {
352 return mlir::hash_value(val);
353 }
354 static bool isEqual(mlir::AffineMap LHS, mlir::AffineMap RHS) {
355 return LHS == RHS;
356 }
357 };
358
359 } // namespace llvm
360
361 #endif // MLIR_IR_AFFINE_MAP_H
362