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
26namespace llvm {
27class SmallBitVector;
28} // namespace llvm
29
30namespace mlir {
31
32namespace detail {
33struct AffineMapStorage;
34} // namespace detail
35
36class Attribute;
37class Builder;
38struct LogicalResult;
39class OpFoldResult;
40class 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.
47class AffineMap {
48public:
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
392private:
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.
400inline ::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.
405struct MutableAffineMap {
406public:
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
434private:
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.
445AffineMap simplifyAffineMap(AffineMap map);
446
447/// Drop the dims that are listed in `unusedDims`.
448AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims);
449
450/// Drop the dims that are not used.
451AffineMap 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.
456SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps);
457
458/// Drop the symbols that are listed in `unusedSymbols`.
459AffineMap compressSymbols(AffineMap map,
460 const llvm::SmallBitVector &unusedSymbols);
461
462/// Drop the symbols that are not used.
463AffineMap 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.
468SmallVector<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`.
472AffineMap 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`.
478AffineMap 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/// ```
514AffineMap 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/// ```
567AffineMap 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/// ```
591AffineMap 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.
610AffineMap 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.
615AffineMap 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`.
620AffineMap 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`.
627llvm::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`.
631llvm::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)`.
636AffineMap expandDimsToRank(AffineMap map, int64_t rank,
637 const llvm::SmallBitVector &projectedDimensions);
638
639inline 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.
649template <typename T>
650SmallVector<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.
671template <typename AffineExprContainer>
672static 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
688namespace llvm {
689
690// AffineExpr hash just like pointers
691template <>
692struct 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

source code of mlir/include/mlir/IR/AffineMap.h