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_AFFINEMAP_H |
15 | #define MLIR_IR_AFFINEMAP_H |
16 | |
17 | #include "mlir/IR/AffineExpr.h" |
18 | #include "mlir/IR/Value.h" |
19 | #include "mlir/Support/LLVM.h" |
20 | #include "llvm/ADT/ArrayRef.h" |
21 | #include "llvm/ADT/DenseMapInfo.h" |
22 | #include "llvm/ADT/SmallBitVector.h" |
23 | #include "llvm/ADT/SmallVectorExtras.h" |
24 | #include <optional> |
25 | |
26 | namespace llvm { |
27 | class SmallBitVector; |
28 | } // namespace llvm |
29 | |
30 | namespace mlir { |
31 | |
32 | namespace detail { |
33 | struct AffineMapStorage; |
34 | } // namespace detail |
35 | |
36 | class Attribute; |
37 | class Builder; |
38 | struct LogicalResult; |
39 | class OpFoldResult; |
40 | class MLIRContext; |
41 | |
42 | /// A multi-dimensional affine map |
43 | /// Affine map's are immutable like Type's, and they are uniqued. |
44 | /// Eg: (d0, d1) -> (d0/128, d0 mod 128, d1) |
45 | /// The names used (d0, d1) don't matter - it's the mathematical function that |
46 | /// is unique to this affine map. |
47 | class AffineMap { |
48 | public: |
49 | using ImplType = detail::AffineMapStorage; |
50 | |
51 | constexpr AffineMap() = default; |
52 | explicit AffineMap(ImplType *map) : map(map) {} |
53 | |
54 | /// Returns a zero result affine map with no dimensions or symbols: () -> (). |
55 | static AffineMap get(MLIRContext *context); |
56 | |
57 | /// Returns a zero result affine map with `dimCount` dimensions and |
58 | /// `symbolCount` symbols, e.g.: `(...) -> ()`. |
59 | static AffineMap get(unsigned dimCount, unsigned symbolCount, |
60 | MLIRContext *context); |
61 | |
62 | /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping |
63 | /// to a single output dimension |
64 | static AffineMap get(unsigned dimCount, unsigned symbolCount, |
65 | AffineExpr result); |
66 | |
67 | /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping |
68 | /// to the given results. |
69 | static AffineMap get(unsigned dimCount, unsigned symbolCount, |
70 | ArrayRef<AffineExpr> results, MLIRContext *context); |
71 | |
72 | /// Returns a single constant result affine map. |
73 | static AffineMap getConstantMap(int64_t val, MLIRContext *context); |
74 | |
75 | /// Returns an AffineMap with 'numDims' identity result dim exprs. |
76 | static AffineMap getMultiDimIdentityMap(unsigned numDims, |
77 | MLIRContext *context); |
78 | |
79 | /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most |
80 | /// minor dimensions. |
81 | static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, |
82 | MLIRContext *context); |
83 | |
84 | /// Returns an identity affine map with `numDims` input dimensions and |
85 | /// filtered results using `keepDimFilter`. If `keepDimFilter` returns true |
86 | /// for a dimension, the dimension is kept in the affine map results. |
87 | /// Otherwise, the dimension is dropped from the results. |
88 | /// |
89 | /// Examples: |
90 | /// * getFilteredIdentityMap(4, [false, true, false, true]) |
91 | /// -> affine_map<(d0, d1, d2, d3) -> (d1, d3)> |
92 | /// * getFilteredIdentityMap(3, [false, false, true]) |
93 | /// -> affine_map<(d0, d1, d2) -> (d2)> |
94 | static AffineMap |
95 | getFilteredIdentityMap(MLIRContext *ctx, unsigned numDims, |
96 | llvm::function_ref<bool(AffineDimExpr)> keepDimFilter); |
97 | |
98 | /// Returns an AffineMap representing a permutation. |
99 | /// The permutation is expressed as a non-empty vector of integers. |
100 | /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with |
101 | /// `permutation = [1,2,0]`. All values in `permutation` must be |
102 | /// integers, in the range 0..`permutation.size()-1` without duplications |
103 | /// (i.e. `[1,1,2]` is an invalid permutation). |
104 | static AffineMap getPermutationMap(ArrayRef<unsigned> permutation, |
105 | MLIRContext *context); |
106 | static AffineMap getPermutationMap(ArrayRef<int64_t> permutation, |
107 | MLIRContext *context); |
108 | |
109 | /// Returns an affine map with `numDims` input dimensions and results |
110 | /// specified by `targets`. |
111 | /// |
112 | /// Examples: |
113 | /// * getMultiDimMapWithTargets(3, [0, 2, 1]) |
114 | /// -> affine_map<(d0, d1, d2) -> (d0, d2, d1)> |
115 | /// * getMultiDimMapWithTargets(3, [2, 1]) |
116 | /// -> affine_map<(d0, d1, d2) -> (d2, d1)> |
117 | static AffineMap getMultiDimMapWithTargets(unsigned numDims, |
118 | ArrayRef<unsigned> targets, |
119 | MLIRContext *context); |
120 | |
121 | /// Returns a vector of AffineMaps; each with as many results as |
122 | /// `exprs.size()`, as many dims as the largest dim in `exprs` and as many |
123 | /// symbols as the largest symbol in `exprs`. |
124 | static SmallVector<AffineMap, 4> |
125 | inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList, |
126 | MLIRContext *context); |
127 | static SmallVector<AffineMap, 4> |
128 | inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList, |
129 | MLIRContext *context); |
130 | |
131 | MLIRContext *getContext() const; |
132 | |
133 | explicit operator bool() const { return map != nullptr; } |
134 | bool operator==(AffineMap other) const { return other.map == map; } |
135 | bool operator!=(AffineMap other) const { return !(other.map == map); } |
136 | |
137 | /// Returns true if this affine map is an identity affine map. |
138 | /// An identity affine map corresponds to an identity affine function on the |
139 | /// dimensional identifiers. |
140 | bool isIdentity() const; |
141 | |
142 | /// Returns true if this affine map is an identity affine map on the symbol |
143 | /// identifiers. |
144 | bool isSymbolIdentity() const; |
145 | |
146 | /// Returns true if this affine map is a minor identity, i.e. an identity |
147 | /// affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions. |
148 | bool isMinorIdentity() const; |
149 | |
150 | /// Returns true if this affine map is a minor identity up to broadcasted |
151 | /// dimensions which are indicated by value 0 in the result. If |
152 | /// `broadcastedDims` is not null, it will be populated with the indices of |
153 | /// the broadcasted dimensions in the result array. |
154 | /// Example: affine_map<(d0, d1, d2, d3, d4) -> (0, d2, 0, d4)> |
155 | /// (`broadcastedDims` will contain [0, 2]) |
156 | bool isMinorIdentityWithBroadcasting( |
157 | SmallVectorImpl<unsigned> *broadcastedDims = nullptr) const; |
158 | |
159 | /// Return true if this affine map can be converted to a minor identity with |
160 | /// broadcast by doing a permute. Return a permutation (there may be |
161 | /// several) to apply to get to a minor identity with broadcasts. |
162 | /// Ex: |
163 | /// * (d0, d1, d2) -> (0, d1) maps to minor identity (d1, 0 = d2) with |
164 | /// perm = [1, 0] and broadcast d2 |
165 | /// * (d0, d1, d2) -> (d0, 0) cannot be mapped to a minor identity by |
166 | /// permutation + broadcast |
167 | /// * (d0, d1, d2, d3) -> (0, d1, d3) maps to minor identity (d1, 0 = d2, d3) |
168 | /// with perm = [1, 0, 2] and broadcast d2 |
169 | /// * (d0, d1) -> (d1, 0, 0, d0) maps to minor identity (d0, d1) with extra |
170 | /// leading broadcat dimensions. The map returned would be (0, 0, d0, d1) |
171 | /// with perm = [3, 0, 1, 2] |
172 | bool isPermutationOfMinorIdentityWithBroadcasting( |
173 | SmallVectorImpl<unsigned> &permutedDims) const; |
174 | |
175 | /// Returns true if this affine map is an empty map, i.e., () -> (). |
176 | bool isEmpty() const; |
177 | |
178 | /// Returns true if this affine map is a single result constant function. |
179 | bool isSingleConstant() const; |
180 | |
181 | /// Returns true if this affine map has only constant results. |
182 | bool isConstant() const; |
183 | |
184 | /// Returns the constant result of this map. This methods asserts that the map |
185 | /// has a single constant result. |
186 | int64_t getSingleConstantResult() const; |
187 | |
188 | /// Returns the constant results of this map. This method asserts that the map |
189 | /// has all constant results. |
190 | SmallVector<int64_t> getConstantResults() const; |
191 | |
192 | // Prints affine map to 'os'. |
193 | void print(raw_ostream &os) const; |
194 | void dump() const; |
195 | |
196 | unsigned getNumDims() const; |
197 | unsigned getNumSymbols() const; |
198 | unsigned getNumResults() const; |
199 | unsigned getNumInputs() const; |
200 | |
201 | ArrayRef<AffineExpr> getResults() const; |
202 | AffineExpr getResult(unsigned idx) const; |
203 | |
204 | /// Extracts the position of the dimensional expression at the given result, |
205 | /// when the caller knows it is safe to do so. |
206 | unsigned getDimPosition(unsigned idx) const; |
207 | |
208 | /// Extracts the first result position where `input` dimension resides. |
209 | /// Returns `std::nullopt` if `input` is not a dimension expression or cannot |
210 | /// be found in results. |
211 | std::optional<unsigned> getResultPosition(AffineExpr input) const; |
212 | |
213 | /// Return true if any affine expression involves AffineDimExpr `position`. |
214 | bool isFunctionOfDim(unsigned position) const { |
215 | return llvm::any_of(Range: getResults(), P: [&](AffineExpr e) { |
216 | return e.isFunctionOfDim(position); |
217 | }); |
218 | } |
219 | |
220 | /// Return true if any affine expression involves AffineSymbolExpr `position`. |
221 | bool isFunctionOfSymbol(unsigned position) const { |
222 | return llvm::any_of(Range: getResults(), P: [&](AffineExpr e) { |
223 | return e.isFunctionOfSymbol(position); |
224 | }); |
225 | } |
226 | |
227 | /// Walk all of the AffineExpr's in this mapping. Each node in an expression |
228 | /// tree is visited in postorder. |
229 | void walkExprs(llvm::function_ref<void(AffineExpr)> callback) const; |
230 | |
231 | /// This method substitutes any uses of dimensions and symbols (e.g. |
232 | /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified |
233 | /// expression mapping. Because this can be used to eliminate dims and |
234 | /// symbols, the client needs to specify the number of dims and symbols in |
235 | /// the result. The returned map always has the same number of results. |
236 | AffineMap replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, |
237 | ArrayRef<AffineExpr> symReplacements, |
238 | unsigned numResultDims, |
239 | unsigned numResultSyms) const; |
240 | |
241 | /// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to |
242 | /// each of the results and return a new AffineMap with the new results and |
243 | /// with the specified number of dims and symbols. |
244 | AffineMap replace(AffineExpr expr, AffineExpr replacement, |
245 | unsigned numResultDims, unsigned numResultSyms) const; |
246 | |
247 | /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the |
248 | /// results and return a new AffineMap with the new results and with inferred |
249 | /// number of dims and symbols. |
250 | AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map) const; |
251 | |
252 | /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the |
253 | /// results and return a new AffineMap with the new results and with the |
254 | /// specified number of dims and symbols. |
255 | AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map, |
256 | unsigned numResultDims, unsigned numResultSyms) const; |
257 | |
258 | /// Replace dims[offset ... numDims) |
259 | /// by dims[offset + shift ... shift + numDims). |
260 | AffineMap shiftDims(unsigned shift, unsigned offset = 0) const { |
261 | assert(offset <= getNumDims()); |
262 | return AffineMap::get(dimCount: getNumDims() + shift, symbolCount: getNumSymbols(), |
263 | results: llvm::map_to_vector<4>( |
264 | C: getResults(), |
265 | F: [&](AffineExpr e) { |
266 | return e.shiftDims(numDims: getNumDims(), shift, offset); |
267 | }), |
268 | context: getContext()); |
269 | } |
270 | |
271 | /// Replace symbols[offset ... numSymbols) |
272 | /// by symbols[offset + shift ... shift + numSymbols). |
273 | AffineMap shiftSymbols(unsigned shift, unsigned offset = 0) const { |
274 | return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols() + shift, |
275 | results: llvm::map_to_vector<4>(C: getResults(), |
276 | F: [&](AffineExpr e) { |
277 | return e.shiftSymbols( |
278 | numSymbols: getNumSymbols(), shift, |
279 | offset); |
280 | }), |
281 | context: getContext()); |
282 | } |
283 | |
284 | /// Returns a new AffineMap with the same number of dims and symbols and one |
285 | /// less result at `pos`, dropped. |
286 | AffineMap dropResult(int64_t pos) const { |
287 | return dropResults(positions: ArrayRef({pos})); |
288 | } |
289 | |
290 | // Returns a new AffineMap with the same number of dims and symbols, but all |
291 | // results in `positions` dropped. |
292 | AffineMap dropResults(ArrayRef<int64_t> positions) const { |
293 | SmallVector<int64_t> reverse_sorted_positions = llvm::to_vector(Range&: positions); |
294 | llvm::sort(C&: reverse_sorted_positions, Comp: std::greater<int64_t>()); |
295 | |
296 | auto exprs = llvm::to_vector<4>(Range: getResults()); |
297 | for (int64_t pos : reverse_sorted_positions) |
298 | exprs.erase(CI: exprs.begin() + pos); |
299 | return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: exprs, context: getContext()); |
300 | } |
301 | |
302 | // Returns a new AffineMap with the same number of dims and symbols, but all |
303 | // results in `positions` dropped. |
304 | AffineMap dropResults(const llvm::SmallBitVector &positions) const; |
305 | |
306 | /// Returns a new AffineMap with the same number of dims and symbols and an |
307 | /// extra result inserted at `pos`. |
308 | AffineMap insertResult(AffineExpr expr, unsigned pos) const { |
309 | auto exprs = llvm::to_vector<4>(Range: getResults()); |
310 | exprs.insert(I: exprs.begin() + pos, Elt: expr); |
311 | return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: exprs, context: getContext()); |
312 | } |
313 | |
314 | /// Folds the results of the application of an affine map on the provided |
315 | /// operands to a constant if possible. |
316 | LogicalResult constantFold(ArrayRef<Attribute> operandConstants, |
317 | SmallVectorImpl<Attribute> &results, |
318 | bool *hasPoison = nullptr) const; |
319 | |
320 | /// Propagates the constant operands into this affine map. Operands are |
321 | /// allowed to be null, at which point they are treated as non-constant. This |
322 | /// does not change the number of symbols and dimensions. Returns a new map, |
323 | /// which may be equal to the old map if no folding happened. If `results` is |
324 | /// provided and if all expressions in the map were folded to constants, |
325 | /// `results` will contain the values of these constants. |
326 | AffineMap partialConstantFold(ArrayRef<Attribute> operandConstants, |
327 | SmallVectorImpl<int64_t> *results = nullptr, |
328 | bool *hasPoison = nullptr) const; |
329 | |
330 | /// Returns the AffineMap resulting from composing `this` with `map`. |
331 | /// The resulting AffineMap has as many AffineDimExpr as `map` and as many |
332 | /// AffineSymbolExpr as the concatenation of `this` and `map` (in which case |
333 | /// the symbols of `this` map come first). |
334 | /// |
335 | /// Prerequisites: |
336 | /// The maps are composable, i.e. that the number of AffineDimExpr of `this` |
337 | /// matches the number of results of `map`. |
338 | /// |
339 | /// Example: |
340 | /// map1: `(d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0)` |
341 | /// map2: `(d0)[s0] -> (d0 + s0, d0 - s0)` |
342 | /// map1.compose(map2): |
343 | /// `(d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)` |
344 | AffineMap compose(AffineMap map) const; |
345 | |
346 | /// Applies composition by the dims of `this` to the integer `values` and |
347 | /// returns the resulting values. `this` must be symbol-less. |
348 | SmallVector<int64_t, 4> compose(ArrayRef<int64_t> values) const; |
349 | |
350 | /// Returns true if the AffineMap represents a subset (i.e. a projection) of a |
351 | /// symbol-less permutation map. `allowZeroInResults` allows projected |
352 | /// permutation maps with constant zero result expressions. |
353 | /// TODO: Remove `allowZeroInResults` when constant zero result expressions |
354 | /// are broadly supported. |
355 | bool isProjectedPermutation(bool allowZeroInResults = false) const; |
356 | |
357 | /// Returns true if the AffineMap represents a symbol-less permutation map. |
358 | bool isPermutation() const; |
359 | |
360 | /// Returns the map consisting of the `resultPos` subset. |
361 | AffineMap getSubMap(ArrayRef<unsigned> resultPos) const; |
362 | |
363 | /// Returns the map consisting of `length` expressions starting from `start`. |
364 | AffineMap getSliceMap(unsigned start, unsigned length) const; |
365 | |
366 | /// Returns the map consisting of the most major `numResults` results. |
367 | /// Returns the null AffineMap if `numResults` == 0. |
368 | /// Returns `*this` if `numResults` >= `this->getNumResults()`. |
369 | AffineMap getMajorSubMap(unsigned numResults) const; |
370 | |
371 | /// Returns the map consisting of the most minor `numResults` results. |
372 | /// Returns the null AffineMap if `numResults` == 0. |
373 | /// Returns `*this` if `numResults` >= `this->getNumResults()`. |
374 | AffineMap getMinorSubMap(unsigned numResults) const; |
375 | |
376 | /// Get the largest known divisor of all map expressions. |
377 | /// For eg: for (d0, d1) -> (8*d0 + 4, 4*d1 + 2), the result is 2. |
378 | /// In the case of maps with no expressions or all zero constant expressions, |
379 | /// the largest known divisor is trivially the max uint64_t value. |
380 | uint64_t getLargestKnownDivisorOfMapExprs(); |
381 | |
382 | friend ::llvm::hash_code hash_value(AffineMap arg); |
383 | |
384 | /// Methods supporting C API. |
385 | const void *getAsOpaquePointer() const { |
386 | return static_cast<const void *>(map); |
387 | } |
388 | static AffineMap getFromOpaquePointer(const void *pointer) { |
389 | return AffineMap(reinterpret_cast<ImplType *>(const_cast<void *>(pointer))); |
390 | } |
391 | |
392 | private: |
393 | ImplType *map{nullptr}; |
394 | |
395 | static AffineMap getImpl(unsigned dimCount, unsigned symbolCount, |
396 | ArrayRef<AffineExpr> results, MLIRContext *context); |
397 | }; |
398 | |
399 | // Make AffineExpr hashable. |
400 | inline ::llvm::hash_code hash_value(AffineMap arg) { |
401 | return ::llvm::hash_value(ptr: arg.map); |
402 | } |
403 | |
404 | /// A mutable affine map. Its affine expressions are however unique. |
405 | struct MutableAffineMap { |
406 | public: |
407 | MutableAffineMap() = default; |
408 | MutableAffineMap(AffineMap map); |
409 | |
410 | ArrayRef<AffineExpr> getResults() const { return results; } |
411 | AffineExpr getResult(unsigned idx) const { return results[idx]; } |
412 | void setResult(unsigned idx, AffineExpr result) { results[idx] = result; } |
413 | unsigned getNumResults() const { return results.size(); } |
414 | unsigned getNumDims() const { return numDims; } |
415 | void setNumDims(unsigned d) { numDims = d; } |
416 | unsigned getNumSymbols() const { return numSymbols; } |
417 | void setNumSymbols(unsigned d) { numSymbols = d; } |
418 | MLIRContext *getContext() const { return context; } |
419 | |
420 | /// Returns true if the idx'th result expression is a multiple of factor. |
421 | bool isMultipleOf(unsigned idx, int64_t factor) const; |
422 | |
423 | /// Resets this MutableAffineMap with 'map'. |
424 | void reset(AffineMap map); |
425 | |
426 | /// Simplify the (result) expressions in this map using analysis (used by |
427 | //-simplify-affine-expr pass). |
428 | void simplify(); |
429 | /// Get the AffineMap corresponding to this MutableAffineMap. Note that an |
430 | /// AffineMap will be uniqued and stored in context, while a mutable one |
431 | /// isn't. |
432 | AffineMap getAffineMap() const; |
433 | |
434 | private: |
435 | // Same meaning as AffineMap's fields. |
436 | SmallVector<AffineExpr, 8> results; |
437 | unsigned numDims = 0; |
438 | unsigned numSymbols = 0; |
439 | /// A pointer to the IR's context to store all newly created |
440 | /// AffineExprStorage's. |
441 | MLIRContext *context = nullptr; |
442 | }; |
443 | |
444 | /// Simplifies an affine map by simplifying its underlying AffineExpr results. |
445 | AffineMap simplifyAffineMap(AffineMap map); |
446 | |
447 | /// Drop the dims that are listed in `unusedDims`. |
448 | AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims); |
449 | |
450 | /// Drop the dims that are not used. |
451 | AffineMap compressUnusedDims(AffineMap map); |
452 | |
453 | /// Drop the dims that are not used by any of the individual maps in `maps`. |
454 | /// Asserts that all maps in `maps` are normalized to the same number of |
455 | /// dims and symbols. |
456 | SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps); |
457 | |
458 | /// Drop the symbols that are listed in `unusedSymbols`. |
459 | AffineMap compressSymbols(AffineMap map, |
460 | const llvm::SmallBitVector &unusedSymbols); |
461 | |
462 | /// Drop the symbols that are not used. |
463 | AffineMap compressUnusedSymbols(AffineMap map); |
464 | |
465 | /// Drop the symbols that are not used by any of the individual maps in `maps`. |
466 | /// Asserts that all maps in `maps` are normalized to the same number of |
467 | /// dims and symbols. |
468 | SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps); |
469 | |
470 | /// Fold all attributes among the given operands into the affine map. Return the |
471 | /// folded affine map. Return all remaining values via `remainingValues`. |
472 | AffineMap foldAttributesIntoMap(Builder &b, AffineMap map, |
473 | ArrayRef<OpFoldResult> operands, |
474 | SmallVector<Value> &remainingValues); |
475 | |
476 | /// Returns a map with the same dimension and symbol count as `map`, but whose |
477 | /// results are the unique affine expressions of `map`. |
478 | AffineMap removeDuplicateExprs(AffineMap map); |
479 | |
480 | /// Returns a map of codomain to domain dimensions such that the first codomain |
481 | /// dimension for a particular domain dimension is selected. |
482 | /// Returns an empty map if the input map is empty. |
483 | /// Returns null map (not empty map) if `map` is not invertible (i.e. `map` does |
484 | /// not contain a subset that is a permutation of full domain rank). |
485 | /// |
486 | /// Prerequisites: |
487 | /// 1. `map` has no symbols. |
488 | /// |
489 | /// Example 1: |
490 | /// |
491 | /// ```mlir |
492 | /// (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0) |
493 | /// 0 2 3 |
494 | /// ``` |
495 | /// |
496 | /// returns: |
497 | /// |
498 | /// ```mlir |
499 | /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) |
500 | /// ``` |
501 | /// |
502 | /// Example 2: |
503 | /// |
504 | /// ```mlir |
505 | /// (d0, d1, d2) -> (d1, d0 + d1, d0, d2, d1, d2, d1, d0) |
506 | /// 0 2 3 |
507 | /// ``` |
508 | /// |
509 | /// returns: |
510 | /// |
511 | /// ```mlir |
512 | /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) |
513 | /// ``` |
514 | AffineMap inversePermutation(AffineMap map); |
515 | |
516 | /// Return the reverse map of a projected permutation where the projected |
517 | /// dimensions are transformed into 0s. |
518 | /// |
519 | /// Prerequisites: `map` must be a projected permutation. |
520 | /// |
521 | /// Example 1: |
522 | /// |
523 | /// ```mlir |
524 | /// affine_map<(d0, d1, d2, d3) -> (d2, d0)> |
525 | /// ``` |
526 | /// |
527 | /// returns: |
528 | /// |
529 | /// ```mlir |
530 | /// affine_map<(d0, d1) -> (d1, 0, d0, 0)> |
531 | /// ``` |
532 | /// |
533 | /// Example 2: |
534 | /// |
535 | /// ```mlir |
536 | /// affine_map<(d0, d1, d2, d3) -> (d0, d3)> |
537 | /// ``` |
538 | /// |
539 | /// returns: |
540 | /// |
541 | /// ```mlir |
542 | /// affine_map<(d0, d1) -> (d0, 0, 0, d1)> |
543 | /// ``` |
544 | /// |
545 | /// Example 3: |
546 | /// |
547 | /// ```mlir |
548 | /// affine_map<(d0, d1, d2, d3) -> (d2)> |
549 | /// ``` |
550 | /// |
551 | /// returns: |
552 | /// |
553 | /// ```mlir |
554 | /// affine_map<(d0) -> (0, 0, d0, 0)> |
555 | /// ``` |
556 | /// Example 4: |
557 | /// |
558 | /// ```mlir |
559 | /// affine_map<(d0, d1, d2) -> (d0, 0)> |
560 | /// ``` |
561 | /// |
562 | /// returns: |
563 | /// |
564 | /// ```mlir |
565 | /// affine_map<(d0, d1) -> (d0, 0, 0)> |
566 | /// ``` |
567 | AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map); |
568 | |
569 | /// Concatenates a list of `maps` into a single AffineMap, stepping over |
570 | /// potentially empty maps. Assumes each of the underlying map has 0 symbols. |
571 | /// The resulting map has a number of dims equal to the max of `maps`' dims and |
572 | /// the concatenated results as its results. |
573 | /// Returns an empty map if all input `maps` are empty. |
574 | /// |
575 | /// Example: |
576 | /// When applied to the following list of 3 affine maps, |
577 | /// |
578 | /// ```mlir |
579 | /// { |
580 | /// (i, j, k) -> (i, k), |
581 | /// (i, j, k) -> (k, j), |
582 | /// (i, j, k) -> (i, j) |
583 | /// } |
584 | /// ``` |
585 | /// |
586 | /// Returns the map: |
587 | /// |
588 | /// ```mlir |
589 | /// (i, j, k) -> (i, k, k, j, i, j) |
590 | /// ``` |
591 | AffineMap concatAffineMaps(ArrayRef<AffineMap> maps); |
592 | |
593 | /// Returns the map that results from projecting out the dimensions specified in |
594 | /// `projectedDimensions`. The projected dimensions are set to 0. |
595 | /// |
596 | /// Example: |
597 | /// 1) map : affine_map<(d0, d1, d2) -> (d0, d1)> |
598 | /// projected_dimensions : {2} |
599 | /// result : affine_map<(d0, d1) -> (d0, d1)> |
600 | /// |
601 | /// 2) map : affine_map<(d0, d1) -> (d0 + d1)> |
602 | /// projected_dimensions : {1} |
603 | /// result : affine_map<(d0) -> (d0)> |
604 | /// |
605 | /// 3) map : affine_map<(d0, d1, d2) -> (d0, d1)> |
606 | /// projected_dimensions : {1} |
607 | /// result : affine_map<(d0, d1) -> (d0, 0)> |
608 | /// |
609 | /// This function also compresses the dims when the boolean flag is true. |
610 | AffineMap projectDims(AffineMap map, |
611 | const llvm::SmallBitVector &projectedDimensions, |
612 | bool compressDimsFlag = false); |
613 | /// Symbol counterpart of `projectDims`. |
614 | /// This function also compresses the symbols when the boolean flag is true. |
615 | AffineMap projectSymbols(AffineMap map, |
616 | const llvm::SmallBitVector &projectedSymbols, |
617 | bool compressSymbolsFlag = false); |
618 | /// Calls `projectDims(map, projectedDimensions, compressDimsFlag)`. |
619 | /// If `compressSymbolsFlag` is true, additionally call `compressUnusedSymbols`. |
620 | AffineMap getProjectedMap(AffineMap map, |
621 | const llvm::SmallBitVector &projectedDimensions, |
622 | bool compressDimsFlag = true, |
623 | bool compressSymbolsFlag = true); |
624 | |
625 | // Return a bitvector where each bit set indicates a dimension that is not used |
626 | // by any of the maps in the input array `maps`. |
627 | llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef<AffineMap> maps); |
628 | |
629 | // Return a bitvector where each bit set indicates a symbol that is not used |
630 | // by any of the maps in the input array `maps`. |
631 | llvm::SmallBitVector getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps); |
632 | |
633 | /// Expand `map` to operate on `rank` dims while projecting out the dims in |
634 | /// `projectedDimensions`. This amounts to composing `map` with |
635 | /// `id(rank).dropResults(projectedDimensions)`. |
636 | AffineMap expandDimsToRank(AffineMap map, int64_t rank, |
637 | const llvm::SmallBitVector &projectedDimensions); |
638 | |
639 | inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) { |
640 | map.print(os); |
641 | return os; |
642 | } |
643 | |
644 | //===----------------------------------------------------------------------===// |
645 | // Templated helper functions. |
646 | //===----------------------------------------------------------------------===// |
647 | |
648 | /// Apply a permutation from `map` to `source` and return the result. |
649 | template <typename T> |
650 | SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) { |
651 | assert(map.isProjectedPermutation()); |
652 | assert(map.getNumInputs() == source.size()); |
653 | SmallVector<T> result; |
654 | result.reserve(map.getNumResults()); |
655 | for (AffineExpr expr : map.getResults()) { |
656 | if (auto dimExpr = dyn_cast<AffineDimExpr>(Val&: expr)) { |
657 | result.push_back(source[dimExpr.getPosition()]); |
658 | } else if (auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr)) { |
659 | assert(constExpr.getValue() == 0 && |
660 | "Unexpected constant in projected permutation map" ); |
661 | result.push_back(0); |
662 | } else { |
663 | llvm_unreachable("Unexpected result in projected permutation map" ); |
664 | } |
665 | } |
666 | return result; |
667 | } |
668 | |
669 | /// Calculates maximum dimension and symbol positions from the expressions |
670 | /// in `exprsLists` and stores them in `maxDim` and `maxSym` respectively. |
671 | template <typename AffineExprContainer> |
672 | static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList, |
673 | int64_t &maxDim, int64_t &maxSym) { |
674 | for (const auto &exprs : exprsList) { |
675 | for (auto expr : exprs) { |
676 | expr.walk([&maxDim, &maxSym](AffineExpr e) { |
677 | if (auto d = dyn_cast<AffineDimExpr>(Val&: e)) |
678 | maxDim = std::max(a: maxDim, b: static_cast<int64_t>(d.getPosition())); |
679 | if (auto s = dyn_cast<AffineSymbolExpr>(Val&: e)) |
680 | maxSym = std::max(a: maxSym, b: static_cast<int64_t>(s.getPosition())); |
681 | }); |
682 | } |
683 | } |
684 | } |
685 | |
686 | } // namespace mlir |
687 | |
688 | namespace llvm { |
689 | |
690 | // AffineExpr hash just like pointers |
691 | template <> |
692 | struct DenseMapInfo<mlir::AffineMap> { |
693 | static mlir::AffineMap getEmptyKey() { |
694 | auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey(); |
695 | return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer)); |
696 | } |
697 | static mlir::AffineMap getTombstoneKey() { |
698 | auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey(); |
699 | return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer)); |
700 | } |
701 | static unsigned getHashValue(mlir::AffineMap val) { |
702 | return mlir::hash_value(arg: val); |
703 | } |
704 | static bool isEqual(mlir::AffineMap LHS, mlir::AffineMap RHS) { |
705 | return LHS == RHS; |
706 | } |
707 | }; |
708 | |
709 | } // namespace llvm |
710 | |
711 | #endif // MLIR_IR_AFFINEMAP_H |
712 | |