MLIR  17.0.0git
AffineMap.cpp
Go to the documentation of this file.
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"
12 #include "mlir/IR/BuiltinTypes.h"
15 #include "llvm/ADT/SmallBitVector.h"
16 #include "llvm/ADT/SmallSet.h"
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/Support/raw_ostream.h"
19 #include <numeric>
20 #include <optional>
21 
22 using namespace mlir;
23 
24 namespace {
25 
26 // AffineExprConstantFolder evaluates an affine expression using constant
27 // operands passed in 'operandConsts'. Returns an IntegerAttr attribute
28 // representing the constant value of the affine expression evaluated on
29 // constant 'operandConsts', or nullptr if it can't be folded.
30 class AffineExprConstantFolder {
31 public:
32  AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
33  : numDims(numDims), operandConsts(operandConsts) {}
34 
35  /// Attempt to constant fold the specified affine expr, or return null on
36  /// failure.
37  IntegerAttr constantFold(AffineExpr expr) {
38  if (auto result = constantFoldImpl(expr))
39  return IntegerAttr::get(IndexType::get(expr.getContext()), *result);
40  return nullptr;
41  }
42 
43 private:
44  std::optional<int64_t> constantFoldImpl(AffineExpr expr) {
45  switch (expr.getKind()) {
47  return constantFoldBinExpr(
48  expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
50  return constantFoldBinExpr(
51  expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
53  return constantFoldBinExpr(
54  expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); });
56  return constantFoldBinExpr(
57  expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); });
59  return constantFoldBinExpr(
60  expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); });
62  return expr.cast<AffineConstantExpr>().getValue();
64  if (auto attr = operandConsts[expr.cast<AffineDimExpr>().getPosition()]
65  .dyn_cast_or_null<IntegerAttr>())
66  return attr.getInt();
67  return std::nullopt;
69  if (auto attr = operandConsts[numDims +
71  .dyn_cast_or_null<IntegerAttr>())
72  return attr.getInt();
73  return std::nullopt;
74  }
75  llvm_unreachable("Unknown AffineExpr");
76  }
77 
78  // TODO: Change these to operate on APInts too.
79  std::optional<int64_t> constantFoldBinExpr(AffineExpr expr,
80  int64_t (*op)(int64_t, int64_t)) {
81  auto binOpExpr = expr.cast<AffineBinaryOpExpr>();
82  if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))
83  if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))
84  return op(*lhs, *rhs);
85  return std::nullopt;
86  }
87 
88  // The number of dimension operands in AffineMap containing this expression.
89  unsigned numDims;
90  // The constant valued operands used to evaluate this AffineExpr.
91  ArrayRef<Attribute> operandConsts;
92 };
93 
94 } // namespace
95 
96 /// Returns a single constant result affine map.
98  return get(/*dimCount=*/0, /*symbolCount=*/0,
99  {getAffineConstantExpr(val, context)});
100 }
101 
102 /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
103 /// minor dimensions.
104 AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
105  MLIRContext *context) {
106  assert(dims >= results && "Dimension mismatch");
107  auto id = AffineMap::getMultiDimIdentityMap(dims, context);
108  return AffineMap::get(dims, 0, id.getResults().take_back(results), context);
109 }
110 
112  return getNumDims() >= getNumResults() &&
113  *this ==
115 }
116 
117 /// Returns true if this affine map is a minor identity up to broadcasted
118 /// dimensions which are indicated by value 0 in the result.
120  SmallVectorImpl<unsigned> *broadcastedDims) const {
121  if (broadcastedDims)
122  broadcastedDims->clear();
123  if (getNumDims() < getNumResults())
124  return false;
125  unsigned suffixStart = getNumDims() - getNumResults();
126  for (const auto &idxAndExpr : llvm::enumerate(getResults())) {
127  unsigned resIdx = idxAndExpr.index();
128  AffineExpr expr = idxAndExpr.value();
129  if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
130  // Each result may be either a constant 0 (broadcasted dimension).
131  if (constExpr.getValue() != 0)
132  return false;
133  if (broadcastedDims)
134  broadcastedDims->push_back(resIdx);
135  } else if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
136  // Or it may be the input dimension corresponding to this result position.
137  if (dimExpr.getPosition() != suffixStart + resIdx)
138  return false;
139  } else {
140  return false;
141  }
142  }
143  return true;
144 }
145 
146 /// Return true if this affine map can be converted to a minor identity with
147 /// broadcast by doing a permute. Return a permutation (there may be
148 /// several) to apply to get to a minor identity with broadcasts.
149 /// Ex:
150 /// * (d0, d1, d2) -> (0, d1) maps to minor identity (d1, 0 = d2) with
151 /// perm = [1, 0] and broadcast d2
152 /// * (d0, d1, d2) -> (d0, 0) cannot be mapped to a minor identity by
153 /// permutation + broadcast
154 /// * (d0, d1, d2, d3) -> (0, d1, d3) maps to minor identity (d1, 0 = d2, d3)
155 /// with perm = [1, 0, 2] and broadcast d2
156 /// * (d0, d1) -> (d1, 0, 0, d0) maps to minor identity (d0, d1) with extra
157 /// leading broadcat dimensions. The map returned would be (0, 0, d0, d1) with
158 /// perm = [3, 0, 1, 2]
160  SmallVectorImpl<unsigned> &permutedDims) const {
161  unsigned projectionStart =
163  permutedDims.clear();
164  SmallVector<unsigned> broadcastDims;
165  permutedDims.resize(getNumResults(), 0);
166  // If there are more results than input dimensions we want the new map to
167  // start with broadcast dimensions in order to be a minor identity with
168  // broadcasting.
169  unsigned leadingBroadcast =
171  llvm::SmallBitVector dimFound(std::max(getNumInputs(), getNumResults()),
172  false);
173  for (const auto &idxAndExpr : llvm::enumerate(getResults())) {
174  unsigned resIdx = idxAndExpr.index();
175  AffineExpr expr = idxAndExpr.value();
176  // Each result may be either a constant 0 (broadcast dimension) or a
177  // dimension.
178  if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
179  if (constExpr.getValue() != 0)
180  return false;
181  broadcastDims.push_back(resIdx);
182  } else if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
183  if (dimExpr.getPosition() < projectionStart)
184  return false;
185  unsigned newPosition =
186  dimExpr.getPosition() - projectionStart + leadingBroadcast;
187  permutedDims[resIdx] = newPosition;
188  dimFound[newPosition] = true;
189  } else {
190  return false;
191  }
192  }
193  // Find a permuation for the broadcast dimension. Since they are broadcasted
194  // any valid permutation is acceptable. We just permute the dim into a slot
195  // without an existing dimension.
196  unsigned pos = 0;
197  for (auto dim : broadcastDims) {
198  while (pos < dimFound.size() && dimFound[pos]) {
199  pos++;
200  }
201  permutedDims[dim] = pos++;
202  }
203  return true;
204 }
205 
206 /// Returns an AffineMap representing a permutation.
208  MLIRContext *context) {
209  assert(!permutation.empty() &&
210  "Cannot create permutation map from empty permutation vector");
212  for (auto index : permutation)
213  affExprs.push_back(getAffineDimExpr(index, context));
214  const auto *m = std::max_element(permutation.begin(), permutation.end());
215  auto permutationMap = AffineMap::get(*m + 1, 0, affExprs, context);
216  assert(permutationMap.isPermutation() && "Invalid permutation vector");
217  return permutationMap;
218 }
219 
220 template <typename AffineExprContainer>
223  assert(!exprsList.empty());
224  assert(!exprsList[0].empty());
225  auto context = exprsList[0][0].getContext();
226  int64_t maxDim = -1, maxSym = -1;
227  getMaxDimAndSymbol(exprsList, maxDim, maxSym);
229  maps.reserve(exprsList.size());
230  for (const auto &exprs : exprsList)
231  maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
232  /*symbolCount=*/maxSym + 1, exprs, context));
233  return maps;
234 }
235 
238  return ::inferFromExprList(exprsList);
239 }
240 
243  return ::inferFromExprList(exprsList);
244 }
245 
247  uint64_t gcd = 0;
248  for (AffineExpr resultExpr : getResults()) {
249  uint64_t thisGcd = resultExpr.getLargestKnownDivisor();
250  gcd = std::gcd(gcd, thisGcd);
251  }
252  if (gcd == 0)
254  return gcd;
255 }
256 
258  MLIRContext *context) {
260  dimExprs.reserve(numDims);
261  for (unsigned i = 0; i < numDims; ++i)
262  dimExprs.push_back(mlir::getAffineDimExpr(i, context));
263  return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context);
264 }
265 
266 MLIRContext *AffineMap::getContext() const { return map->context; }
267 
268 bool AffineMap::isIdentity() const {
269  if (getNumDims() != getNumResults())
270  return false;
271  ArrayRef<AffineExpr> results = getResults();
272  for (unsigned i = 0, numDims = getNumDims(); i < numDims; ++i) {
273  auto expr = results[i].dyn_cast<AffineDimExpr>();
274  if (!expr || expr.getPosition() != i)
275  return false;
276  }
277  return true;
278 }
279 
281  if (getNumSymbols() != getNumResults())
282  return false;
283  ArrayRef<AffineExpr> results = getResults();
284  for (unsigned i = 0, numSymbols = getNumSymbols(); i < numSymbols; ++i) {
285  auto expr = results[i].dyn_cast<AffineDimExpr>();
286  if (!expr || expr.getPosition() != i)
287  return false;
288  }
289  return true;
290 }
291 
292 bool AffineMap::isEmpty() const {
293  return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
294 }
295 
297  return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
298 }
299 
300 bool AffineMap::isConstant() const {
301  return llvm::all_of(getResults(), [](AffineExpr expr) {
302  return expr.isa<AffineConstantExpr>();
303  });
304 }
305 
307  assert(isSingleConstant() && "map must have a single constant result");
308  return getResult(0).cast<AffineConstantExpr>().getValue();
309 }
310 
312  assert(isConstant() && "map must have only constant results");
313  SmallVector<int64_t> result;
314  for (auto expr : getResults())
315  result.emplace_back(expr.cast<AffineConstantExpr>().getValue());
316  return result;
317 }
318 
319 unsigned AffineMap::getNumDims() const {
320  assert(map && "uninitialized map storage");
321  return map->numDims;
322 }
323 unsigned AffineMap::getNumSymbols() const {
324  assert(map && "uninitialized map storage");
325  return map->numSymbols;
326 }
327 unsigned AffineMap::getNumResults() const { return getResults().size(); }
328 unsigned AffineMap::getNumInputs() const {
329  assert(map && "uninitialized map storage");
330  return map->numDims + map->numSymbols;
331 }
333  assert(map && "uninitialized map storage");
334  return map->results();
335 }
336 AffineExpr AffineMap::getResult(unsigned idx) const {
337  return getResults()[idx];
338 }
339 
340 unsigned AffineMap::getDimPosition(unsigned idx) const {
341  return getResult(idx).cast<AffineDimExpr>().getPosition();
342 }
343 
344 std::optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const {
345  if (!input.isa<AffineDimExpr>())
346  return std::nullopt;
347 
348  for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {
349  if (getResult(i) == input)
350  return i;
351  }
352 
353  return std::nullopt;
354 }
355 
356 /// Folds the results of the application of an affine map on the provided
357 /// operands to a constant if possible. Returns false if the folding happens,
358 /// true otherwise.
361  SmallVectorImpl<Attribute> &results) const {
362  // Attempt partial folding.
363  SmallVector<int64_t, 2> integers;
364  partialConstantFold(operandConstants, &integers);
365 
366  // If all expressions folded to a constant, populate results with attributes
367  // containing those constants.
368  if (integers.empty())
369  return failure();
370 
371  auto range = llvm::map_range(integers, [this](int64_t i) {
372  return IntegerAttr::get(IndexType::get(getContext()), i);
373  });
374  results.append(range.begin(), range.end());
375  return success();
376 }
377 
378 AffineMap
380  SmallVectorImpl<int64_t> *results) const {
381  assert(getNumInputs() == operandConstants.size());
382 
383  // Fold each of the result expressions.
384  AffineExprConstantFolder exprFolder(getNumDims(), operandConstants);
386  exprs.reserve(getNumResults());
387 
388  for (auto expr : getResults()) {
389  auto folded = exprFolder.constantFold(expr);
390  // If did not fold to a constant, keep the original expression, and clear
391  // the integer results vector.
392  if (folded) {
393  exprs.push_back(
394  getAffineConstantExpr(folded.getInt(), folded.getContext()));
395  if (results)
396  results->push_back(folded.getInt());
397  } else {
398  exprs.push_back(expr);
399  if (results) {
400  results->clear();
401  results = nullptr;
402  }
403  }
404  }
405 
406  return get(getNumDims(), getNumSymbols(), exprs, getContext());
407 }
408 
409 /// Walk all of the AffineExpr's in this mapping. Each node in an expression
410 /// tree is visited in postorder.
412  for (auto expr : getResults())
413  expr.walk(callback);
414 }
415 
416 /// This method substitutes any uses of dimensions and symbols (e.g.
417 /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
418 /// expression mapping. Because this can be used to eliminate dims and
419 /// symbols, the client needs to specify the number of dims and symbols in
420 /// the result. The returned map always has the same number of results.
422  ArrayRef<AffineExpr> symReplacements,
423  unsigned numResultDims,
424  unsigned numResultSyms) const {
426  results.reserve(getNumResults());
427  for (auto expr : getResults())
428  results.push_back(
429  expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
430  return get(numResultDims, numResultSyms, results, getContext());
431 }
432 
433 /// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to
434 /// each of the results and return a new AffineMap with the new results and
435 /// with the specified number of dims and symbols.
437  unsigned numResultDims,
438  unsigned numResultSyms) const {
439  SmallVector<AffineExpr, 4> newResults;
440  newResults.reserve(getNumResults());
441  for (AffineExpr e : getResults())
442  newResults.push_back(e.replace(expr, replacement));
443  return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());
444 }
445 
446 /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
447 /// results and return a new AffineMap with the new results and with the
448 /// specified number of dims and symbols.
450  unsigned numResultDims,
451  unsigned numResultSyms) const {
452  SmallVector<AffineExpr, 4> newResults;
453  newResults.reserve(getNumResults());
454  for (AffineExpr e : getResults())
455  newResults.push_back(e.replace(map));
456  return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());
457 }
458 
459 AffineMap
461  SmallVector<AffineExpr, 4> newResults;
462  newResults.reserve(getNumResults());
463  for (AffineExpr e : getResults())
464  newResults.push_back(e.replace(map));
465  return AffineMap::inferFromExprList(newResults).front();
466 }
467 
469  assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
470  // Prepare `map` by concatenating the symbols and rewriting its exprs.
471  unsigned numDims = map.getNumDims();
472  unsigned numSymbolsThisMap = getNumSymbols();
473  unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
474  SmallVector<AffineExpr, 8> newDims(numDims);
475  for (unsigned idx = 0; idx < numDims; ++idx) {
476  newDims[idx] = getAffineDimExpr(idx, getContext());
477  }
478  SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap);
479  for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
480  newSymbols[idx - numSymbolsThisMap] =
482  }
483  auto newMap =
484  map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
486  exprs.reserve(getResults().size());
487  for (auto expr : getResults())
488  exprs.push_back(expr.compose(newMap));
489  return AffineMap::get(numDims, numSymbols, exprs, map.getContext());
490 }
491 
493  assert(getNumSymbols() == 0 && "Expected symbol-less map");
495  exprs.reserve(values.size());
496  MLIRContext *ctx = getContext();
497  for (auto v : values)
498  exprs.push_back(getAffineConstantExpr(v, ctx));
499  auto resMap = compose(AffineMap::get(0, 0, exprs, ctx));
501  res.reserve(resMap.getNumResults());
502  for (auto e : resMap.getResults())
503  res.push_back(e.cast<AffineConstantExpr>().getValue());
504  return res;
505 }
506 
507 bool AffineMap::isProjectedPermutation(bool allowZeroInResults) const {
508  if (getNumSymbols() > 0)
509  return false;
510 
511  // Having more results than inputs means that results have duplicated dims or
512  // zeros that can't be mapped to input dims.
513  if (getNumResults() > getNumInputs())
514  return false;
515 
516  SmallVector<bool, 8> seen(getNumInputs(), false);
517  // A projected permutation can have, at most, only one instance of each input
518  // dimension in the result expressions. Zeros are allowed as long as the
519  // number of result expressions is lower or equal than the number of input
520  // expressions.
521  for (auto expr : getResults()) {
522  if (auto dim = expr.dyn_cast<AffineDimExpr>()) {
523  if (seen[dim.getPosition()])
524  return false;
525  seen[dim.getPosition()] = true;
526  } else {
527  auto constExpr = expr.dyn_cast<AffineConstantExpr>();
528  if (!allowZeroInResults || !constExpr || constExpr.getValue() != 0)
529  return false;
530  }
531  }
532 
533  // Results are either dims or zeros and zeros can be mapped to input dims.
534  return true;
535 }
536 
538  if (getNumDims() != getNumResults())
539  return false;
540  return isProjectedPermutation();
541 }
542 
545  exprs.reserve(resultPos.size());
546  for (auto idx : resultPos)
547  exprs.push_back(getResult(idx));
548  return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
549 }
550 
551 AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const {
553  getResults().slice(start, length), getContext());
554 }
555 
556 AffineMap AffineMap::getMajorSubMap(unsigned numResults) const {
557  if (numResults == 0)
558  return AffineMap();
559  if (numResults > getNumResults())
560  return *this;
561  return getSliceMap(0, numResults);
562 }
563 
564 AffineMap AffineMap::getMinorSubMap(unsigned numResults) const {
565  if (numResults == 0)
566  return AffineMap();
567  if (numResults > getNumResults())
568  return *this;
569  return getSliceMap(getNumResults() - numResults, numResults);
570 }
571 
573  const llvm::SmallBitVector &unusedDims) {
574  unsigned numDims = 0;
575  SmallVector<AffineExpr> dimReplacements;
576  dimReplacements.reserve(map.getNumDims());
577  MLIRContext *context = map.getContext();
578  for (unsigned dim = 0, e = map.getNumDims(); dim < e; ++dim) {
579  if (unusedDims.test(dim))
580  dimReplacements.push_back(getAffineConstantExpr(0, context));
581  else
582  dimReplacements.push_back(getAffineDimExpr(numDims++, context));
583  }
584  SmallVector<AffineExpr> resultExprs;
585  resultExprs.reserve(map.getNumResults());
586  for (auto e : map.getResults())
587  resultExprs.push_back(e.replaceDims(dimReplacements));
588  return AffineMap::get(numDims, map.getNumSymbols(), resultExprs, context);
589 }
590 
592  return compressDims(map, getUnusedDimsBitVector({map}));
593 }
594 
597  llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
598  if (maps.empty())
599  return SmallVector<AffineMap>();
600  SmallVector<AffineExpr> allExprs;
601  allExprs.reserve(maps.size() * maps.front().getNumResults());
602  unsigned numDims = maps.front().getNumDims(),
603  numSymbols = maps.front().getNumSymbols();
604  for (auto m : maps) {
605  assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&
606  "expected maps with same num dims and symbols");
607  llvm::append_range(allExprs, m.getResults());
608  }
609  AffineMap unifiedMap = compressionFun(
610  AffineMap::get(numDims, numSymbols, allExprs, maps.front().getContext()));
611  unsigned unifiedNumDims = unifiedMap.getNumDims(),
612  unifiedNumSymbols = unifiedMap.getNumSymbols();
613  ArrayRef<AffineExpr> unifiedResults = unifiedMap.getResults();
615  res.reserve(maps.size());
616  for (auto m : maps) {
617  res.push_back(AffineMap::get(unifiedNumDims, unifiedNumSymbols,
618  unifiedResults.take_front(m.getNumResults()),
619  m.getContext()));
620  unifiedResults = unifiedResults.drop_front(m.getNumResults());
621  }
622  return res;
623 }
624 
626  return compressUnusedImpl(maps,
627  [](AffineMap m) { return compressUnusedDims(m); });
628 }
629 
631  const llvm::SmallBitVector &unusedSymbols) {
632  unsigned numSymbols = 0;
633  SmallVector<AffineExpr> symReplacements;
634  symReplacements.reserve(map.getNumSymbols());
635  MLIRContext *context = map.getContext();
636  for (unsigned sym = 0, e = map.getNumSymbols(); sym < e; ++sym) {
637  if (unusedSymbols.test(sym))
638  symReplacements.push_back(getAffineConstantExpr(0, context));
639  else
640  symReplacements.push_back(getAffineSymbolExpr(numSymbols++, context));
641  }
642  SmallVector<AffineExpr> resultExprs;
643  resultExprs.reserve(map.getNumResults());
644  for (auto e : map.getResults())
645  resultExprs.push_back(e.replaceSymbols(symReplacements));
646  return AffineMap::get(map.getNumDims(), numSymbols, resultExprs, context);
647 }
648 
650  llvm::SmallBitVector unusedSymbols(map.getNumSymbols(), true);
651  map.walkExprs([&](AffineExpr expr) {
652  if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>())
653  unusedSymbols.reset(symExpr.getPosition());
654  });
655  return compressSymbols(map, unusedSymbols);
656 }
657 
659  return compressUnusedImpl(
660  maps, [](AffineMap m) { return compressUnusedSymbols(m); });
661 }
662 
665  for (auto e : map.getResults()) {
666  exprs.push_back(
667  simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
668  }
669  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,
670  map.getContext());
671 }
672 
674  auto results = map.getResults();
675  SmallVector<AffineExpr, 4> uniqueExprs(results.begin(), results.end());
676  uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()),
677  uniqueExprs.end());
678  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs,
679  map.getContext());
680 }
681 
683  if (map.isEmpty())
684  return map;
685  assert(map.getNumSymbols() == 0 && "expected map without symbols");
687  for (const auto &en : llvm::enumerate(map.getResults())) {
688  auto expr = en.value();
689  // Skip non-permutations.
690  if (auto d = expr.dyn_cast<AffineDimExpr>()) {
691  if (exprs[d.getPosition()])
692  continue;
693  exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());
694  }
695  }
696  SmallVector<AffineExpr, 4> seenExprs;
697  seenExprs.reserve(map.getNumDims());
698  for (auto expr : exprs)
699  if (expr)
700  seenExprs.push_back(expr);
701  if (seenExprs.size() != map.getNumInputs())
702  return AffineMap();
703  return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext());
704 }
705 
707  assert(map.isProjectedPermutation(/*allowZeroInResults=*/true));
708  MLIRContext *context = map.getContext();
709  AffineExpr zero = mlir::getAffineConstantExpr(0, context);
710  // Start with all the results as 0.
711  SmallVector<AffineExpr, 4> exprs(map.getNumInputs(), zero);
712  for (unsigned i : llvm::seq(unsigned(0), map.getNumResults())) {
713  // Skip zeros from input map. 'exprs' is already initialized to zero.
714  if (auto constExpr = map.getResult(i).dyn_cast<AffineConstantExpr>()) {
715  assert(constExpr.getValue() == 0 &&
716  "Unexpected constant in projected permutation");
717  (void)constExpr;
718  continue;
719  }
720 
721  // Reverse each dimension existing in the original map result.
722  exprs[map.getDimPosition(i)] = getAffineDimExpr(i, context);
723  }
724  return AffineMap::get(map.getNumResults(), /*symbolCount=*/0, exprs, context);
725 }
726 
728  unsigned numResults = 0, numDims = 0, numSymbols = 0;
729  for (auto m : maps)
730  numResults += m.getNumResults();
732  results.reserve(numResults);
733  for (auto m : maps) {
734  for (auto res : m.getResults())
735  results.push_back(res.shiftSymbols(m.getNumSymbols(), numSymbols));
736 
737  numSymbols += m.getNumSymbols();
738  numDims = std::max(m.getNumDims(), numDims);
739  }
740  return AffineMap::get(numDims, numSymbols, results,
741  maps.front().getContext());
742 }
743 
745  const llvm::SmallBitVector &unusedDims) {
746  return compressUnusedSymbols(compressDims(map, unusedDims));
747 }
748 
750  unsigned numDims = maps[0].getNumDims();
751  llvm::SmallBitVector numDimsBitVector(numDims, true);
752  for (const auto &m : maps) {
753  for (unsigned i = 0; i < numDims; ++i) {
754  if (m.isFunctionOfDim(i))
755  numDimsBitVector.reset(i);
756  }
757  }
758  return numDimsBitVector;
759 }
760 
761 //===----------------------------------------------------------------------===//
762 // MutableAffineMap.
763 //===----------------------------------------------------------------------===//
764 
766  : results(map.getResults().begin(), map.getResults().end()),
767  numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
768  context(map.getContext()) {}
769 
771  results.clear();
772  numDims = map.getNumDims();
773  numSymbols = map.getNumSymbols();
774  context = map.getContext();
775  llvm::append_range(results, map.getResults());
776 }
777 
778 bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
779  if (results[idx].isMultipleOf(factor))
780  return true;
781 
782  // TODO: use simplifyAffineExpr and FlatAffineConstraints to
783  // complete this (for a more powerful analysis).
784  return false;
785 }
786 
787 // Simplifies the result affine expressions of this map. The expressions have to
788 // be pure for the simplification implemented.
790  // Simplify each of the results if possible.
791  // TODO: functional-style map
792  for (unsigned i = 0, e = getNumResults(); i < e; i++) {
793  results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
794  }
795 }
796 
798  return AffineMap::get(numDims, numSymbols, results, context);
799 }
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< AffineExprContainer > exprsList)
Definition: AffineMap.cpp:222
static SmallVector< AffineMap > compressUnusedImpl(ArrayRef< AffineMap > maps, llvm::function_ref< AffineMap(AffineMap)> compressionFun)
Definition: AffineMap.cpp:596
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Affine binary operation expression.
Definition: AffineExpr.h:207
An integer constant appearing in affine expression.
Definition: AffineExpr.h:232
int64_t getValue() const
Definition: AffineExpr.cpp:519
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:216
unsigned getPosition() const
Definition: AffineExpr.cpp:325
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineExpr.cpp:66
U cast() const
Definition: AffineExpr.h:291
void walk(std::function< void(AffineExpr)> callback) const
Walk all of the AffineExpr's in this expression in postorder.
Definition: AffineExpr.cpp:30
AffineExprKind getKind() const
Return the classification for this type.
Definition: AffineExpr.cpp:27
AffineExpr compose(AffineMap map) const
Compose with an AffineMap.
Definition: AffineExpr.cpp:889
constexpr bool isa() const
Definition: AffineExpr.h:270
MLIRContext * getContext() const
Definition: AffineExpr.cpp:25
U dyn_cast() const
Definition: AffineExpr.h:281
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:43
int64_t getSingleConstantResult() const
Returns the constant result of this map.
Definition: AffineMap.cpp:306
static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, MLIRContext *context)
Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
Definition: AffineMap.cpp:104
AffineMap getSliceMap(unsigned start, unsigned length) const
Returns the map consisting of length expressions starting from start.
Definition: AffineMap.cpp:551
AffineMap getMajorSubMap(unsigned numResults) const
Returns the map consisting of the most major numResults results.
Definition: AffineMap.cpp:556
MLIRContext * getContext() const
Definition: AffineMap.cpp:266
AffineMap partialConstantFold(ArrayRef< Attribute > operandConstants, SmallVectorImpl< int64_t > *results=nullptr) const
Propagates the constant operands into this affine map.
Definition: AffineMap.cpp:379
bool isMinorIdentity() const
Returns true if this affine map is a minor identity, i.e.
Definition: AffineMap.cpp:111
unsigned getDimPosition(unsigned idx) const
Extracts the position of the dimensional expression at the given result, when the caller knows it is ...
Definition: AffineMap.cpp:340
bool isConstant() const
Returns true if this affine map has only constant results.
Definition: AffineMap.cpp:300
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
Definition: AffineMap.cpp:257
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
bool isSingleConstant() const
Returns true if this affine map is a single result constant function.
Definition: AffineMap.cpp:296
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
Definition: AffineMap.cpp:507
AffineMap getMinorSubMap(unsigned numResults) const
Returns the map consisting of the most minor numResults results.
Definition: AffineMap.cpp:564
uint64_t getLargestKnownDivisorOfMapExprs()
Get the largest known divisor of all map expressions.
Definition: AffineMap.cpp:246
constexpr AffineMap()=default
bool isEmpty() const
Returns true if this affine map is an empty map, i.e., () -> ().
Definition: AffineMap.cpp:292
std::optional< unsigned > getResultPosition(AffineExpr input) const
Extracts the first result position where input dimension resides.
Definition: AffineMap.cpp:344
unsigned getNumSymbols() const
Definition: AffineMap.cpp:323
bool isMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > *broadcastedDims=nullptr) const
Returns true if this affine map is a minor identity up to broadcasted dimensions which are indicated ...
Definition: AffineMap.cpp:119
unsigned getNumDims() const
Definition: AffineMap.cpp:319
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:332
SmallVector< int64_t > getConstantResults() const
Returns the constant results of this map.
Definition: AffineMap.cpp:311
bool isPermutationOfMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > &permutedDims) const
Return true if this affine map can be converted to a minor identity with broadcast by doing a permute...
Definition: AffineMap.cpp:159
LogicalResult constantFold(ArrayRef< Attribute > operandConstants, SmallVectorImpl< Attribute > &results) const
Folds the results of the application of an affine map on the provided operands to a constant if possi...
Definition: AffineMap.cpp:360
bool isSymbolIdentity() const
Returns true if this affine map is an identity affine map on the symbol identifiers.
Definition: AffineMap.cpp:280
unsigned getNumResults() const
Definition: AffineMap.cpp:327
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineMap.cpp:421
unsigned getNumInputs() const
Definition: AffineMap.cpp:328
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< ArrayRef< AffineExpr >> exprsList)
Returns a vector of AffineMaps; each with as many results as exprs.size(), as many dims as the larges...
Definition: AffineMap.cpp:237
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:336
AffineMap replace(AffineExpr expr, AffineExpr replacement, unsigned numResultDims, unsigned numResultSyms) const
Sparse replace method.
Definition: AffineMap.cpp:436
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:207
void walkExprs(llvm::function_ref< void(AffineExpr)> callback) const
Walk all of the AffineExpr's in this mapping.
Definition: AffineMap.cpp:411
static AffineMap getConstantMap(int64_t val, MLIRContext *context)
Returns a single constant result affine map.
Definition: AffineMap.cpp:97
AffineMap getSubMap(ArrayRef< unsigned > resultPos) const
Returns the map consisting of the resultPos subset.
Definition: AffineMap.cpp:543
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
Definition: AffineMap.cpp:468
bool isIdentity() const
Returns true if this affine map is an identity affine map.
Definition: AffineMap.cpp:268
bool isPermutation() const
Returns true if the AffineMap represents a symbol-less permutation map.
Definition: AffineMap.cpp:537
A symbolic identifier appearing in an affine expression.
Definition: AffineExpr.h:224
unsigned getPosition() const
Definition: AffineExpr.cpp:508
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:56
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:223
LLVM_ATTRIBUTE_ALWAYS_INLINE MPInt gcd(const MPInt &a, const MPInt &b)
Definition: MPInt.h:399
Include the generated interface declarations.
AffineMap simplifyAffineMap(AffineMap map)
Simplifies an affine map by simplifying its underlying AffineExpr results.
Definition: AffineMap.cpp:663
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
AffineMap removeDuplicateExprs(AffineMap map)
Returns a map with the same dimension and symbol count as map, but whose results are the unique affin...
Definition: AffineMap.cpp:673
AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map)
Return the reverse map of a projected permutation where the projected dimensions are transformed into...
Definition: AffineMap.cpp:706
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's floordiv operation on constants.
Definition: MathExtras.h:33
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's ceildiv operation on constants.
Definition: MathExtras.h:23
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:682
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
AffineMap concatAffineMaps(ArrayRef< AffineMap > maps)
Concatenates a list of maps into a single AffineMap, stepping over potentially empty maps.
Definition: AffineMap.cpp:727
@ CeilDiv
RHS of ceildiv is always a constant or a symbolic expression.
@ Mul
RHS of mul is always a constant or a symbolic expression.
@ Mod
RHS of mod is always a constant or a symbolic expression with a positive value.
@ DimId
Dimensional identifier.
@ FloorDiv
RHS of floordiv is always a constant or a symbolic expression.
@ Constant
Constant integer.
@ SymbolId
Symbolic identifier.
AffineMap compressSymbols(AffineMap map, const llvm::SmallBitVector &unusedSymbols)
Drop the symbols that are not listed in unusedSymbols.
Definition: AffineMap.cpp:630
static void getMaxDimAndSymbol(ArrayRef< AffineExprContainer > exprsList, int64_t &maxDim, int64_t &maxSym)
Calculates maxmimum dimension and symbol positions from the expressions in exprsLists and stores them...
Definition: AffineMap.h:590
AffineMap compressUnusedDims(AffineMap map)
Drop the dims that are not used.
Definition: AffineMap.cpp:591
AffineMap getProjectedMap(AffineMap map, const llvm::SmallBitVector &projectedDimensions)
Returns the map that results from projecting out the dimensions specified in projectedDimensions.
Definition: AffineMap.cpp:744
AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims)
Drop the dims that are not listed in unusedDims.
Definition: AffineMap.cpp:572
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:527
llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef< AffineMap > maps)
Definition: AffineMap.cpp:749
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify an affine expression by flattening and some amount of simple analysis.
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:502
AffineMap compressUnusedSymbols(AffineMap map)
Drop the symbols that are not used.
Definition: AffineMap.cpp:649
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)
Definition: AffineExpr.cpp:512
int64_t mod(int64_t lhs, int64_t rhs)
Returns MLIR's mod operation on constants.
Definition: MathExtras.h:45
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
void reset(AffineMap map)
Resets this MutableAffineMap with 'map'.
Definition: AffineMap.cpp:770
AffineMap getAffineMap() const
Get the AffineMap corresponding to this MutableAffineMap.
Definition: AffineMap.cpp:797
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.h:370
bool isMultipleOf(unsigned idx, int64_t factor) const
Returns true if the idx'th result expression is a multiple of factor.
Definition: AffineMap.cpp:778
unsigned getNumResults() const
Definition: AffineMap.h:372
void simplify()
Simplify the (result) expressions in this map using analysis (used by.
Definition: AffineMap.cpp:789
ArrayRef< AffineExpr > results() const
The affine expressions for this (multi-dimensional) map.