MLIR  19.0.0git
AffineExpr.cpp
Go to the documentation of this file.
1 //===- AffineExpr.cpp - MLIR Affine Expr 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 <utility>
10 
11 #include "AffineExprDetail.h"
12 #include "mlir/IR/AffineExpr.h"
14 #include "mlir/IR/AffineMap.h"
15 #include "mlir/IR/IntegerSet.h"
16 #include "mlir/Support/TypeID.h"
17 #include "llvm/ADT/STLExtras.h"
18 #include "llvm/Support/MathExtras.h"
19 #include <numeric>
20 #include <optional>
21 
22 using namespace mlir;
23 using namespace mlir::detail;
24 
25 using llvm::divideCeilSigned;
26 using llvm::divideFloorSigned;
27 using llvm::mod;
28 
29 MLIRContext *AffineExpr::getContext() const { return expr->context; }
30 
31 AffineExprKind AffineExpr::getKind() const { return expr->kind; }
32 
33 /// Walk all of the AffineExprs in `e` in postorder. This is a private factory
34 /// method to help handle lambda walk functions. Users should use the regular
35 /// (non-static) `walk` method.
36 template <typename WalkRetTy>
38  function_ref<WalkRetTy(AffineExpr)> callback) {
39  struct AffineExprWalker
40  : public AffineExprVisitor<AffineExprWalker, WalkRetTy> {
41  function_ref<WalkRetTy(AffineExpr)> callback;
42 
43  AffineExprWalker(function_ref<WalkRetTy(AffineExpr)> callback)
44  : callback(callback) {}
45 
46  WalkRetTy visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
47  return callback(expr);
48  }
49  WalkRetTy visitConstantExpr(AffineConstantExpr expr) {
50  return callback(expr);
51  }
52  WalkRetTy visitDimExpr(AffineDimExpr expr) { return callback(expr); }
53  WalkRetTy visitSymbolExpr(AffineSymbolExpr expr) { return callback(expr); }
54  };
55 
56  return AffineExprWalker(callback).walkPostOrder(e);
57 }
58 // Explicitly instantiate for the two supported return types.
59 template void mlir::AffineExpr::walk(AffineExpr e,
60  function_ref<void(AffineExpr)> callback);
61 template WalkResult
64 
65 // Dispatch affine expression construction based on kind.
67  AffineExpr rhs) {
68  if (kind == AffineExprKind::Add)
69  return lhs + rhs;
70  if (kind == AffineExprKind::Mul)
71  return lhs * rhs;
72  if (kind == AffineExprKind::FloorDiv)
73  return lhs.floorDiv(rhs);
74  if (kind == AffineExprKind::CeilDiv)
75  return lhs.ceilDiv(rhs);
76  if (kind == AffineExprKind::Mod)
77  return lhs % rhs;
78 
79  llvm_unreachable("unknown binary operation on affine expressions");
80 }
81 
82 /// This method substitutes any uses of dimensions and symbols (e.g.
83 /// dim#0 with dimReplacements[0]) and returns the modified expression tree.
86  ArrayRef<AffineExpr> symReplacements) const {
87  switch (getKind()) {
89  return *this;
90  case AffineExprKind::DimId: {
91  unsigned dimId = llvm::cast<AffineDimExpr>(*this).getPosition();
92  if (dimId >= dimReplacements.size())
93  return *this;
94  return dimReplacements[dimId];
95  }
97  unsigned symId = llvm::cast<AffineSymbolExpr>(*this).getPosition();
98  if (symId >= symReplacements.size())
99  return *this;
100  return symReplacements[symId];
101  }
102  case AffineExprKind::Add:
103  case AffineExprKind::Mul:
106  case AffineExprKind::Mod:
107  auto binOp = llvm::cast<AffineBinaryOpExpr>(*this);
108  auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
109  auto newLHS = lhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
110  auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
111  if (newLHS == lhs && newRHS == rhs)
112  return *this;
113  return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
114  }
115  llvm_unreachable("Unknown AffineExpr");
116 }
117 
119  return replaceDimsAndSymbols(dimReplacements, {});
120 }
121 
124  return replaceDimsAndSymbols({}, symReplacements);
125 }
126 
127 /// Replace dims[offset ... numDims)
128 /// by dims[offset + shift ... shift + numDims).
129 AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift,
130  unsigned offset) const {
132  for (unsigned idx = 0; idx < offset; ++idx)
133  dims.push_back(getAffineDimExpr(idx, getContext()));
134  for (unsigned idx = offset; idx < numDims; ++idx)
135  dims.push_back(getAffineDimExpr(idx + shift, getContext()));
136  return replaceDimsAndSymbols(dims, {});
137 }
138 
139 /// Replace symbols[offset ... numSymbols)
140 /// by symbols[offset + shift ... shift + numSymbols).
141 AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift,
142  unsigned offset) const {
144  for (unsigned idx = 0; idx < offset; ++idx)
145  symbols.push_back(getAffineSymbolExpr(idx, getContext()));
146  for (unsigned idx = offset; idx < numSymbols; ++idx)
147  symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
148  return replaceDimsAndSymbols({}, symbols);
149 }
150 
151 /// Sparse replace method. Return the modified expression tree.
154  auto it = map.find(*this);
155  if (it != map.end())
156  return it->second;
157  switch (getKind()) {
158  default:
159  return *this;
160  case AffineExprKind::Add:
161  case AffineExprKind::Mul:
164  case AffineExprKind::Mod:
165  auto binOp = llvm::cast<AffineBinaryOpExpr>(*this);
166  auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
167  auto newLHS = lhs.replace(map);
168  auto newRHS = rhs.replace(map);
169  if (newLHS == lhs && newRHS == rhs)
170  return *this;
171  return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
172  }
173  llvm_unreachable("Unknown AffineExpr");
174 }
175 
176 /// Sparse replace method. Return the modified expression tree.
179  map.insert(std::make_pair(expr, replacement));
180  return replace(map);
181 }
182 /// Returns true if this expression is made out of only symbols and
183 /// constants (no dimensional identifiers).
185  switch (getKind()) {
187  return true;
189  return false;
191  return true;
192 
193  case AffineExprKind::Add:
194  case AffineExprKind::Mul:
197  case AffineExprKind::Mod: {
198  auto expr = llvm::cast<AffineBinaryOpExpr>(*this);
199  return expr.getLHS().isSymbolicOrConstant() &&
200  expr.getRHS().isSymbolicOrConstant();
201  }
202  }
203  llvm_unreachable("Unknown AffineExpr");
204 }
205 
206 /// Returns true if this is a pure affine expression, i.e., multiplication,
207 /// floordiv, ceildiv, and mod is only allowed w.r.t constants.
209  switch (getKind()) {
213  return true;
214  case AffineExprKind::Add: {
215  auto op = llvm::cast<AffineBinaryOpExpr>(*this);
216  return op.getLHS().isPureAffine() && op.getRHS().isPureAffine();
217  }
218 
219  case AffineExprKind::Mul: {
220  // TODO: Canonicalize the constants in binary operators to the RHS when
221  // possible, allowing this to merge into the next case.
222  auto op = llvm::cast<AffineBinaryOpExpr>(*this);
223  return op.getLHS().isPureAffine() && op.getRHS().isPureAffine() &&
224  (llvm::isa<AffineConstantExpr>(op.getLHS()) ||
225  llvm::isa<AffineConstantExpr>(op.getRHS()));
226  }
229  case AffineExprKind::Mod: {
230  auto op = llvm::cast<AffineBinaryOpExpr>(*this);
231  return op.getLHS().isPureAffine() &&
232  llvm::isa<AffineConstantExpr>(op.getRHS());
233  }
234  }
235  llvm_unreachable("Unknown AffineExpr");
236 }
237 
238 // Returns the greatest known integral divisor of this affine expression.
240  AffineBinaryOpExpr binExpr(nullptr);
241  switch (getKind()) {
243  [[fallthrough]];
245  return 1;
247  [[fallthrough]];
249  // If the RHS is a constant and divides the known divisor on the LHS, the
250  // quotient is a known divisor of the expression.
251  binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
252  auto rhs = llvm::dyn_cast<AffineConstantExpr>(binExpr.getRHS());
253  // Leave alone undefined expressions.
254  if (rhs && rhs.getValue() != 0) {
255  int64_t lhsDiv = binExpr.getLHS().getLargestKnownDivisor();
256  if (lhsDiv % rhs.getValue() == 0)
257  return lhsDiv / rhs.getValue();
258  }
259  return 1;
260  }
262  return std::abs(llvm::cast<AffineConstantExpr>(*this).getValue());
263  case AffineExprKind::Mul: {
264  binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
265  return binExpr.getLHS().getLargestKnownDivisor() *
266  binExpr.getRHS().getLargestKnownDivisor();
267  }
268  case AffineExprKind::Add:
269  [[fallthrough]];
270  case AffineExprKind::Mod: {
271  binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
272  return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(),
273  (uint64_t)binExpr.getRHS().getLargestKnownDivisor());
274  }
275  }
276  llvm_unreachable("Unknown AffineExpr");
277 }
278 
279 bool AffineExpr::isMultipleOf(int64_t factor) const {
280  AffineBinaryOpExpr binExpr(nullptr);
281  uint64_t l, u;
282  switch (getKind()) {
284  [[fallthrough]];
286  return factor * factor == 1;
288  return llvm::cast<AffineConstantExpr>(*this).getValue() % factor == 0;
289  case AffineExprKind::Mul: {
290  binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
291  // It's probably not worth optimizing this further (to not traverse the
292  // whole sub-tree under - it that would require a version of isMultipleOf
293  // that on a 'false' return also returns the largest known divisor).
294  return (l = binExpr.getLHS().getLargestKnownDivisor()) % factor == 0 ||
295  (u = binExpr.getRHS().getLargestKnownDivisor()) % factor == 0 ||
296  (l * u) % factor == 0;
297  }
298  case AffineExprKind::Add:
301  case AffineExprKind::Mod: {
302  binExpr = llvm::cast<AffineBinaryOpExpr>(*this);
303  return std::gcd((uint64_t)binExpr.getLHS().getLargestKnownDivisor(),
304  (uint64_t)binExpr.getRHS().getLargestKnownDivisor()) %
305  factor ==
306  0;
307  }
308  }
309  llvm_unreachable("Unknown AffineExpr");
310 }
311 
312 bool AffineExpr::isFunctionOfDim(unsigned position) const {
313  if (getKind() == AffineExprKind::DimId) {
314  return *this == mlir::getAffineDimExpr(position, getContext());
315  }
316  if (auto expr = llvm::dyn_cast<AffineBinaryOpExpr>(*this)) {
317  return expr.getLHS().isFunctionOfDim(position) ||
318  expr.getRHS().isFunctionOfDim(position);
319  }
320  return false;
321 }
322 
323 bool AffineExpr::isFunctionOfSymbol(unsigned position) const {
324  if (getKind() == AffineExprKind::SymbolId) {
325  return *this == mlir::getAffineSymbolExpr(position, getContext());
326  }
327  if (auto expr = llvm::dyn_cast<AffineBinaryOpExpr>(*this)) {
328  return expr.getLHS().isFunctionOfSymbol(position) ||
329  expr.getRHS().isFunctionOfSymbol(position);
330  }
331  return false;
332 }
333 
335  : AffineExpr(ptr) {}
337  return static_cast<ImplType *>(expr)->lhs;
338 }
340  return static_cast<ImplType *>(expr)->rhs;
341 }
342 
344 unsigned AffineDimExpr::getPosition() const {
345  return static_cast<ImplType *>(expr)->position;
346 }
347 
348 /// Returns true if the expression is divisible by the given symbol with
349 /// position `symbolPos`. The argument `opKind` specifies here what kind of
350 /// division or mod operation called this division. It helps in implementing the
351 /// commutative property of the floordiv and ceildiv operations. If the argument
352 ///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv
353 /// operation, then the commutative property can be used otherwise, the floordiv
354 /// operation is not divisible. The same argument holds for ceildiv operation.
355 static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos,
356  AffineExprKind opKind) {
357  // The argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
358  assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
359  opKind == AffineExprKind::CeilDiv) &&
360  "unexpected opKind");
361  switch (expr.getKind()) {
363  return cast<AffineConstantExpr>(expr).getValue() == 0;
365  return false;
367  return (cast<AffineSymbolExpr>(expr).getPosition() == symbolPos);
368  // Checks divisibility by the given symbol for both operands.
369  case AffineExprKind::Add: {
370  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
371  return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) &&
372  isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
373  }
374  // Checks divisibility by the given symbol for both operands. Consider the
375  // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`,
376  // this is a division by s1 and both the operands of modulo are divisible by
377  // s1 but it is not divisible by s1 always. The third argument is
378  // `AffineExprKind::Mod` for this reason.
379  case AffineExprKind::Mod: {
380  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
381  return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos,
383  isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos,
385  }
386  // Checks if any of the operand divisible by the given symbol.
387  case AffineExprKind::Mul: {
388  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
389  return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) ||
390  isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
391  }
392  // Floordiv and ceildiv are divisible by the given symbol when the first
393  // operand is divisible, and the affine expression kind of the argument expr
394  // is same as the argument `opKind`. This can be inferred from commutative
395  // property of floordiv and ceildiv operations and are as follow:
396  // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
397  // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
398  // It will fail if operations are not same. For example:
399  // (exps1 ceildiv exp2) floordiv exp3 can not be simplified.
402  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
403  if (opKind != expr.getKind())
404  return false;
405  return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind());
406  }
407  }
408  llvm_unreachable("Unknown AffineExpr");
409 }
410 
411 /// Divides the given expression by the given symbol at position `symbolPos`. It
412 /// considers the divisibility condition is checked before calling itself. A
413 /// null expression is returned whenever the divisibility condition fails.
414 static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos,
415  AffineExprKind opKind) {
416  // THe argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
417  assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
418  opKind == AffineExprKind::CeilDiv) &&
419  "unexpected opKind");
420  switch (expr.getKind()) {
422  if (cast<AffineConstantExpr>(expr).getValue() != 0)
423  return nullptr;
424  return getAffineConstantExpr(0, expr.getContext());
426  return nullptr;
428  return getAffineConstantExpr(1, expr.getContext());
429  // Dividing both operands by the given symbol.
430  case AffineExprKind::Add: {
431  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
432  return getAffineBinaryOpExpr(
433  expr.getKind(), symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind),
434  symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind));
435  }
436  // Dividing both operands by the given symbol.
437  case AffineExprKind::Mod: {
438  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
439  return getAffineBinaryOpExpr(
440  expr.getKind(),
441  symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
442  symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind()));
443  }
444  // Dividing any of the operand by the given symbol.
445  case AffineExprKind::Mul: {
446  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
447  if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind))
448  return binaryExpr.getLHS() *
449  symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind);
450  return symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind) *
451  binaryExpr.getRHS();
452  }
453  // Dividing first operand only by the given symbol.
456  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
457  return getAffineBinaryOpExpr(
458  expr.getKind(),
459  symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
460  binaryExpr.getRHS());
461  }
462  }
463  llvm_unreachable("Unknown AffineExpr");
464 }
465 
466 /// Populate `result` with all summand operands of given (potentially nested)
467 /// addition. If the given expression is not an addition, just populate the
468 /// expression itself.
469 /// Example: Add(Add(7, 8), Mul(9, 10)) will return [7, 8, Mul(9, 10)].
471  auto addExpr = dyn_cast<AffineBinaryOpExpr>(expr);
472  if (!addExpr || addExpr.getKind() != AffineExprKind::Add) {
473  result.push_back(expr);
474  return;
475  }
476  getSummandExprs(addExpr.getLHS(), result);
477  getSummandExprs(addExpr.getRHS(), result);
478 }
479 
480 /// Return "true" if `candidate` is a negated expression, i.e., Mul(-1, expr).
481 /// If so, also return the non-negated expression via `expr`.
482 static bool isNegatedAffineExpr(AffineExpr candidate, AffineExpr &expr) {
483  auto mulExpr = dyn_cast<AffineBinaryOpExpr>(candidate);
484  if (!mulExpr || mulExpr.getKind() != AffineExprKind::Mul)
485  return false;
486  if (auto lhs = dyn_cast<AffineConstantExpr>(mulExpr.getLHS())) {
487  if (lhs.getValue() == -1) {
488  expr = mulExpr.getRHS();
489  return true;
490  }
491  }
492  if (auto rhs = dyn_cast<AffineConstantExpr>(mulExpr.getRHS())) {
493  if (rhs.getValue() == -1) {
494  expr = mulExpr.getLHS();
495  return true;
496  }
497  }
498  return false;
499 }
500 
501 /// Return "true" if `lhs` % `rhs` is guaranteed to evaluate to zero based on
502 /// the fact that `lhs` contains another modulo expression that ensures that
503 /// `lhs` is divisible by `rhs`. This is a common pattern in the resulting IR
504 /// after loop peeling.
505 ///
506 /// Example: lhs = ub - ub % step
507 /// rhs = step
508 /// => (ub - ub % step) % step is guaranteed to evaluate to 0.
510  unsigned numDims, unsigned numSymbols) {
511  // TODO: Try to unify this function with `getBoundForAffineExpr`.
512  // Collect all summands in lhs.
513  SmallVector<AffineExpr> summands;
514  getSummandExprs(lhs, summands);
515  // Look for Mul(-1, Mod(x, rhs)) among the summands. If x matches the
516  // remaining summands, then lhs % rhs is guaranteed to evaluate to 0.
517  for (int64_t i = 0, e = summands.size(); i < e; ++i) {
518  AffineExpr current = summands[i];
519  AffineExpr beforeNegation;
520  if (!isNegatedAffineExpr(current, beforeNegation))
521  continue;
522  AffineBinaryOpExpr innerMod = dyn_cast<AffineBinaryOpExpr>(beforeNegation);
523  if (!innerMod || innerMod.getKind() != AffineExprKind::Mod)
524  continue;
525  if (innerMod.getRHS() != rhs)
526  continue;
527  // Sum all remaining summands and subtract x. If that expression can be
528  // simplified to zero, then the remaining summands and x are equal.
530  for (int64_t j = 0; j < e; ++j)
531  if (i != j)
532  diff = diff + summands[j];
533  diff = diff - innerMod.getLHS();
534  diff = simplifyAffineExpr(diff, numDims, numSymbols);
535  auto constExpr = dyn_cast<AffineConstantExpr>(diff);
536  if (constExpr && constExpr.getValue() == 0)
537  return true;
538  }
539  return false;
540 }
541 
542 /// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv
543 /// operations when the second operand simplifies to a symbol and the first
544 /// operand is divisible by that symbol. It can be applied to any semi-affine
545 /// expression. Returned expression can either be a semi-affine or pure affine
546 /// expression.
547 static AffineExpr simplifySemiAffine(AffineExpr expr, unsigned numDims,
548  unsigned numSymbols) {
549  switch (expr.getKind()) {
553  return expr;
554  case AffineExprKind::Add:
555  case AffineExprKind::Mul: {
556  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
557  return getAffineBinaryOpExpr(
558  expr.getKind(),
559  simplifySemiAffine(binaryExpr.getLHS(), numDims, numSymbols),
560  simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols));
561  }
562  // Check if the simplification of the second operand is a symbol, and the
563  // first operand is divisible by it. If the operation is a modulo, a constant
564  // zero expression is returned. In the case of floordiv and ceildiv, the
565  // symbol from the simplification of the second operand divides the first
566  // operand. Otherwise, simplification is not possible.
569  case AffineExprKind::Mod: {
570  AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
571  AffineExpr sLHS =
572  simplifySemiAffine(binaryExpr.getLHS(), numDims, numSymbols);
573  AffineExpr sRHS =
574  simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols);
575  if (isModOfModSubtraction(sLHS, sRHS, numDims, numSymbols))
576  return getAffineConstantExpr(0, expr.getContext());
577  AffineSymbolExpr symbolExpr = dyn_cast<AffineSymbolExpr>(
578  simplifySemiAffine(binaryExpr.getRHS(), numDims, numSymbols));
579  if (!symbolExpr)
580  return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
581  unsigned symbolPos = symbolExpr.getPosition();
582  if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind()))
583  return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
584  if (expr.getKind() == AffineExprKind::Mod)
585  return getAffineConstantExpr(0, expr.getContext());
586  return symbolicDivide(sLHS, symbolPos, expr.getKind());
587  }
588  }
589  llvm_unreachable("Unknown AffineExpr");
590 }
591 
592 static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
593  MLIRContext *context) {
594  auto assignCtx = [context](AffineDimExprStorage *storage) {
595  storage->context = context;
596  };
597 
598  StorageUniquer &uniquer = context->getAffineUniquer();
599  return uniquer.get<AffineDimExprStorage>(
600  assignCtx, static_cast<unsigned>(kind), position);
601 }
602 
603 AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) {
604  return getAffineDimOrSymbol(AffineExprKind::DimId, position, context);
605 }
606 
608  : AffineExpr(ptr) {}
610  return static_cast<ImplType *>(expr)->position;
611 }
612 
613 AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) {
614  return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context);
615 }
616 
618  : AffineExpr(ptr) {}
620  return static_cast<ImplType *>(expr)->constant;
621 }
622 
623 bool AffineExpr::operator==(int64_t v) const {
624  return *this == getAffineConstantExpr(v, getContext());
625 }
626 
628  auto assignCtx = [context](AffineConstantExprStorage *storage) {
629  storage->context = context;
630  };
631 
632  StorageUniquer &uniquer = context->getAffineUniquer();
633  return uniquer.get<AffineConstantExprStorage>(assignCtx, constant);
634 }
635 
638  MLIRContext *context) {
639  return llvm::to_vector(llvm::map_range(constants, [&](int64_t constant) {
640  return getAffineConstantExpr(constant, context);
641  }));
642 }
643 
644 /// Simplify add expression. Return nullptr if it can't be simplified.
646  auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
647  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
648  // Fold if both LHS, RHS are a constant.
649  if (lhsConst && rhsConst)
650  return getAffineConstantExpr(lhsConst.getValue() + rhsConst.getValue(),
651  lhs.getContext());
652 
653  // Canonicalize so that only the RHS is a constant. (4 + d0 becomes d0 + 4).
654  // If only one of them is a symbolic expressions, make it the RHS.
655  if (isa<AffineConstantExpr>(lhs) ||
656  (lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())) {
657  return rhs + lhs;
658  }
659 
660  // At this point, if there was a constant, it would be on the right.
661 
662  // Addition with a zero is a noop, return the other input.
663  if (rhsConst) {
664  if (rhsConst.getValue() == 0)
665  return lhs;
666  }
667  // Fold successive additions like (d0 + 2) + 3 into d0 + 5.
668  auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
669  if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) {
670  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS()))
671  return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue());
672  }
673 
674  // Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr".
675  // c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their
676  // respective multiplicands.
677  std::optional<int64_t> rLhsConst, rRhsConst;
678  AffineExpr firstExpr, secondExpr;
679  AffineConstantExpr rLhsConstExpr;
680  auto lBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lhs);
681  if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul &&
682  (rLhsConstExpr = dyn_cast<AffineConstantExpr>(lBinOpExpr.getRHS()))) {
683  rLhsConst = rLhsConstExpr.getValue();
684  firstExpr = lBinOpExpr.getLHS();
685  } else {
686  rLhsConst = 1;
687  firstExpr = lhs;
688  }
689 
690  auto rBinOpExpr = dyn_cast<AffineBinaryOpExpr>(rhs);
691  AffineConstantExpr rRhsConstExpr;
692  if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul &&
693  (rRhsConstExpr = dyn_cast<AffineConstantExpr>(rBinOpExpr.getRHS()))) {
694  rRhsConst = rRhsConstExpr.getValue();
695  secondExpr = rBinOpExpr.getLHS();
696  } else {
697  rRhsConst = 1;
698  secondExpr = rhs;
699  }
700 
701  if (rLhsConst && rRhsConst && firstExpr == secondExpr)
702  return getAffineBinaryOpExpr(
703  AffineExprKind::Mul, firstExpr,
704  getAffineConstantExpr(*rLhsConst + *rRhsConst, lhs.getContext()));
705 
706  // When doing successive additions, bring constant to the right: turn (d0 + 2)
707  // + d1 into (d0 + d1) + 2.
708  if (lBin && lBin.getKind() == AffineExprKind::Add) {
709  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
710  return lBin.getLHS() + rhs + lrhs;
711  }
712  }
713 
714  // Detect and transform "expr - q * (expr floordiv q)" to "expr mod q", where
715  // q may be a constant or symbolic expression. This leads to a much more
716  // efficient form when 'c' is a power of two, and in general a more compact
717  // and readable form.
718 
719  // Process '(expr floordiv c) * (-c)'.
720  if (!rBinOpExpr)
721  return nullptr;
722 
723  auto lrhs = rBinOpExpr.getLHS();
724  auto rrhs = rBinOpExpr.getRHS();
725 
726  AffineExpr llrhs, rlrhs;
727 
728  // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a
729  // symbolic expression.
730  auto lrhsBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
731  // Check rrhsConstOpExpr = -1.
732  auto rrhsConstOpExpr = dyn_cast<AffineConstantExpr>(rrhs);
733  if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr &&
734  lrhsBinOpExpr.getKind() == AffineExprKind::Mul) {
735  // Check llrhs = expr floordiv q.
736  llrhs = lrhsBinOpExpr.getLHS();
737  // Check rlrhs = q.
738  rlrhs = lrhsBinOpExpr.getRHS();
739  auto llrhsBinOpExpr = dyn_cast<AffineBinaryOpExpr>(llrhs);
740  if (!llrhsBinOpExpr || llrhsBinOpExpr.getKind() != AffineExprKind::FloorDiv)
741  return nullptr;
742  if (llrhsBinOpExpr.getRHS() == rlrhs && lhs == llrhsBinOpExpr.getLHS())
743  return lhs % rlrhs;
744  }
745 
746  // Process lrhs, which is 'expr floordiv c'.
747  AffineBinaryOpExpr lrBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
748  if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
749  return nullptr;
750 
751  llrhs = lrBinOpExpr.getLHS();
752  rlrhs = lrBinOpExpr.getRHS();
753 
754  if (lhs == llrhs && rlrhs == -rrhs) {
755  return lhs % rlrhs;
756  }
757  return nullptr;
758 }
759 
761  return *this + getAffineConstantExpr(v, getContext());
762 }
764  if (auto simplified = simplifyAdd(*this, other))
765  return simplified;
766 
768  return uniquer.get<AffineBinaryOpExprStorage>(
769  /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
770 }
771 
772 /// Simplify a multiply expression. Return nullptr if it can't be simplified.
774  auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
775  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
776 
777  if (lhsConst && rhsConst)
778  return getAffineConstantExpr(lhsConst.getValue() * rhsConst.getValue(),
779  lhs.getContext());
780 
781  if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())
782  return nullptr;
783 
784  // Canonicalize the mul expression so that the constant/symbolic term is the
785  // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a
786  // constant. (Note that a constant is trivially symbolic).
787  if (!rhs.isSymbolicOrConstant() || isa<AffineConstantExpr>(lhs)) {
788  // At least one of them has to be symbolic.
789  return rhs * lhs;
790  }
791 
792  // At this point, if there was a constant, it would be on the right.
793 
794  // Multiplication with a one is a noop, return the other input.
795  if (rhsConst) {
796  if (rhsConst.getValue() == 1)
797  return lhs;
798  // Multiplication with zero.
799  if (rhsConst.getValue() == 0)
800  return rhsConst;
801  }
802 
803  // Fold successive multiplications: eg: (d0 * 2) * 3 into d0 * 6.
804  auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
805  if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Mul) {
806  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS()))
807  return lBin.getLHS() * (lrhs.getValue() * rhsConst.getValue());
808  }
809 
810  // When doing successive multiplication, bring constant to the right: turn (d0
811  // * 2) * d1 into (d0 * d1) * 2.
812  if (lBin && lBin.getKind() == AffineExprKind::Mul) {
813  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
814  return (lBin.getLHS() * rhs) * lrhs;
815  }
816  }
817 
818  return nullptr;
819 }
820 
822  return *this * getAffineConstantExpr(v, getContext());
823 }
825  if (auto simplified = simplifyMul(*this, other))
826  return simplified;
827 
829  return uniquer.get<AffineBinaryOpExprStorage>(
830  /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
831 }
832 
833 // Unary minus, delegate to operator*.
835  return *this * getAffineConstantExpr(-1, getContext());
836 }
837 
838 // Delegate to operator+.
839 AffineExpr AffineExpr::operator-(int64_t v) const { return *this + (-v); }
841  return *this + (-other);
842 }
843 
845  auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
846  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
847 
848  // mlir floordiv by zero or negative numbers is undefined and preserved as is.
849  if (!rhsConst || rhsConst.getValue() < 1)
850  return nullptr;
851 
852  if (lhsConst)
853  return getAffineConstantExpr(
854  divideFloorSigned(lhsConst.getValue(), rhsConst.getValue()),
855  lhs.getContext());
856 
857  // Fold floordiv of a multiply with a constant that is a multiple of the
858  // divisor. Eg: (i * 128) floordiv 64 = i * 2.
859  if (rhsConst == 1)
860  return lhs;
861 
862  // Simplify (expr * const) floordiv divConst when expr is known to be a
863  // multiple of divConst.
864  auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
865  if (lBin && lBin.getKind() == AffineExprKind::Mul) {
866  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
867  // rhsConst is known to be a positive constant.
868  if (lrhs.getValue() % rhsConst.getValue() == 0)
869  return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
870  }
871  }
872 
873  // Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is
874  // known to be a multiple of divConst.
875  if (lBin && lBin.getKind() == AffineExprKind::Add) {
876  int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
877  int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
878  // rhsConst is known to be a positive constant.
879  if (llhsDiv % rhsConst.getValue() == 0 ||
880  lrhsDiv % rhsConst.getValue() == 0)
881  return lBin.getLHS().floorDiv(rhsConst.getValue()) +
882  lBin.getRHS().floorDiv(rhsConst.getValue());
883  }
884 
885  return nullptr;
886 }
887 
888 AffineExpr AffineExpr::floorDiv(uint64_t v) const {
890 }
892  if (auto simplified = simplifyFloorDiv(*this, other))
893  return simplified;
894 
896  return uniquer.get<AffineBinaryOpExprStorage>(
897  /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this,
898  other);
899 }
900 
902  auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
903  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
904 
905  if (!rhsConst || rhsConst.getValue() < 1)
906  return nullptr;
907 
908  if (lhsConst)
909  return getAffineConstantExpr(
910  divideCeilSigned(lhsConst.getValue(), rhsConst.getValue()),
911  lhs.getContext());
912 
913  // Fold ceildiv of a multiply with a constant that is a multiple of the
914  // divisor. Eg: (i * 128) ceildiv 64 = i * 2.
915  if (rhsConst.getValue() == 1)
916  return lhs;
917 
918  // Simplify (expr * const) ceildiv divConst when const is known to be a
919  // multiple of divConst.
920  auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
921  if (lBin && lBin.getKind() == AffineExprKind::Mul) {
922  if (auto lrhs = dyn_cast<AffineConstantExpr>(lBin.getRHS())) {
923  // rhsConst is known to be a positive constant.
924  if (lrhs.getValue() % rhsConst.getValue() == 0)
925  return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
926  }
927  }
928 
929  return nullptr;
930 }
931 
932 AffineExpr AffineExpr::ceilDiv(uint64_t v) const {
934 }
936  if (auto simplified = simplifyCeilDiv(*this, other))
937  return simplified;
938 
940  return uniquer.get<AffineBinaryOpExprStorage>(
941  /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this,
942  other);
943 }
944 
946  auto lhsConst = dyn_cast<AffineConstantExpr>(lhs);
947  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
948 
949  // mod w.r.t zero or negative numbers is undefined and preserved as is.
950  if (!rhsConst || rhsConst.getValue() < 1)
951  return nullptr;
952 
953  if (lhsConst)
954  return getAffineConstantExpr(mod(lhsConst.getValue(), rhsConst.getValue()),
955  lhs.getContext());
956 
957  // Fold modulo of an expression that is known to be a multiple of a constant
958  // to zero if that constant is a multiple of the modulo factor. Eg: (i * 128)
959  // mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0.
960  if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0)
961  return getAffineConstantExpr(0, lhs.getContext());
962 
963  // Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is
964  // known to be a multiple of divConst.
965  auto lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
966  if (lBin && lBin.getKind() == AffineExprKind::Add) {
967  int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
968  int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
969  // rhsConst is known to be a positive constant.
970  if (llhsDiv % rhsConst.getValue() == 0)
971  return lBin.getRHS() % rhsConst.getValue();
972  if (lrhsDiv % rhsConst.getValue() == 0)
973  return lBin.getLHS() % rhsConst.getValue();
974  }
975 
976  // Simplify (e % a) % b to e % b when b evenly divides a
977  if (lBin && lBin.getKind() == AffineExprKind::Mod) {
978  auto intermediate = dyn_cast<AffineConstantExpr>(lBin.getRHS());
979  if (intermediate && intermediate.getValue() >= 1 &&
980  mod(intermediate.getValue(), rhsConst.getValue()) == 0) {
981  return lBin.getLHS() % rhsConst.getValue();
982  }
983  }
984 
985  return nullptr;
986 }
987 
989  return *this % getAffineConstantExpr(v, getContext());
990 }
992  if (auto simplified = simplifyMod(*this, other))
993  return simplified;
994 
996  return uniquer.get<AffineBinaryOpExprStorage>(
997  /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
998 }
999 
1001  SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(),
1002  map.getResults().end());
1003  return replaceDimsAndSymbols(dimReplacements, {});
1004 }
1005 raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) {
1006  expr.print(os);
1007  return os;
1008 }
1009 
1010 /// Constructs an affine expression from a flat ArrayRef. If there are local
1011 /// identifiers (neither dimensional nor symbolic) that appear in the sum of
1012 /// products expression, `localExprs` is expected to have the AffineExpr
1013 /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be
1014 /// in the format [dims, symbols, locals, constant term].
1016  unsigned numDims,
1017  unsigned numSymbols,
1018  ArrayRef<AffineExpr> localExprs,
1019  MLIRContext *context) {
1020  // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1.
1021  assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
1022  "unexpected number of local expressions");
1023 
1024  auto expr = getAffineConstantExpr(0, context);
1025  // Dimensions and symbols.
1026  for (unsigned j = 0; j < numDims + numSymbols; j++) {
1027  if (flatExprs[j] == 0)
1028  continue;
1029  auto id = j < numDims ? getAffineDimExpr(j, context)
1030  : getAffineSymbolExpr(j - numDims, context);
1031  expr = expr + id * flatExprs[j];
1032  }
1033 
1034  // Local identifiers.
1035  for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e;
1036  j++) {
1037  if (flatExprs[j] == 0)
1038  continue;
1039  auto term = localExprs[j - numDims - numSymbols] * flatExprs[j];
1040  expr = expr + term;
1041  }
1042 
1043  // Constant term.
1044  int64_t constTerm = flatExprs[flatExprs.size() - 1];
1045  if (constTerm != 0)
1046  expr = expr + constTerm;
1047  return expr;
1048 }
1049 
1050 /// Constructs a semi-affine expression from a flat ArrayRef. If there are
1051 /// local identifiers (neither dimensional nor symbolic) that appear in the sum
1052 /// of products expression, `localExprs` is expected to have the AffineExprs for
1053 /// it, and is substituted into. The ArrayRef `flatExprs` is expected to be in
1054 /// the format [dims, symbols, locals, constant term]. The semi-affine
1055 /// expression is constructed in the sorted order of dimension and symbol
1056 /// position numbers. Note: local expressions/ids are used for mod, div as well
1057 /// as symbolic RHS terms for terms that are not pure affine.
1059  unsigned numDims,
1060  unsigned numSymbols,
1061  ArrayRef<AffineExpr> localExprs,
1062  MLIRContext *context) {
1063  assert(!flatExprs.empty() && "flatExprs cannot be empty");
1064 
1065  // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1.
1066  assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
1067  "unexpected number of local expressions");
1068 
1069  AffineExpr expr = getAffineConstantExpr(0, context);
1070 
1071  // We design indices as a pair which help us present the semi-affine map as
1072  // sum of product where terms are sorted based on dimension or symbol
1073  // position: <keyA, keyB> for expressions of the form dimension * symbol,
1074  // where keyA is the position number of the dimension and keyB is the
1075  // position number of the symbol. For dimensional expressions we set the index
1076  // as (position number of the dimension, -1), as we want dimensional
1077  // expressions to appear before symbolic and product of dimensional and
1078  // symbolic expressions having the dimension with the same position number.
1079  // For symbolic expression set the index as (position number of the symbol,
1080  // maximum of last dimension and symbol position) number. For example, we want
1081  // the expression we are constructing to look something like: d0 + d0 * s0 +
1082  // s0 + d1*s1 + s1.
1083 
1084  // Stores the affine expression corresponding to a given index.
1086  // Stores the constant coefficient value corresponding to a given
1087  // dimension, symbol or a non-pure affine expression stored in `localExprs`.
1088  DenseMap<std::pair<unsigned, signed>, int64_t> coefficients;
1089  // Stores the indices as defined above, and later sorted to produce
1090  // the semi-affine expression in the desired form.
1092 
1093  // Example: expression = d0 + d0 * s0 + 2 * s0.
1094  // indices = [{0,-1}, {0, 0}, {0, 1}]
1095  // coefficients = [{{0, -1}, 1}, {{0, 0}, 1}, {{0, 1}, 2}]
1096  // indexToExprMap = [{{0, -1}, d0}, {{0, 0}, d0 * s0}, {{0, 1}, s0}]
1097 
1098  // Adds entries to `indexToExprMap`, `coefficients` and `indices`.
1099  auto addEntry = [&](std::pair<unsigned, signed> index, int64_t coefficient,
1100  AffineExpr expr) {
1101  assert(!llvm::is_contained(indices, index) &&
1102  "Key is already present in indices vector and overwriting will "
1103  "happen in `indexToExprMap` and `coefficients`!");
1104 
1105  indices.push_back(index);
1106  coefficients.insert({index, coefficient});
1107  indexToExprMap.insert({index, expr});
1108  };
1109 
1110  // Design indices for dimensional or symbolic terms, and store the indices,
1111  // constant coefficient corresponding to the indices in `coefficients` map,
1112  // and affine expression corresponding to indices in `indexToExprMap` map.
1113 
1114  // Ensure we do not have duplicate keys in `indexToExpr` map.
1115  unsigned offsetSym = 0;
1116  signed offsetDim = -1;
1117  for (unsigned j = numDims; j < numDims + numSymbols; ++j) {
1118  if (flatExprs[j] == 0)
1119  continue;
1120  // For symbolic expression set the index as <position number
1121  // of the symbol, max(dimCount, symCount)> number,
1122  // as we want symbolic expressions with the same positional number to
1123  // appear after dimensional expressions having the same positional number.
1124  std::pair<unsigned, signed> indexEntry(
1125  j - numDims, std::max(numDims, numSymbols) + offsetSym++);
1126  addEntry(indexEntry, flatExprs[j],
1127  getAffineSymbolExpr(j - numDims, context));
1128  }
1129 
1130  // Denotes semi-affine product, modulo or division terms, which has been added
1131  // to the `indexToExpr` map.
1132  SmallVector<bool, 4> addedToMap(flatExprs.size() - numDims - numSymbols - 1,
1133  false);
1134  unsigned lhsPos, rhsPos;
1135  // Construct indices for product terms involving dimension, symbol or constant
1136  // as lhs/rhs, and store the indices, constant coefficient corresponding to
1137  // the indices in `coefficients` map, and affine expression corresponding to
1138  // in indices in `indexToExprMap` map.
1139  for (const auto &it : llvm::enumerate(localExprs)) {
1140  AffineExpr expr = it.value();
1141  if (flatExprs[numDims + numSymbols + it.index()] == 0)
1142  continue;
1143  AffineExpr lhs = cast<AffineBinaryOpExpr>(expr).getLHS();
1144  AffineExpr rhs = cast<AffineBinaryOpExpr>(expr).getRHS();
1145  if (!((isa<AffineDimExpr>(lhs) || isa<AffineSymbolExpr>(lhs)) &&
1146  (isa<AffineDimExpr>(rhs) || isa<AffineSymbolExpr>(rhs) ||
1147  isa<AffineConstantExpr>(rhs)))) {
1148  continue;
1149  }
1150  if (isa<AffineConstantExpr>(rhs)) {
1151  // For product/modulo/division expressions, when rhs of modulo/division
1152  // expression is constant, we put 0 in place of keyB, because we want
1153  // them to appear earlier in the semi-affine expression we are
1154  // constructing. When rhs is constant, we place 0 in place of keyB.
1155  if (isa<AffineDimExpr>(lhs)) {
1156  lhsPos = cast<AffineDimExpr>(lhs).getPosition();
1157  std::pair<unsigned, signed> indexEntry(lhsPos, offsetDim--);
1158  addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
1159  expr);
1160  } else {
1161  lhsPos = cast<AffineSymbolExpr>(lhs).getPosition();
1162  std::pair<unsigned, signed> indexEntry(
1163  lhsPos, std::max(numDims, numSymbols) + offsetSym++);
1164  addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
1165  expr);
1166  }
1167  } else if (isa<AffineDimExpr>(lhs)) {
1168  // For product/modulo/division expressions having lhs as dimension and rhs
1169  // as symbol, we order the terms in the semi-affine expression based on
1170  // the pair: <keyA, keyB> for expressions of the form dimension * symbol,
1171  // where keyA is the position number of the dimension and keyB is the
1172  // position number of the symbol.
1173  lhsPos = cast<AffineDimExpr>(lhs).getPosition();
1174  rhsPos = cast<AffineSymbolExpr>(rhs).getPosition();
1175  std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos);
1176  addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
1177  } else {
1178  // For product/modulo/division expressions having both lhs and rhs as
1179  // symbol, we design indices as a pair: <keyA, keyB> for expressions
1180  // of the form dimension * symbol, where keyA is the position number of
1181  // the dimension and keyB is the position number of the symbol.
1182  lhsPos = cast<AffineSymbolExpr>(lhs).getPosition();
1183  rhsPos = cast<AffineSymbolExpr>(rhs).getPosition();
1184  std::pair<unsigned, signed> indexEntry(
1185  lhsPos, std::max(numDims, numSymbols) + offsetSym++);
1186  addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
1187  }
1188  addedToMap[it.index()] = true;
1189  }
1190 
1191  for (unsigned j = 0; j < numDims; ++j) {
1192  if (flatExprs[j] == 0)
1193  continue;
1194  // For dimensional expressions we set the index as <position number of the
1195  // dimension, 0>, as we want dimensional expressions to appear before
1196  // symbolic ones and products of dimensional and symbolic expressions
1197  // having the dimension with the same position number.
1198  std::pair<unsigned, signed> indexEntry(j, offsetDim--);
1199  addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context));
1200  }
1201 
1202  // Constructing the simplified semi-affine sum of product/division/mod
1203  // expression from the flattened form in the desired sorted order of indices
1204  // of the various individual product/division/mod expressions.
1205  llvm::sort(indices);
1206  for (const std::pair<unsigned, unsigned> index : indices) {
1207  assert(indexToExprMap.lookup(index) &&
1208  "cannot find key in `indexToExprMap` map");
1209  expr = expr + indexToExprMap.lookup(index) * coefficients.lookup(index);
1210  }
1211 
1212  // Local identifiers.
1213  for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e;
1214  j++) {
1215  // If the coefficient of the local expression is 0, continue as we need not
1216  // add it in out final expression.
1217  if (flatExprs[j] == 0 || addedToMap[j - numDims - numSymbols])
1218  continue;
1219  auto term = localExprs[j - numDims - numSymbols] * flatExprs[j];
1220  expr = expr + term;
1221  }
1222 
1223  // Constant term.
1224  int64_t constTerm = flatExprs.back();
1225  if (constTerm != 0)
1226  expr = expr + constTerm;
1227  return expr;
1228 }
1229 
1231  unsigned numSymbols)
1232  : numDims(numDims), numSymbols(numSymbols), numLocals(0) {
1233  operandExprStack.reserve(8);
1234 }
1235 
1236 // In pure affine t = expr * c, we multiply each coefficient of lhs with c.
1237 //
1238 // In case of semi affine multiplication expressions, t = expr * symbolic_expr,
1239 // introduce a local variable p (= expr * symbolic_expr), and the affine
1240 // expression expr * symbolic_expr is added to `localExprs`.
1242  assert(operandExprStack.size() >= 2);
1244  operandExprStack.pop_back();
1246 
1247  // Flatten semi-affine multiplication expressions by introducing a local
1248  // variable in place of the product; the affine expression
1249  // corresponding to the quantifier is added to `localExprs`.
1250  if (!isa<AffineConstantExpr>(expr.getRHS())) {
1251  SmallVector<int64_t, 8> mulLhs(lhs);
1252  MLIRContext *context = expr.getContext();
1254  localExprs, context);
1256  localExprs, context);
1257  return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
1258  }
1259 
1260  // Get the RHS constant.
1261  int64_t rhsConst = rhs[getConstantIndex()];
1262  for (int64_t &lhsElt : lhs)
1263  lhsElt *= rhsConst;
1264 
1265  return success();
1266 }
1267 
1269  assert(operandExprStack.size() >= 2);
1270  const auto &rhs = operandExprStack.back();
1271  auto &lhs = operandExprStack[operandExprStack.size() - 2];
1272  assert(lhs.size() == rhs.size());
1273  // Update the LHS in place.
1274  for (unsigned i = 0, e = rhs.size(); i < e; i++) {
1275  lhs[i] += rhs[i];
1276  }
1277  // Pop off the RHS.
1278  operandExprStack.pop_back();
1279  return success();
1280 }
1281 
1282 //
1283 // t = expr mod c <=> t = expr - c*q and c*q <= expr <= c*q + c - 1
1284 //
1285 // A mod expression "expr mod c" is thus flattened by introducing a new local
1286 // variable q (= expr floordiv c), such that expr mod c is replaced with
1287 // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
1288 //
1289 // In case of semi-affine modulo expressions, t = expr mod symbolic_expr,
1290 // introduce a local variable m (= expr mod symbolic_expr), and the affine
1291 // expression expr mod symbolic_expr is added to `localExprs`.
1293  assert(operandExprStack.size() >= 2);
1294 
1296  operandExprStack.pop_back();
1298  MLIRContext *context = expr.getContext();
1299 
1300  // Flatten semi affine modulo expressions by introducing a local
1301  // variable in place of the modulo value, and the affine expression
1302  // corresponding to the quantifier is added to `localExprs`.
1303  if (!isa<AffineConstantExpr>(expr.getRHS())) {
1304  SmallVector<int64_t, 8> modLhs(lhs);
1305  AffineExpr dividendExpr = getAffineExprFromFlatForm(
1306  lhs, numDims, numSymbols, localExprs, context);
1308  localExprs, context);
1309  AffineExpr modExpr = dividendExpr % divisorExpr;
1310  return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
1311  }
1312 
1313  int64_t rhsConst = rhs[getConstantIndex()];
1314  if (rhsConst <= 0)
1315  return failure();
1316 
1317  // Check if the LHS expression is a multiple of modulo factor.
1318  unsigned i, e;
1319  for (i = 0, e = lhs.size(); i < e; i++)
1320  if (lhs[i] % rhsConst != 0)
1321  break;
1322  // If yes, modulo expression here simplifies to zero.
1323  if (i == lhs.size()) {
1324  std::fill(lhs.begin(), lhs.end(), 0);
1325  return success();
1326  }
1327 
1328  // Add a local variable for the quotient, i.e., expr % c is replaced by
1329  // (expr - q * c) where q = expr floordiv c. Do this while canceling out
1330  // the GCD of expr and c.
1331  SmallVector<int64_t, 8> floorDividend(lhs);
1332  uint64_t gcd = rhsConst;
1333  for (int64_t lhsElt : lhs)
1334  gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
1335  // Simplify the numerator and the denominator.
1336  if (gcd != 1) {
1337  for (int64_t &floorDividendElt : floorDividend)
1338  floorDividendElt = floorDividendElt / static_cast<int64_t>(gcd);
1339  }
1340  int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd);
1341 
1342  // Construct the AffineExpr form of the floordiv to store in localExprs.
1343 
1344  AffineExpr dividendExpr = getAffineExprFromFlatForm(
1345  floorDividend, numDims, numSymbols, localExprs, context);
1346  AffineExpr divisorExpr = getAffineConstantExpr(floorDivisor, context);
1347  AffineExpr floorDivExpr = dividendExpr.floorDiv(divisorExpr);
1348  int loc;
1349  if ((loc = findLocalId(floorDivExpr)) == -1) {
1350  addLocalFloorDivId(floorDividend, floorDivisor, floorDivExpr);
1351  // Set result at top of stack to "lhs - rhsConst * q".
1352  lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst;
1353  } else {
1354  // Reuse the existing local id.
1355  lhs[getLocalVarStartIndex() + loc] = -rhsConst;
1356  }
1357  return success();
1358 }
1359 
1362  return visitDivExpr(expr, /*isCeil=*/true);
1363 }
1366  return visitDivExpr(expr, /*isCeil=*/false);
1367 }
1368 
1370  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
1371  auto &eq = operandExprStack.back();
1372  assert(expr.getPosition() < numDims && "Inconsistent number of dims");
1373  eq[getDimStartIndex() + expr.getPosition()] = 1;
1374  return success();
1375 }
1376 
1379  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
1380  auto &eq = operandExprStack.back();
1381  assert(expr.getPosition() < numSymbols && "inconsistent number of symbols");
1382  eq[getSymbolStartIndex() + expr.getPosition()] = 1;
1383  return success();
1384 }
1385 
1388  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
1389  auto &eq = operandExprStack.back();
1390  eq[getConstantIndex()] = expr.getValue();
1391  return success();
1392 }
1393 
1394 LogicalResult SimpleAffineExprFlattener::addLocalVariableSemiAffine(
1395  ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr,
1396  SmallVectorImpl<int64_t> &result, unsigned long resultSize) {
1397  assert(result.size() == resultSize &&
1398  "`result` vector passed is not of correct size");
1399  int loc;
1400  if ((loc = findLocalId(localExpr)) == -1) {
1401  if (failed(addLocalIdSemiAffine(lhs, rhs, localExpr)))
1402  return failure();
1403  }
1404  std::fill(result.begin(), result.end(), 0);
1405  if (loc == -1)
1406  result[getLocalVarStartIndex() + numLocals - 1] = 1;
1407  else
1408  result[getLocalVarStartIndex() + loc] = 1;
1409  return success();
1410 }
1411 
1412 // t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1
1413 // A floordiv is thus flattened by introducing a new local variable q, and
1414 // replacing that expression with 'q' while adding the constraints
1415 // c * q <= expr <= c * q + c - 1 to localVarCst (done by
1416 // IntegerRelation::addLocalFloorDiv).
1417 //
1418 // A ceildiv is similarly flattened:
1419 // t = expr ceildiv c <=> t = (expr + c - 1) floordiv c
1420 //
1421 // In case of semi affine division expressions, t = expr floordiv symbolic_expr
1422 // or t = expr ceildiv symbolic_expr, introduce a local variable q (= expr
1423 // floordiv/ceildiv symbolic_expr), and the affine floordiv/ceildiv is added to
1424 // `localExprs`.
1425 LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
1426  bool isCeil) {
1427  assert(operandExprStack.size() >= 2);
1428 
1429  MLIRContext *context = expr.getContext();
1431  operandExprStack.pop_back();
1433 
1434  // Flatten semi affine division expressions by introducing a local
1435  // variable in place of the quotient, and the affine expression corresponding
1436  // to the quantifier is added to `localExprs`.
1437  if (!isa<AffineConstantExpr>(expr.getRHS())) {
1438  SmallVector<int64_t, 8> divLhs(lhs);
1440  localExprs, context);
1442  localExprs, context);
1443  AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
1444  return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
1445  }
1446 
1447  // This is a pure affine expr; the RHS is a positive constant.
1448  int64_t rhsConst = rhs[getConstantIndex()];
1449  if (rhsConst <= 0)
1450  return failure();
1451 
1452  // Simplify the floordiv, ceildiv if possible by canceling out the greatest
1453  // common divisors of the numerator and denominator.
1454  uint64_t gcd = std::abs(rhsConst);
1455  for (int64_t lhsElt : lhs)
1456  gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
1457  // Simplify the numerator and the denominator.
1458  if (gcd != 1) {
1459  for (int64_t &lhsElt : lhs)
1460  lhsElt = lhsElt / static_cast<int64_t>(gcd);
1461  }
1462  int64_t divisor = rhsConst / static_cast<int64_t>(gcd);
1463  // If the divisor becomes 1, the updated LHS is the result. (The
1464  // divisor can't be negative since rhsConst is positive).
1465  if (divisor == 1)
1466  return success();
1467 
1468  // If the divisor cannot be simplified to one, we will have to retain
1469  // the ceil/floor expr (simplified up until here). Add an existential
1470  // quantifier to express its result, i.e., expr1 div expr2 is replaced
1471  // by a new identifier, q.
1472  AffineExpr a =
1474  AffineExpr b = getAffineConstantExpr(divisor, context);
1475 
1476  int loc;
1477  AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
1478  if ((loc = findLocalId(divExpr)) == -1) {
1479  if (!isCeil) {
1480  SmallVector<int64_t, 8> dividend(lhs);
1481  addLocalFloorDivId(dividend, divisor, divExpr);
1482  } else {
1483  // lhs ceildiv c <=> (lhs + c - 1) floordiv c
1484  SmallVector<int64_t, 8> dividend(lhs);
1485  dividend.back() += divisor - 1;
1486  addLocalFloorDivId(dividend, divisor, divExpr);
1487  }
1488  }
1489  // Set the expression on stack to the local var introduced to capture the
1490  // result of the division (floor or ceil).
1491  std::fill(lhs.begin(), lhs.end(), 0);
1492  if (loc == -1)
1493  lhs[getLocalVarStartIndex() + numLocals - 1] = 1;
1494  else
1495  lhs[getLocalVarStartIndex() + loc] = 1;
1496  return success();
1497 }
1498 
1499 // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
1500 // The local identifier added is always a floordiv of a pure add/mul affine
1501 // function of other identifiers, coefficients of which are specified in
1502 // dividend and with respect to a positive constant divisor. localExpr is the
1503 // simplified tree expression (AffineExpr) corresponding to the quantifier.
1505  int64_t divisor,
1506  AffineExpr localExpr) {
1507  assert(divisor > 0 && "positive constant divisor expected");
1508  for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
1509  subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
1510  localExprs.push_back(localExpr);
1511  numLocals++;
1512  // dividend and divisor are not used here; an override of this method uses it.
1513 }
1514 
1516  ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) {
1517  for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
1518  subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
1519  localExprs.push_back(localExpr);
1520  ++numLocals;
1521  // lhs and rhs are not used here; an override of this method uses them.
1522  return success();
1523 }
1524 
1525 int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
1527  if ((it = llvm::find(localExprs, localExpr)) == localExprs.end())
1528  return -1;
1529  return it - localExprs.begin();
1530 }
1531 
1532 /// Simplify the affine expression by flattening it and reconstructing it.
1534  unsigned numSymbols) {
1535  // Simplify semi-affine expressions separately.
1536  if (!expr.isPureAffine())
1537  expr = simplifySemiAffine(expr, numDims, numSymbols);
1538 
1539  SimpleAffineExprFlattener flattener(numDims, numSymbols);
1540  // has poison expression
1541  if (failed(flattener.walkPostOrder(expr)))
1542  return expr;
1543  ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
1544  if (!expr.isPureAffine() &&
1545  expr == getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
1546  flattener.localExprs,
1547  expr.getContext()))
1548  return expr;
1549  AffineExpr simplifiedExpr =
1550  expr.isPureAffine()
1551  ? getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
1552  flattener.localExprs, expr.getContext())
1553  : getSemiAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
1554  flattener.localExprs,
1555  expr.getContext());
1556 
1557  flattener.operandExprStack.pop_back();
1558  assert(flattener.operandExprStack.empty());
1559  return simplifiedExpr;
1560 }
1561 
1562 std::optional<int64_t> mlir::getBoundForAffineExpr(
1563  AffineExpr expr, unsigned numDims, unsigned numSymbols,
1564  ArrayRef<std::optional<int64_t>> constLowerBounds,
1565  ArrayRef<std::optional<int64_t>> constUpperBounds, bool isUpper) {
1566  // Handle divs and mods.
1567  if (auto binOpExpr = dyn_cast<AffineBinaryOpExpr>(expr)) {
1568  // If the LHS of a floor or ceil is bounded and the RHS is a constant, we
1569  // can compute an upper bound.
1570  if (binOpExpr.getKind() == AffineExprKind::FloorDiv) {
1571  auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
1572  if (!rhsConst || rhsConst.getValue() < 1)
1573  return std::nullopt;
1574  auto bound =
1575  getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
1576  constLowerBounds, constUpperBounds, isUpper);
1577  if (!bound)
1578  return std::nullopt;
1579  return divideFloorSigned(*bound, rhsConst.getValue());
1580  }
1581  if (binOpExpr.getKind() == AffineExprKind::CeilDiv) {
1582  auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
1583  if (rhsConst && rhsConst.getValue() >= 1) {
1584  auto bound =
1585  getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
1586  constLowerBounds, constUpperBounds, isUpper);
1587  if (!bound)
1588  return std::nullopt;
1589  return divideCeilSigned(*bound, rhsConst.getValue());
1590  }
1591  return std::nullopt;
1592  }
1593  if (binOpExpr.getKind() == AffineExprKind::Mod) {
1594  // lhs mod c is always <= c - 1 and non-negative. In addition, if `lhs` is
1595  // bounded such that lb <= lhs <= ub and lb floordiv c == ub floordiv c
1596  // (same "interval"), then lb mod c <= lhs mod c <= ub mod c.
1597  auto rhsConst = dyn_cast<AffineConstantExpr>(binOpExpr.getRHS());
1598  if (rhsConst && rhsConst.getValue() >= 1) {
1599  int64_t rhsConstVal = rhsConst.getValue();
1600  auto lb = getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
1601  constLowerBounds, constUpperBounds,
1602  /*isUpper=*/false);
1603  auto ub =
1604  getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
1605  constLowerBounds, constUpperBounds, isUpper);
1606  if (ub && lb &&
1607  divideFloorSigned(*lb, rhsConstVal) ==
1608  divideFloorSigned(*ub, rhsConstVal))
1609  return isUpper ? mod(*ub, rhsConstVal) : mod(*lb, rhsConstVal);
1610  return isUpper ? rhsConstVal - 1 : 0;
1611  }
1612  }
1613  }
1614  // Flatten the expression.
1615  SimpleAffineExprFlattener flattener(numDims, numSymbols);
1616  auto simpleResult = flattener.walkPostOrder(expr);
1617  // has poison expression
1618  if (failed(simpleResult))
1619  return std::nullopt;
1620  ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
1621  // TODO: Handle local variables. We can get hold of flattener.localExprs and
1622  // get bound on the local expr recursively.
1623  if (flattener.numLocals > 0)
1624  return std::nullopt;
1625  int64_t bound = 0;
1626  // Substitute the constant lower or upper bound for the dimensional or
1627  // symbolic input depending on `isUpper` to determine the bound.
1628  for (unsigned i = 0, e = numDims + numSymbols; i < e; ++i) {
1629  if (flattenedExpr[i] > 0) {
1630  auto &constBound = isUpper ? constUpperBounds[i] : constLowerBounds[i];
1631  if (!constBound)
1632  return std::nullopt;
1633  bound += *constBound * flattenedExpr[i];
1634  } else if (flattenedExpr[i] < 0) {
1635  auto &constBound = isUpper ? constLowerBounds[i] : constUpperBounds[i];
1636  if (!constBound)
1637  return std::nullopt;
1638  bound += *constBound * flattenedExpr[i];
1639  }
1640  }
1641  // Constant term.
1642  bound += flattenedExpr.back();
1643  return bound;
1644 }
static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos, AffineExprKind opKind)
Divides the given expression by the given symbol at position symbolPos.
Definition: AffineExpr.cpp:414
static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs)
Simplify a multiply expression. Return nullptr if it can't be simplified.
Definition: AffineExpr.cpp:773
static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs)
Definition: AffineExpr.cpp:945
static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs)
Simplify add expression. Return nullptr if it can't be simplified.
Definition: AffineExpr.cpp:645
static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef< int64_t > flatExprs, unsigned numDims, unsigned numSymbols, ArrayRef< AffineExpr > localExprs, MLIRContext *context)
Constructs a semi-affine expression from a flat ArrayRef.
static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs)
Definition: AffineExpr.cpp:901
static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs)
Definition: AffineExpr.cpp:844
static bool isNegatedAffineExpr(AffineExpr candidate, AffineExpr &expr)
Return "true" if candidate is a negated expression, i.e., Mul(-1, expr).
Definition: AffineExpr.cpp:482
static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position, MLIRContext *context)
Definition: AffineExpr.cpp:592
static bool isModOfModSubtraction(AffineExpr lhs, AffineExpr rhs, unsigned numDims, unsigned numSymbols)
Return "true" if lhs % rhs is guaranteed to evaluate to zero based on the fact that lhs contains anot...
Definition: AffineExpr.cpp:509
static void getSummandExprs(AffineExpr expr, SmallVector< AffineExpr > &result)
Populate result with all summand operands of given (potentially nested) addition.
Definition: AffineExpr.cpp:470
static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos, AffineExprKind opKind)
Returns true if the expression is divisible by the given symbol with position symbolPos.
Definition: AffineExpr.cpp:355
static AffineExpr simplifySemiAffine(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv operations when the second...
Definition: AffineExpr.cpp:547
static MLIRContext * getContext(OpFoldResult val)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Affine binary operation expression.
Definition: AffineExpr.h:228
AffineExpr getLHS() const
Definition: AffineExpr.cpp:336
AffineBinaryOpExpr(AffineExpr::ImplType *ptr)
Definition: AffineExpr.cpp:334
AffineExpr getRHS() const
Definition: AffineExpr.cpp:339
An integer constant appearing in affine expression.
Definition: AffineExpr.h:253
AffineConstantExpr(AffineExpr::ImplType *ptr=nullptr)
Definition: AffineExpr.cpp:617
int64_t getValue() const
Definition: AffineExpr.cpp:619
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:237
AffineDimExpr(AffineExpr::ImplType *ptr)
Definition: AffineExpr.cpp:343
unsigned getPosition() const
Definition: AffineExpr.cpp:344
See documentation for AffineExprVisitorBase.
RetTy walkPostOrder(AffineExpr expr)
Base type for affine expression.
Definition: AffineExpr.h:69
AffineExpr replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineExpr.cpp:85
AffineExpr shiftDims(unsigned numDims, unsigned shift, unsigned offset=0) const
Replace dims[offset ...
Definition: AffineExpr.cpp:129
AffineExpr operator+(int64_t v) const
Definition: AffineExpr.cpp:760
bool isSymbolicOrConstant() const
Returns true if this expression is made out of only symbols and constants, i.e., it does not involve ...
Definition: AffineExpr.cpp:184
AffineExpr operator*(int64_t v) const
Definition: AffineExpr.cpp:821
bool operator==(AffineExpr other) const
Definition: AffineExpr.h:77
bool isPureAffine() const
Returns true if this is a pure affine expression, i.e., multiplication, floordiv, ceildiv,...
Definition: AffineExpr.cpp:208
AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift, unsigned offset=0) const
Replace symbols[offset ...
Definition: AffineExpr.cpp:141
AffineExpr operator-() const
Definition: AffineExpr.cpp:834
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:888
ImplType * expr
Definition: AffineExpr.h:210
RetT walk(FnT &&callback) const
Walk all of the AffineExpr's in this expression in postorder.
Definition: AffineExpr.h:131
AffineExprKind getKind() const
Return the classification for this type.
Definition: AffineExpr.cpp:31
bool isMultipleOf(int64_t factor) const
Return true if the affine expression is a multiple of 'factor'.
Definition: AffineExpr.cpp:279
int64_t getLargestKnownDivisor() const
Returns the greatest known integral divisor of this affine expression.
Definition: AffineExpr.cpp:239
AffineExpr compose(AffineMap map) const
Compose with an AffineMap.
bool isFunctionOfDim(unsigned position) const
Return true if the affine expression involves AffineDimExpr position.
Definition: AffineExpr.cpp:312
bool isFunctionOfSymbol(unsigned position) const
Return true if the affine expression involves AffineSymbolExpr position.
Definition: AffineExpr.cpp:323
AffineExpr replaceDims(ArrayRef< AffineExpr > dimReplacements) const
Dim-only version of replaceDimsAndSymbols.
Definition: AffineExpr.cpp:118
AffineExpr operator%(uint64_t v) const
Definition: AffineExpr.cpp:988
MLIRContext * getContext() const
Definition: AffineExpr.cpp:29
AffineExpr replace(AffineExpr expr, AffineExpr replacement) const
Sparse replace method.
Definition: AffineExpr.cpp:177
AffineExpr replaceSymbols(ArrayRef< AffineExpr > symReplacements) const
Symbol-only version of replaceDimsAndSymbols.
Definition: AffineExpr.cpp:123
AffineExpr ceilDiv(uint64_t v) const
Definition: AffineExpr.cpp:932
void print(raw_ostream &os) const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:395
A symbolic identifier appearing in an affine expression.
Definition: AffineExpr.h:245
AffineSymbolExpr(AffineExpr::ImplType *ptr)
Definition: AffineExpr.cpp:607
unsigned getPosition() const
Definition: AffineExpr.cpp:609
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
StorageUniquer & getAffineUniquer()
Returns the storage uniquer used for creating affine constructs.
virtual void addLocalFloorDivId(ArrayRef< int64_t > dividend, int64_t divisor, AffineExpr localExpr)
LogicalResult visitSymbolExpr(AffineSymbolExpr expr)
std::vector< SmallVector< int64_t, 8 > > operandExprStack
LogicalResult visitDimExpr(AffineDimExpr expr)
LogicalResult visitFloorDivExpr(AffineBinaryOpExpr expr)
LogicalResult visitConstantExpr(AffineConstantExpr expr)
virtual LogicalResult addLocalIdSemiAffine(ArrayRef< int64_t > lhs, ArrayRef< int64_t > rhs, AffineExpr localExpr)
Add a local identifier (needed to flatten a mod, floordiv, ceildiv, mul expr) when the rhs is a symbo...
LogicalResult visitModExpr(AffineBinaryOpExpr expr)
LogicalResult visitAddExpr(AffineBinaryOpExpr expr)
LogicalResult visitCeilDivExpr(AffineBinaryOpExpr expr)
LogicalResult visitMulExpr(AffineBinaryOpExpr expr)
SmallVector< AffineExpr, 4 > localExprs
SimpleAffineExprFlattener(unsigned numDims, unsigned numSymbols)
A utility class to get or create instances of "storage classes".
Storage * get(function_ref< void(Storage *)> initFn, TypeID id, Args &&...args)
Gets a uniqued instance of 'Storage'.
A utility result that is used to signal how to proceed with an ongoing walk:
Definition: Visitors.h:34
Detect if any of the given parameter types has a sub-element handler.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
Fraction abs(const Fraction &f)
Definition: Fraction.h:106
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::optional< int64_t > getBoundForAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, ArrayRef< std::optional< int64_t >> constLowerBounds, ArrayRef< std::optional< int64_t >> constUpperBounds, bool isUpper)
Get a lower or upper (depending on isUpper) bound for expr while using the constant lower and upper b...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
AffineExprKind
Definition: AffineExpr.h:41
@ 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.
AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr rhs)
Definition: AffineExpr.cpp:66
AffineExpr getAffineExprFromFlatForm(ArrayRef< int64_t > flatExprs, unsigned numDims, unsigned numSymbols, ArrayRef< AffineExpr > localExprs, MLIRContext *context)
Constructs an affine expression from a flat ArrayRef.
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:627
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify an affine expression by flattening and some amount of simple analysis.
SmallVector< AffineExpr > getAffineConstantExprs(ArrayRef< int64_t > constants, MLIRContext *context)
Definition: AffineExpr.cpp:637
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:603
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)
Definition: AffineExpr.cpp:613
raw_ostream & operator<<(raw_ostream &os, const AliasResult &result)
Definition: AliasAnalysis.h:78
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
A binary operation appearing in an affine expression.
An integer constant appearing in affine expression.
A dimensional or symbolic identifier appearing in an affine expression.
Base storage class appearing in an affine expression.
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.