1//===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//
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#include "mlir/IR/AffineMap.h"
10#include "AffineMapDetail.h"
11#include "mlir/IR/AffineExpr.h"
12#include "mlir/IR/Builders.h"
13#include "mlir/IR/BuiltinAttributes.h"
14#include "mlir/IR/BuiltinTypes.h"
15#include "llvm/ADT/STLExtras.h"
16#include "llvm/ADT/SmallBitVector.h"
17#include "llvm/ADT/SmallVector.h"
18#include "llvm/Support/MathExtras.h"
19#include <numeric>
20#include <optional>
21#include <type_traits>
22
23using namespace mlir;
24
25using llvm::divideCeilSigned;
26using llvm::divideFloorSigned;
27using llvm::mod;
28
29namespace {
30
31// AffineExprConstantFolder evaluates an affine expression using constant
32// operands passed in 'operandConsts'. Returns an IntegerAttr attribute
33// representing the constant value of the affine expression evaluated on
34// constant 'operandConsts', or nullptr if it can't be folded.
35class AffineExprConstantFolder {
36public:
37 AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
38 : numDims(numDims), operandConsts(operandConsts) {}
39
40 /// Attempt to constant fold the specified affine expr, or return null on
41 /// failure.
42 IntegerAttr constantFold(AffineExpr expr) {
43 if (auto result = constantFoldImpl(expr))
44 return IntegerAttr::get(type: IndexType::get(context: expr.getContext()), value: *result);
45 return nullptr;
46 }
47
48 bool hasPoison() const { return hasPoison_; }
49
50private:
51 std::optional<int64_t> constantFoldImpl(AffineExpr expr) {
52 switch (expr.getKind()) {
53 case AffineExprKind::Add:
54 return constantFoldBinExpr(
55 expr, op: [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
56 case AffineExprKind::Mul:
57 return constantFoldBinExpr(
58 expr, op: [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
59 case AffineExprKind::Mod:
60 return constantFoldBinExpr(
61 expr, op: [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
62 if (rhs < 1) {
63 hasPoison_ = true;
64 return std::nullopt;
65 }
66 return mod(Numerator: lhs, Denominator: rhs);
67 });
68 case AffineExprKind::FloorDiv:
69 return constantFoldBinExpr(
70 expr, op: [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
71 if (rhs == 0) {
72 hasPoison_ = true;
73 return std::nullopt;
74 }
75 return divideFloorSigned(Numerator: lhs, Denominator: rhs);
76 });
77 case AffineExprKind::CeilDiv:
78 return constantFoldBinExpr(
79 expr, op: [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
80 if (rhs == 0) {
81 hasPoison_ = true;
82 return std::nullopt;
83 }
84 return divideCeilSigned(Numerator: lhs, Denominator: rhs);
85 });
86 case AffineExprKind::Constant:
87 return cast<AffineConstantExpr>(Val&: expr).getValue();
88 case AffineExprKind::DimId:
89 if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(
90 Val: operandConsts[cast<AffineDimExpr>(Val&: expr).getPosition()]))
91 return attr.getInt();
92 return std::nullopt;
93 case AffineExprKind::SymbolId:
94 if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(
95 Val: operandConsts[numDims +
96 cast<AffineSymbolExpr>(Val&: expr).getPosition()]))
97 return attr.getInt();
98 return std::nullopt;
99 }
100 llvm_unreachable("Unknown AffineExpr");
101 }
102
103 // TODO: Change these to operate on APInts too.
104 std::optional<int64_t> constantFoldBinExpr(
105 AffineExpr expr,
106 llvm::function_ref<std::optional<int64_t>(int64_t, int64_t)> op) {
107 auto binOpExpr = cast<AffineBinaryOpExpr>(Val&: expr);
108 if (auto lhs = constantFoldImpl(expr: binOpExpr.getLHS()))
109 if (auto rhs = constantFoldImpl(expr: binOpExpr.getRHS()))
110 return op(*lhs, *rhs);
111 return std::nullopt;
112 }
113
114 // The number of dimension operands in AffineMap containing this expression.
115 unsigned numDims;
116 // The constant valued operands used to evaluate this AffineExpr.
117 ArrayRef<Attribute> operandConsts;
118 bool hasPoison_{false};
119};
120
121} // namespace
122
123/// Returns a single constant result affine map.
124AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
125 return get(/*dimCount=*/0, /*symbolCount=*/0,
126 result: {getAffineConstantExpr(constant: val, context)});
127}
128
129/// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
130/// minor dimensions.
131AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
132 MLIRContext *context) {
133 assert(dims >= results && "Dimension mismatch");
134 auto id = AffineMap::getMultiDimIdentityMap(numDims: dims, context);
135 return AffineMap::get(dimCount: dims, symbolCount: 0, results: id.getResults().take_back(N: results), context);
136}
137
138AffineMap AffineMap::getFilteredIdentityMap(
139 MLIRContext *ctx, unsigned numDims,
140 llvm::function_ref<bool(AffineDimExpr)> keepDimFilter) {
141 auto identityMap = getMultiDimIdentityMap(numDims, context: ctx);
142
143 // Apply filter to results.
144 llvm::SmallBitVector dropDimResults(numDims);
145 for (auto [idx, resultExpr] : llvm::enumerate(First: identityMap.getResults()))
146 dropDimResults[idx] = !keepDimFilter(cast<AffineDimExpr>(Val: resultExpr));
147
148 return identityMap.dropResults(positions: dropDimResults);
149}
150
151bool AffineMap::isMinorIdentity() const {
152 return getNumDims() >= getNumResults() &&
153 *this ==
154 getMinorIdentityMap(dims: getNumDims(), results: getNumResults(), context: getContext());
155}
156
157SmallVector<unsigned> AffineMap::getBroadcastDims() const {
158 SmallVector<unsigned> broadcastedDims;
159 for (const auto &[resIdx, expr] : llvm::enumerate(First: getResults())) {
160 if (auto constExpr = dyn_cast<AffineConstantExpr>(Val: expr)) {
161 if (constExpr.getValue() != 0)
162 continue;
163 broadcastedDims.push_back(Elt: resIdx);
164 }
165 }
166
167 return broadcastedDims;
168}
169
170/// Returns true if this affine map is a minor identity up to broadcasted
171/// dimensions which are indicated by value 0 in the result.
172bool AffineMap::isMinorIdentityWithBroadcasting(
173 SmallVectorImpl<unsigned> *broadcastedDims) const {
174 if (broadcastedDims)
175 broadcastedDims->clear();
176 if (getNumDims() < getNumResults())
177 return false;
178 unsigned suffixStart = getNumDims() - getNumResults();
179 for (const auto &idxAndExpr : llvm::enumerate(First: getResults())) {
180 unsigned resIdx = idxAndExpr.index();
181 AffineExpr expr = idxAndExpr.value();
182 if (auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr)) {
183 // Each result may be either a constant 0 (broadcasted dimension).
184 if (constExpr.getValue() != 0)
185 return false;
186 if (broadcastedDims)
187 broadcastedDims->push_back(Elt: resIdx);
188 } else if (auto dimExpr = dyn_cast<AffineDimExpr>(Val&: expr)) {
189 // Or it may be the input dimension corresponding to this result position.
190 if (dimExpr.getPosition() != suffixStart + resIdx)
191 return false;
192 } else {
193 return false;
194 }
195 }
196 return true;
197}
198
199/// Return true if this affine map can be converted to a minor identity with
200/// broadcast by doing a permute. Return a permutation (there may be
201/// several) to apply to get to a minor identity with broadcasts.
202/// Ex:
203/// * (d0, d1, d2) -> (0, d1) maps to minor identity (d1, 0 = d2) with
204/// perm = [1, 0] and broadcast d2
205/// * (d0, d1, d2) -> (d0, 0) cannot be mapped to a minor identity by
206/// permutation + broadcast
207/// * (d0, d1, d2, d3) -> (0, d1, d3) maps to minor identity (d1, 0 = d2, d3)
208/// with perm = [1, 0, 2] and broadcast d2
209/// * (d0, d1) -> (d1, 0, 0, d0) maps to minor identity (d0, d1) with extra
210/// leading broadcat dimensions. The map returned would be (0, 0, d0, d1) with
211/// perm = [3, 0, 1, 2]
212bool AffineMap::isPermutationOfMinorIdentityWithBroadcasting(
213 SmallVectorImpl<unsigned> &permutedDims) const {
214 unsigned projectionStart =
215 getNumResults() < getNumInputs() ? getNumInputs() - getNumResults() : 0;
216 permutedDims.clear();
217 SmallVector<unsigned> broadcastDims;
218 permutedDims.resize(N: getNumResults(), NV: 0);
219 // If there are more results than input dimensions we want the new map to
220 // start with broadcast dimensions in order to be a minor identity with
221 // broadcasting.
222 unsigned leadingBroadcast =
223 getNumResults() > getNumInputs() ? getNumResults() - getNumInputs() : 0;
224 llvm::SmallBitVector dimFound(std::max(a: getNumInputs(), b: getNumResults()),
225 false);
226 for (const auto &idxAndExpr : llvm::enumerate(First: getResults())) {
227 unsigned resIdx = idxAndExpr.index();
228 AffineExpr expr = idxAndExpr.value();
229 // Each result may be either a constant 0 (broadcast dimension) or a
230 // dimension.
231 if (auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr)) {
232 if (constExpr.getValue() != 0)
233 return false;
234 broadcastDims.push_back(Elt: resIdx);
235 } else if (auto dimExpr = dyn_cast<AffineDimExpr>(Val&: expr)) {
236 if (dimExpr.getPosition() < projectionStart)
237 return false;
238 unsigned newPosition =
239 dimExpr.getPosition() - projectionStart + leadingBroadcast;
240 permutedDims[resIdx] = newPosition;
241 dimFound[newPosition] = true;
242 } else {
243 return false;
244 }
245 }
246 // Find a permuation for the broadcast dimension. Since they are broadcasted
247 // any valid permutation is acceptable. We just permute the dim into a slot
248 // without an existing dimension.
249 unsigned pos = 0;
250 for (auto dim : broadcastDims) {
251 while (pos < dimFound.size() && dimFound[pos]) {
252 pos++;
253 }
254 permutedDims[dim] = pos++;
255 }
256 return true;
257}
258
259/// Returns an AffineMap representing a permutation.
260AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
261 MLIRContext *context) {
262 assert(!permutation.empty() &&
263 "Cannot create permutation map from empty permutation vector");
264 const auto *m = llvm::max_element(Range&: permutation);
265 auto permutationMap = getMultiDimMapWithTargets(numDims: *m + 1, targets: permutation, context);
266 assert(permutationMap.isPermutation() && "Invalid permutation vector");
267 return permutationMap;
268}
269AffineMap AffineMap::getPermutationMap(ArrayRef<int64_t> permutation,
270 MLIRContext *context) {
271 SmallVector<unsigned> perm = llvm::map_to_vector(
272 C&: permutation, F: [](int64_t i) { return static_cast<unsigned>(i); });
273 return AffineMap::getPermutationMap(permutation: perm, context);
274}
275
276AffineMap AffineMap::getMultiDimMapWithTargets(unsigned numDims,
277 ArrayRef<unsigned> targets,
278 MLIRContext *context) {
279 SmallVector<AffineExpr, 4> affExprs;
280 for (unsigned t : targets)
281 affExprs.push_back(Elt: getAffineDimExpr(position: t, context));
282 AffineMap result = AffineMap::get(/*dimCount=*/numDims, /*symbolCount=*/0,
283 results: affExprs, context);
284 return result;
285}
286
287/// Creates an affine map each for each list of AffineExpr's in `exprsList`
288/// while inferring the right number of dimensional and symbolic inputs needed
289/// based on the maximum dimensional and symbolic identifier appearing in the
290/// expressions.
291template <typename AffineExprContainer>
292static SmallVector<AffineMap, 4>
293inferFromExprList(ArrayRef<AffineExprContainer> exprsList,
294 MLIRContext *context) {
295 if (exprsList.empty())
296 return {};
297 int64_t maxDim = -1, maxSym = -1;
298 getMaxDimAndSymbol(exprsList, maxDim, maxSym);
299 SmallVector<AffineMap, 4> maps;
300 maps.reserve(N: exprsList.size());
301 for (const auto &exprs : exprsList)
302 maps.push_back(Elt: AffineMap::get(/*dimCount=*/maxDim + 1,
303 /*symbolCount=*/maxSym + 1, exprs, context));
304 return maps;
305}
306
307SmallVector<AffineMap, 4>
308AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList,
309 MLIRContext *context) {
310 return ::inferFromExprList(exprsList, context);
311}
312
313SmallVector<AffineMap, 4>
314AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList,
315 MLIRContext *context) {
316 return ::inferFromExprList(exprsList, context);
317}
318
319uint64_t AffineMap::getLargestKnownDivisorOfMapExprs() {
320 uint64_t gcd = 0;
321 for (AffineExpr resultExpr : getResults()) {
322 uint64_t thisGcd = resultExpr.getLargestKnownDivisor();
323 gcd = std::gcd(m: gcd, n: thisGcd);
324 }
325 if (gcd == 0)
326 gcd = std::numeric_limits<uint64_t>::max();
327 return gcd;
328}
329
330AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
331 MLIRContext *context) {
332 SmallVector<AffineExpr, 4> dimExprs;
333 dimExprs.reserve(N: numDims);
334 for (unsigned i = 0; i < numDims; ++i)
335 dimExprs.push_back(Elt: mlir::getAffineDimExpr(position: i, context));
336 return get(/*dimCount=*/numDims, /*symbolCount=*/0, results: dimExprs, context);
337}
338
339MLIRContext *AffineMap::getContext() const { return map->context; }
340
341bool AffineMap::isIdentity() const {
342 if (getNumDims() != getNumResults())
343 return false;
344 ArrayRef<AffineExpr> results = getResults();
345 for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
346 auto expr = dyn_cast<AffineDimExpr>(Val: results[i]);
347 if (!expr || expr.getPosition() != i)
348 return false;
349 }
350 return true;
351}
352
353bool AffineMap::isSymbolIdentity() const {
354 if (getNumSymbols() != getNumResults())
355 return false;
356 ArrayRef<AffineExpr> results = getResults();
357 for (unsigned i = 0, numSymbols = getNumSymbols(); i < numSymbols; ++i) {
358 auto expr = dyn_cast<AffineDimExpr>(Val: results[i]);
359 if (!expr || expr.getPosition() != i)
360 return false;
361 }
362 return true;
363}
364
365bool AffineMap::isEmpty() const {
366 return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
367}
368
369bool AffineMap::isSingleConstant() const {
370 return getNumResults() == 1 && isa<AffineConstantExpr>(Val: getResult(idx: 0));
371}
372
373bool AffineMap::isConstant() const {
374 return llvm::all_of(Range: getResults(), P: llvm::IsaPred<AffineConstantExpr>);
375}
376
377int64_t AffineMap::getSingleConstantResult() const {
378 assert(isSingleConstant() && "map must have a single constant result");
379 return cast<AffineConstantExpr>(Val: getResult(idx: 0)).getValue();
380}
381
382SmallVector<int64_t> AffineMap::getConstantResults() const {
383 assert(isConstant() && "map must have only constant results");
384 SmallVector<int64_t> result;
385 for (auto expr : getResults())
386 result.emplace_back(Args: cast<AffineConstantExpr>(Val&: expr).getValue());
387 return result;
388}
389
390unsigned AffineMap::getNumDims() const {
391 assert(map && "uninitialized map storage");
392 return map->numDims;
393}
394unsigned AffineMap::getNumSymbols() const {
395 assert(map && "uninitialized map storage");
396 return map->numSymbols;
397}
398unsigned AffineMap::getNumResults() const { return getResults().size(); }
399unsigned AffineMap::getNumInputs() const {
400 assert(map && "uninitialized map storage");
401 return map->numDims + map->numSymbols;
402}
403ArrayRef<AffineExpr> AffineMap::getResults() const {
404 assert(map && "uninitialized map storage");
405 return map->results();
406}
407AffineExpr AffineMap::getResult(unsigned idx) const {
408 return getResults()[idx];
409}
410
411unsigned AffineMap::getDimPosition(unsigned idx) const {
412 return cast<AffineDimExpr>(Val: getResult(idx)).getPosition();
413}
414
415std::optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const {
416 if (!isa<AffineDimExpr>(Val: input))
417 return std::nullopt;
418
419 for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {
420 if (getResult(idx: i) == input)
421 return i;
422 }
423
424 return std::nullopt;
425}
426
427/// Folds the results of the application of an affine map on the provided
428/// operands to a constant if possible. Returns false if the folding happens,
429/// true otherwise.
430LogicalResult AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
431 SmallVectorImpl<Attribute> &results,
432 bool *hasPoison) const {
433 // Attempt partial folding.
434 SmallVector<int64_t, 2> integers;
435 partialConstantFold(operandConstants, results: &integers, hasPoison);
436
437 // If all expressions folded to a constant, populate results with attributes
438 // containing those constants.
439 if (integers.empty())
440 return failure();
441
442 auto range = llvm::map_range(C&: integers, F: [this](int64_t i) {
443 return IntegerAttr::get(type: IndexType::get(context: getContext()), value: i);
444 });
445 results.append(in_start: range.begin(), in_end: range.end());
446 return success();
447}
448
449AffineMap AffineMap::partialConstantFold(ArrayRef<Attribute> operandConstants,
450 SmallVectorImpl<int64_t> *results,
451 bool *hasPoison) const {
452 assert(getNumInputs() == operandConstants.size());
453
454 // Fold each of the result expressions.
455 AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);
456 SmallVector<AffineExpr, 4> exprs;
457 exprs.reserve(N: getNumResults());
458
459 for (auto expr : getResults()) {
460 auto folded = exprFolder.constantFold(expr);
461 if (exprFolder.hasPoison() && hasPoison) {
462 *hasPoison = true;
463 return {};
464 }
465 // If did not fold to a constant, keep the original expression, and clear
466 // the integer results vector.
467 if (folded) {
468 exprs.push_back(
469 Elt: getAffineConstantExpr(constant: folded.getInt(), context: folded.getContext()));
470 if (results)
471 results->push_back(Elt: folded.getInt());
472 } else {
473 exprs.push_back(Elt: expr);
474 if (results) {
475 results->clear();
476 results = nullptr;
477 }
478 }
479 }
480
481 return get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: exprs, context: getContext());
482}
483
484/// Walk all of the AffineExpr's in this mapping. Each node in an expression
485/// tree is visited in postorder.
486void AffineMap::walkExprs(llvm::function_ref<void(AffineExpr)> callback) const {
487 for (auto expr : getResults())
488 expr.walk(callback);
489}
490
491/// This method substitutes any uses of dimensions and symbols (e.g.
492/// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
493/// expression mapping. Because this can be used to eliminate dims and
494/// symbols, the client needs to specify the number of dims and symbols in
495/// the result. The returned map always has the same number of results.
496AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
497 ArrayRef<AffineExpr> symReplacements,
498 unsigned numResultDims,
499 unsigned numResultSyms) const {
500 SmallVector<AffineExpr, 8> results;
501 results.reserve(N: getNumResults());
502 for (auto expr : getResults())
503 results.push_back(
504 Elt: expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
505 return get(dimCount: numResultDims, symbolCount: numResultSyms, results, context: getContext());
506}
507
508/// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to
509/// each of the results and return a new AffineMap with the new results and
510/// with the specified number of dims and symbols.
511AffineMap AffineMap::replace(AffineExpr expr, AffineExpr replacement,
512 unsigned numResultDims,
513 unsigned numResultSyms) const {
514 SmallVector<AffineExpr, 4> newResults;
515 newResults.reserve(N: getNumResults());
516 for (AffineExpr e : getResults())
517 newResults.push_back(Elt: e.replace(expr, replacement));
518 return AffineMap::get(dimCount: numResultDims, symbolCount: numResultSyms, results: newResults, context: getContext());
519}
520
521/// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
522/// results and return a new AffineMap with the new results and with the
523/// specified number of dims and symbols.
524AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map,
525 unsigned numResultDims,
526 unsigned numResultSyms) const {
527 SmallVector<AffineExpr, 4> newResults;
528 newResults.reserve(N: getNumResults());
529 for (AffineExpr e : getResults())
530 newResults.push_back(Elt: e.replace(map));
531 return AffineMap::get(dimCount: numResultDims, symbolCount: numResultSyms, results: newResults, context: getContext());
532}
533
534AffineMap
535AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
536 SmallVector<AffineExpr, 4> newResults;
537 newResults.reserve(N: getNumResults());
538 for (AffineExpr e : getResults())
539 newResults.push_back(Elt: e.replace(map));
540 return AffineMap::inferFromExprList(exprsList: newResults, context: getContext()).front();
541}
542
543AffineMap AffineMap::dropResults(const llvm::SmallBitVector &positions) const {
544 auto exprs = llvm::to_vector<4>(Range: getResults());
545 // TODO: this is a pretty terrible API .. is there anything better?
546 for (auto pos = positions.find_last(); pos != -1;
547 pos = positions.find_prev(PriorTo: pos))
548 exprs.erase(CI: exprs.begin() + pos);
549 return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: exprs, context: getContext());
550}
551
552AffineMap AffineMap::compose(AffineMap map) const {
553 assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
554 // Prepare `map` by concatenating the symbols and rewriting its exprs.
555 unsigned numDims = map.getNumDims();
556 unsigned numSymbolsThisMap = getNumSymbols();
557 unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
558 SmallVector<AffineExpr, 8> newDims(numDims);
559 for (unsigned idx = 0; idx < numDims; ++idx) {
560 newDims[idx] = getAffineDimExpr(position: idx, context: getContext());
561 }
562 SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap);
563 for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
564 newSymbols[idx - numSymbolsThisMap] =
565 getAffineSymbolExpr(position: idx, context: getContext());
566 }
567 auto newMap =
568 map.replaceDimsAndSymbols(dimReplacements: newDims, symReplacements: newSymbols, numResultDims: numDims, numResultSyms: numSymbols);
569 SmallVector<AffineExpr, 8> exprs;
570 exprs.reserve(N: getResults().size());
571 for (auto expr : getResults())
572 exprs.push_back(Elt: expr.compose(map: newMap));
573 return AffineMap::get(dimCount: numDims, symbolCount: numSymbols, results: exprs, context: map.getContext());
574}
575
576SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const {
577 assert(getNumSymbols() == 0 && "Expected symbol-less map");
578 SmallVector<AffineExpr, 4> exprs;
579 MLIRContext *ctx = getContext();
580 for (int64_t value : values)
581 exprs.push_back(Elt: getAffineConstantExpr(constant: value, context: ctx));
582 SmallVector<int64_t, 4> res;
583 res.reserve(N: getNumResults());
584 for (auto e : getResults())
585 res.push_back(Elt: cast<AffineConstantExpr>(Val: e.replaceDims(dimReplacements: exprs)).getValue());
586 return res;
587}
588
589size_t AffineMap::getNumOfZeroResults() const {
590 size_t res = 0;
591 for (auto expr : getResults()) {
592 auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr);
593 if (constExpr && constExpr.getValue() == 0)
594 res++;
595 }
596
597 return res;
598}
599
600AffineMap AffineMap::dropZeroResults() {
601 SmallVector<AffineExpr> newExprs;
602
603 for (auto expr : getResults()) {
604 auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr);
605 if (!constExpr || constExpr.getValue() != 0)
606 newExprs.push_back(Elt: expr);
607 }
608 return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: newExprs, context: getContext());
609}
610
611bool AffineMap::isProjectedPermutation(bool allowZeroInResults) const {
612 if (getNumSymbols() > 0)
613 return false;
614
615 // Having more results than inputs means that results have duplicated dims or
616 // zeros that can't be mapped to input dims.
617 if (getNumResults() > getNumInputs())
618 return false;
619
620 SmallVector<bool, 8> seen(getNumInputs(), false);
621 // A projected permutation can have, at most, only one instance of each input
622 // dimension in the result expressions. Zeros are allowed as long as the
623 // number of result expressions is lower or equal than the number of input
624 // expressions.
625 for (auto expr : getResults()) {
626 if (auto dim = dyn_cast<AffineDimExpr>(Val&: expr)) {
627 if (seen[dim.getPosition()])
628 return false;
629 seen[dim.getPosition()] = true;
630 } else {
631 auto constExpr = dyn_cast<AffineConstantExpr>(Val&: expr);
632 if (!allowZeroInResults || !constExpr || constExpr.getValue() != 0)
633 return false;
634 }
635 }
636
637 // Results are either dims or zeros and zeros can be mapped to input dims.
638 return true;
639}
640
641bool AffineMap::isPermutation() const {
642 if (getNumDims() != getNumResults())
643 return false;
644 return isProjectedPermutation();
645}
646
647AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {
648 SmallVector<AffineExpr, 4> exprs;
649 exprs.reserve(N: resultPos.size());
650 for (auto idx : resultPos)
651 exprs.push_back(Elt: getResult(idx));
652 return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(), results: exprs, context: getContext());
653}
654
655AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const {
656 return AffineMap::get(dimCount: getNumDims(), symbolCount: getNumSymbols(),
657 results: getResults().slice(N: start, M: length), context: getContext());
658}
659
660AffineMap AffineMap::getMajorSubMap(unsigned numResults) const {
661 if (numResults == 0)
662 return AffineMap();
663 if (numResults > getNumResults())
664 return *this;
665 return getSliceMap(start: 0, length: numResults);
666}
667
668AffineMap AffineMap::getMinorSubMap(unsigned numResults) const {
669 if (numResults == 0)
670 return AffineMap();
671 if (numResults > getNumResults())
672 return *this;
673 return getSliceMap(start: getNumResults() - numResults, length: numResults);
674}
675
676/// Implementation detail to compress multiple affine maps with a compressionFun
677/// that is expected to be either compressUnusedDims or compressUnusedSymbols.
678/// The implementation keeps track of num dims and symbols across the different
679/// affine maps.
680static SmallVector<AffineMap> compressUnusedListImpl(
681 ArrayRef<AffineMap> maps,
682 llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
683 if (maps.empty())
684 return SmallVector<AffineMap>();
685 SmallVector<AffineExpr> allExprs;
686 allExprs.reserve(N: maps.size() * maps.front().getNumResults());
687 unsigned numDims = maps.front().getNumDims(),
688 numSymbols = maps.front().getNumSymbols();
689 for (auto m : maps) {
690 assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&
691 "expected maps with same num dims and symbols");
692 llvm::append_range(C&: allExprs, R: m.getResults());
693 }
694 AffineMap unifiedMap = compressionFun(
695 AffineMap::get(dimCount: numDims, symbolCount: numSymbols, results: allExprs, context: maps.front().getContext()));
696 unsigned unifiedNumDims = unifiedMap.getNumDims(),
697 unifiedNumSymbols = unifiedMap.getNumSymbols();
698 ArrayRef<AffineExpr> unifiedResults = unifiedMap.getResults();
699 SmallVector<AffineMap> res;
700 res.reserve(N: maps.size());
701 for (auto m : maps) {
702 res.push_back(Elt: AffineMap::get(dimCount: unifiedNumDims, symbolCount: unifiedNumSymbols,
703 results: unifiedResults.take_front(N: m.getNumResults()),
704 context: m.getContext()));
705 unifiedResults = unifiedResults.drop_front(N: m.getNumResults());
706 }
707 return res;
708}
709
710AffineMap mlir::compressDims(AffineMap map,
711 const llvm::SmallBitVector &unusedDims) {
712 return projectDims(map, projectedDimensions: unusedDims, /*compressDimsFlag=*/true);
713}
714
715AffineMap mlir::compressUnusedDims(AffineMap map) {
716 return compressDims(map, unusedDims: getUnusedDimsBitVector(maps: {map}));
717}
718
719SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) {
720 return compressUnusedListImpl(
721 maps, compressionFun: [](AffineMap m) { return compressUnusedDims(map: m); });
722}
723
724AffineMap mlir::compressSymbols(AffineMap map,
725 const llvm::SmallBitVector &unusedSymbols) {
726 return projectSymbols(map, projectedSymbols: unusedSymbols, /*compressSymbolsFlag=*/true);
727}
728
729AffineMap mlir::compressUnusedSymbols(AffineMap map) {
730 return compressSymbols(map, unusedSymbols: getUnusedSymbolsBitVector(maps: {map}));
731}
732
733SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {
734 return compressUnusedListImpl(
735 maps, compressionFun: [](AffineMap m) { return compressUnusedSymbols(map: m); });
736}
737
738AffineMap mlir::foldAttributesIntoMap(Builder &b, AffineMap map,
739 ArrayRef<OpFoldResult> operands,
740 SmallVector<Value> &remainingValues) {
741 SmallVector<AffineExpr> dimReplacements, symReplacements;
742 int64_t numDims = 0;
743 for (int64_t i = 0; i < map.getNumDims(); ++i) {
744 if (auto attr = dyn_cast<Attribute>(Val: operands[i])) {
745 dimReplacements.push_back(
746 Elt: b.getAffineConstantExpr(constant: cast<IntegerAttr>(Val&: attr).getInt()));
747 } else {
748 dimReplacements.push_back(Elt: b.getAffineDimExpr(position: numDims++));
749 remainingValues.push_back(Elt: cast<Value>(Val: operands[i]));
750 }
751 }
752 int64_t numSymbols = 0;
753 for (int64_t i = 0; i < map.getNumSymbols(); ++i) {
754 if (auto attr = dyn_cast<Attribute>(Val: operands[i + map.getNumDims()])) {
755 symReplacements.push_back(
756 Elt: b.getAffineConstantExpr(constant: cast<IntegerAttr>(Val&: attr).getInt()));
757 } else {
758 symReplacements.push_back(Elt: b.getAffineSymbolExpr(position: numSymbols++));
759 remainingValues.push_back(Elt: cast<Value>(Val: operands[i + map.getNumDims()]));
760 }
761 }
762 return map.replaceDimsAndSymbols(dimReplacements, symReplacements, numResultDims: numDims,
763 numResultSyms: numSymbols);
764}
765
766AffineMap mlir::simplifyAffineMap(AffineMap map) {
767 SmallVector<AffineExpr, 8> exprs;
768 for (auto e : map.getResults()) {
769 exprs.push_back(
770 Elt: simplifyAffineExpr(expr: e, numDims: map.getNumDims(), numSymbols: map.getNumSymbols()));
771 }
772 return AffineMap::get(dimCount: map.getNumDims(), symbolCount: map.getNumSymbols(), results: exprs,
773 context: map.getContext());
774}
775
776AffineMap mlir::removeDuplicateExprs(AffineMap map) {
777 auto results = map.getResults();
778 SmallVector<AffineExpr, 4> uniqueExprs(results);
779 uniqueExprs.erase(CS: llvm::unique(R&: uniqueExprs), CE: uniqueExprs.end());
780 return AffineMap::get(dimCount: map.getNumDims(), symbolCount: map.getNumSymbols(), results: uniqueExprs,
781 context: map.getContext());
782}
783
784AffineMap mlir::inversePermutation(AffineMap map) {
785 if (map.isEmpty())
786 return map;
787 assert(map.getNumSymbols() == 0 && "expected map without symbols");
788 SmallVector<AffineExpr, 4> exprs(map.getNumDims());
789 for (const auto &en : llvm::enumerate(First: map.getResults())) {
790 auto expr = en.value();
791 // Skip non-permutations.
792 if (auto d = dyn_cast<AffineDimExpr>(Val&: expr)) {
793 if (exprs[d.getPosition()])
794 continue;
795 exprs[d.getPosition()] = getAffineDimExpr(position: en.index(), context: d.getContext());
796 }
797 }
798 SmallVector<AffineExpr, 4> seenExprs;
799 seenExprs.reserve(N: map.getNumDims());
800 for (auto expr : exprs)
801 if (expr)
802 seenExprs.push_back(Elt: expr);
803 if (seenExprs.size() != map.getNumInputs())
804 return AffineMap();
805 return AffineMap::get(dimCount: map.getNumResults(), symbolCount: 0, results: seenExprs, context: map.getContext());
806}
807
808AffineMap mlir::inverseAndBroadcastProjectedPermutation(AffineMap map) {
809 assert(map.isProjectedPermutation(/*allowZeroInResults=*/true));
810 MLIRContext *context = map.getContext();
811 AffineExpr zero = mlir::getAffineConstantExpr(constant: 0, context);
812 // Start with all the results as 0.
813 SmallVector<AffineExpr, 4> exprs(map.getNumInputs(), zero);
814 for (unsigned i : llvm::seq(Begin: unsigned(0), End: map.getNumResults())) {
815 // Skip zeros from input map. 'exprs' is already initialized to zero.
816 if (auto constExpr = dyn_cast<AffineConstantExpr>(Val: map.getResult(idx: i))) {
817 assert(constExpr.getValue() == 0 &&
818 "Unexpected constant in projected permutation");
819 (void)constExpr;
820 continue;
821 }
822
823 // Reverse each dimension existing in the original map result.
824 exprs[map.getDimPosition(idx: i)] = getAffineDimExpr(position: i, context);
825 }
826 return AffineMap::get(dimCount: map.getNumResults(), /*symbolCount=*/0, results: exprs, context);
827}
828
829AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps,
830 MLIRContext *context) {
831 if (maps.empty())
832 return AffineMap::get(context);
833 unsigned numResults = 0, numDims = 0, numSymbols = 0;
834 for (auto m : maps)
835 numResults += m.getNumResults();
836 SmallVector<AffineExpr, 8> results;
837 results.reserve(N: numResults);
838 for (auto m : maps) {
839 for (auto res : m.getResults())
840 results.push_back(Elt: res.shiftSymbols(numSymbols: m.getNumSymbols(), shift: numSymbols));
841
842 numSymbols += m.getNumSymbols();
843 numDims = std::max(a: m.getNumDims(), b: numDims);
844 }
845 return AffineMap::get(dimCount: numDims, symbolCount: numSymbols, results, context);
846}
847
848/// Common implementation to project out dimensions or symbols from an affine
849/// map based on the template type.
850/// Additionally, if 'compress' is true, the projected out dimensions or symbols
851/// are also dropped from the resulting map.
852template <typename AffineDimOrSymExpr>
853static AffineMap projectCommonImpl(AffineMap map,
854 const llvm::SmallBitVector &toProject,
855 bool compress) {
856 static_assert(llvm::is_one_of<AffineDimOrSymExpr, AffineDimExpr,
857 AffineSymbolExpr>::value,
858 "expected AffineDimExpr or AffineSymbolExpr");
859
860 constexpr bool isDim = std::is_same<AffineDimOrSymExpr, AffineDimExpr>::value;
861 int64_t numDimOrSym = (isDim) ? map.getNumDims() : map.getNumSymbols();
862 SmallVector<AffineExpr> replacements;
863 replacements.reserve(N: numDimOrSym);
864
865 auto createNewDimOrSym = (isDim) ? getAffineDimExpr : getAffineSymbolExpr;
866
867 using replace_fn_ty =
868 std::function<AffineExpr(AffineExpr, ArrayRef<AffineExpr>)>;
869 replace_fn_ty replaceDims = [](AffineExpr e,
870 ArrayRef<AffineExpr> replacements) {
871 return e.replaceDims(dimReplacements: replacements);
872 };
873 replace_fn_ty replaceSymbols = [](AffineExpr e,
874 ArrayRef<AffineExpr> replacements) {
875 return e.replaceSymbols(symReplacements: replacements);
876 };
877 replace_fn_ty replaceNewDimOrSym = (isDim) ? replaceDims : replaceSymbols;
878
879 MLIRContext *context = map.getContext();
880 int64_t newNumDimOrSym = 0;
881 for (unsigned dimOrSym = 0; dimOrSym < numDimOrSym; ++dimOrSym) {
882 if (toProject.test(Idx: dimOrSym)) {
883 replacements.push_back(Elt: getAffineConstantExpr(constant: 0, context));
884 continue;
885 }
886 int64_t newPos = compress ? newNumDimOrSym++ : dimOrSym;
887 replacements.push_back(Elt: createNewDimOrSym(newPos, context));
888 }
889 SmallVector<AffineExpr> resultExprs;
890 resultExprs.reserve(N: map.getNumResults());
891 for (auto e : map.getResults())
892 resultExprs.push_back(Elt: replaceNewDimOrSym(e, replacements));
893
894 int64_t numDims = (compress && isDim) ? newNumDimOrSym : map.getNumDims();
895 int64_t numSyms = (compress && !isDim) ? newNumDimOrSym : map.getNumSymbols();
896 return AffineMap::get(dimCount: numDims, symbolCount: numSyms, results: resultExprs, context);
897}
898
899AffineMap mlir::projectDims(AffineMap map,
900 const llvm::SmallBitVector &projectedDimensions,
901 bool compressDimsFlag) {
902 return projectCommonImpl<AffineDimExpr>(map, toProject: projectedDimensions,
903 compress: compressDimsFlag);
904}
905
906AffineMap mlir::projectSymbols(AffineMap map,
907 const llvm::SmallBitVector &projectedSymbols,
908 bool compressSymbolsFlag) {
909 return projectCommonImpl<AffineSymbolExpr>(map, toProject: projectedSymbols,
910 compress: compressSymbolsFlag);
911}
912
913AffineMap mlir::getProjectedMap(AffineMap map,
914 const llvm::SmallBitVector &projectedDimensions,
915 bool compressDimsFlag,
916 bool compressSymbolsFlag) {
917 map = projectDims(map, projectedDimensions, compressDimsFlag);
918 if (compressSymbolsFlag)
919 map = compressUnusedSymbols(map);
920 return map;
921}
922
923llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) {
924 unsigned numDims = maps[0].getNumDims();
925 llvm::SmallBitVector numDimsBitVector(numDims, true);
926 for (AffineMap m : maps) {
927 for (unsigned i = 0; i < numDims; ++i) {
928 if (m.isFunctionOfDim(position: i))
929 numDimsBitVector.reset(Idx: i);
930 }
931 }
932 return numDimsBitVector;
933}
934
935llvm::SmallBitVector mlir::getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps) {
936 unsigned numSymbols = maps[0].getNumSymbols();
937 llvm::SmallBitVector numSymbolsBitVector(numSymbols, true);
938 for (AffineMap m : maps) {
939 for (unsigned i = 0; i < numSymbols; ++i) {
940 if (m.isFunctionOfSymbol(position: i))
941 numSymbolsBitVector.reset(Idx: i);
942 }
943 }
944 return numSymbolsBitVector;
945}
946
947AffineMap
948mlir::expandDimsToRank(AffineMap map, int64_t rank,
949 const llvm::SmallBitVector &projectedDimensions) {
950 auto id = AffineMap::getMultiDimIdentityMap(numDims: rank, context: map.getContext());
951 AffineMap proj = id.dropResults(positions: projectedDimensions);
952 return map.compose(map: proj);
953}
954
955//===----------------------------------------------------------------------===//
956// MutableAffineMap.
957//===----------------------------------------------------------------------===//
958
959MutableAffineMap::MutableAffineMap(AffineMap map)
960 : results(map.getResults()), numDims(map.getNumDims()),
961 numSymbols(map.getNumSymbols()), context(map.getContext()) {}
962
963void MutableAffineMap::reset(AffineMap map) {
964 results.clear();
965 numDims = map.getNumDims();
966 numSymbols = map.getNumSymbols();
967 context = map.getContext();
968 llvm::append_range(C&: results, R: map.getResults());
969}
970
971bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
972 return results[idx].isMultipleOf(factor);
973}
974
975// Simplifies the result affine expressions of this map. The expressions
976// have to be pure for the simplification implemented.
977void MutableAffineMap::simplify() {
978 // Simplify each of the results if possible.
979 // TODO: functional-style map
980 for (unsigned i = 0, e = getNumResults(); i < e; i++) {
981 results[i] = simplifyAffineExpr(expr: getResult(idx: i), numDims, numSymbols);
982 }
983}
984
985AffineMap MutableAffineMap::getAffineMap() const {
986 return AffineMap::get(dimCount: numDims, symbolCount: numSymbols, results, context);
987}
988

source code of mlir/lib/IR/AffineMap.cpp