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

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