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 "mlir/IR/AffineExpr.h"
10 #include "AffineExprDetail.h"
11 #include "mlir/IR/AffineExprVisitor.h"
12 #include "mlir/IR/AffineMap.h"
13 #include "mlir/IR/IntegerSet.h"
14 #include "mlir/Support/MathExtras.h"
15 #include "mlir/Support/TypeID.h"
16 #include "llvm/ADT/STLExtras.h"
17 
18 using namespace mlir;
19 using namespace mlir::detail;
20 
21 MLIRContext *AffineExpr::getContext() const { return expr->context; }
22 
23 AffineExprKind AffineExpr::getKind() const { return expr->kind; }
24 
25 /// Walk all of the AffineExprs in this subgraph in postorder.
26 void AffineExpr::walk(std::function<void(AffineExpr)> callback) const {
27   struct AffineExprWalker : public AffineExprVisitor<AffineExprWalker> {
28     std::function<void(AffineExpr)> callback;
29 
30     AffineExprWalker(std::function<void(AffineExpr)> callback)
31         : callback(callback) {}
32 
33     void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) { callback(expr); }
34     void visitConstantExpr(AffineConstantExpr expr) { callback(expr); }
35     void visitDimExpr(AffineDimExpr expr) { callback(expr); }
36     void visitSymbolExpr(AffineSymbolExpr expr) { callback(expr); }
37   };
38 
39   AffineExprWalker(callback).walkPostOrder(*this);
40 }
41 
42 // Dispatch affine expression construction based on kind.
43 AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs,
44                                        AffineExpr rhs) {
45   if (kind == AffineExprKind::Add)
46     return lhs + rhs;
47   if (kind == AffineExprKind::Mul)
48     return lhs * rhs;
49   if (kind == AffineExprKind::FloorDiv)
50     return lhs.floorDiv(rhs);
51   if (kind == AffineExprKind::CeilDiv)
52     return lhs.ceilDiv(rhs);
53   if (kind == AffineExprKind::Mod)
54     return lhs % rhs;
55 
56   llvm_unreachable("unknown binary operation on affine expressions");
57 }
58 
59 /// This method substitutes any uses of dimensions and symbols (e.g.
60 /// dim#0 with dimReplacements[0]) and returns the modified expression tree.
61 AffineExpr
62 AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
63                                   ArrayRef<AffineExpr> symReplacements) const {
64   switch (getKind()) {
65   case AffineExprKind::Constant:
66     return *this;
67   case AffineExprKind::DimId: {
68     unsigned dimId = cast<AffineDimExpr>().getPosition();
69     if (dimId >= dimReplacements.size())
70       return *this;
71     return dimReplacements[dimId];
72   }
73   case AffineExprKind::SymbolId: {
74     unsigned symId = cast<AffineSymbolExpr>().getPosition();
75     if (symId >= symReplacements.size())
76       return *this;
77     return symReplacements[symId];
78   }
79   case AffineExprKind::Add:
80   case AffineExprKind::Mul:
81   case AffineExprKind::FloorDiv:
82   case AffineExprKind::CeilDiv:
83   case AffineExprKind::Mod:
84     auto binOp = cast<AffineBinaryOpExpr>();
85     auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
86     auto newLHS = lhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
87     auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
88     if (newLHS == lhs && newRHS == rhs)
89       return *this;
90     return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
91   }
92   llvm_unreachable("Unknown AffineExpr");
93 }
94 
95 /// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
96 AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift) const {
97   SmallVector<AffineExpr, 4> dims;
98   for (unsigned idx = 0; idx < numDims; ++idx)
99     dims.push_back(getAffineDimExpr(idx + shift, getContext()));
100   return replaceDimsAndSymbols(dims, {});
101 }
102 
103 /// Replace symbols[0 .. numSymbols - 1] by
104 /// symbols[shift .. shift + numSymbols - 1].
105 AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
106   SmallVector<AffineExpr, 4> symbols;
107   for (unsigned idx = 0; idx < numSymbols; ++idx)
108     symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
109   return replaceDimsAndSymbols({}, symbols);
110 }
111 
112 /// Sparse replace method. Return the modified expression tree.
113 AffineExpr
114 AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
115   auto it = map.find(*this);
116   if (it != map.end())
117     return it->second;
118   switch (getKind()) {
119   default:
120     return *this;
121   case AffineExprKind::Add:
122   case AffineExprKind::Mul:
123   case AffineExprKind::FloorDiv:
124   case AffineExprKind::CeilDiv:
125   case AffineExprKind::Mod:
126     auto binOp = cast<AffineBinaryOpExpr>();
127     auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
128     auto newLHS = lhs.replace(map);
129     auto newRHS = rhs.replace(map);
130     if (newLHS == lhs && newRHS == rhs)
131       return *this;
132     return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
133   }
134   llvm_unreachable("Unknown AffineExpr");
135 }
136 
137 /// Sparse replace method. Return the modified expression tree.
138 AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const {
139   DenseMap<AffineExpr, AffineExpr> map;
140   map.insert(std::make_pair(expr, replacement));
141   return replace(map);
142 }
143 /// Returns true if this expression is made out of only symbols and
144 /// constants (no dimensional identifiers).
145 bool AffineExpr::isSymbolicOrConstant() const {
146   switch (getKind()) {
147   case AffineExprKind::Constant:
148     return true;
149   case AffineExprKind::DimId:
150     return false;
151   case AffineExprKind::SymbolId:
152     return true;
153 
154   case AffineExprKind::Add:
155   case AffineExprKind::Mul:
156   case AffineExprKind::FloorDiv:
157   case AffineExprKind::CeilDiv:
158   case AffineExprKind::Mod: {
159     auto expr = this->cast<AffineBinaryOpExpr>();
160     return expr.getLHS().isSymbolicOrConstant() &&
161            expr.getRHS().isSymbolicOrConstant();
162   }
163   }
164   llvm_unreachable("Unknown AffineExpr");
165 }
166 
167 /// Returns true if this is a pure affine expression, i.e., multiplication,
168 /// floordiv, ceildiv, and mod is only allowed w.r.t constants.
169 bool AffineExpr::isPureAffine() const {
170   switch (getKind()) {
171   case AffineExprKind::SymbolId:
172   case AffineExprKind::DimId:
173   case AffineExprKind::Constant:
174     return true;
175   case AffineExprKind::Add: {
176     auto op = cast<AffineBinaryOpExpr>();
177     return op.getLHS().isPureAffine() && op.getRHS().isPureAffine();
178   }
179 
180   case AffineExprKind::Mul: {
181     // TODO: Canonicalize the constants in binary operators to the RHS when
182     // possible, allowing this to merge into the next case.
183     auto op = cast<AffineBinaryOpExpr>();
184     return op.getLHS().isPureAffine() && op.getRHS().isPureAffine() &&
185            (op.getLHS().template isa<AffineConstantExpr>() ||
186             op.getRHS().template isa<AffineConstantExpr>());
187   }
188   case AffineExprKind::FloorDiv:
189   case AffineExprKind::CeilDiv:
190   case AffineExprKind::Mod: {
191     auto op = cast<AffineBinaryOpExpr>();
192     return op.getLHS().isPureAffine() &&
193            op.getRHS().template isa<AffineConstantExpr>();
194   }
195   }
196   llvm_unreachable("Unknown AffineExpr");
197 }
198 
199 // Returns the greatest known integral divisor of this affine expression.
200 int64_t AffineExpr::getLargestKnownDivisor() const {
201   AffineBinaryOpExpr binExpr(nullptr);
202   switch (getKind()) {
203   case AffineExprKind::SymbolId:
204     LLVM_FALLTHROUGH;
205   case AffineExprKind::DimId:
206     return 1;
207   case AffineExprKind::Constant:
208     return std::abs(this->cast<AffineConstantExpr>().getValue());
209   case AffineExprKind::Mul: {
210     binExpr = this->cast<AffineBinaryOpExpr>();
211     return binExpr.getLHS().getLargestKnownDivisor() *
212            binExpr.getRHS().getLargestKnownDivisor();
213   }
214   case AffineExprKind::Add:
215     LLVM_FALLTHROUGH;
216   case AffineExprKind::FloorDiv:
217   case AffineExprKind::CeilDiv:
218   case AffineExprKind::Mod: {
219     binExpr = cast<AffineBinaryOpExpr>();
220     return llvm::GreatestCommonDivisor64(
221         binExpr.getLHS().getLargestKnownDivisor(),
222         binExpr.getRHS().getLargestKnownDivisor());
223   }
224   }
225   llvm_unreachable("Unknown AffineExpr");
226 }
227 
228 bool AffineExpr::isMultipleOf(int64_t factor) const {
229   AffineBinaryOpExpr binExpr(nullptr);
230   uint64_t l, u;
231   switch (getKind()) {
232   case AffineExprKind::SymbolId:
233     LLVM_FALLTHROUGH;
234   case AffineExprKind::DimId:
235     return factor * factor == 1;
236   case AffineExprKind::Constant:
237     return cast<AffineConstantExpr>().getValue() % factor == 0;
238   case AffineExprKind::Mul: {
239     binExpr = cast<AffineBinaryOpExpr>();
240     // It's probably not worth optimizing this further (to not traverse the
241     // whole sub-tree under - it that would require a version of isMultipleOf
242     // that on a 'false' return also returns the largest known divisor).
243     return (l = binExpr.getLHS().getLargestKnownDivisor()) % factor == 0 ||
244            (u = binExpr.getRHS().getLargestKnownDivisor()) % factor == 0 ||
245            (l * u) % factor == 0;
246   }
247   case AffineExprKind::Add:
248   case AffineExprKind::FloorDiv:
249   case AffineExprKind::CeilDiv:
250   case AffineExprKind::Mod: {
251     binExpr = cast<AffineBinaryOpExpr>();
252     return llvm::GreatestCommonDivisor64(
253                binExpr.getLHS().getLargestKnownDivisor(),
254                binExpr.getRHS().getLargestKnownDivisor()) %
255                factor ==
256            0;
257   }
258   }
259   llvm_unreachable("Unknown AffineExpr");
260 }
261 
262 bool AffineExpr::isFunctionOfDim(unsigned position) const {
263   if (getKind() == AffineExprKind::DimId) {
264     return *this == mlir::getAffineDimExpr(position, getContext());
265   }
266   if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) {
267     return expr.getLHS().isFunctionOfDim(position) ||
268            expr.getRHS().isFunctionOfDim(position);
269   }
270   return false;
271 }
272 
273 bool AffineExpr::isFunctionOfSymbol(unsigned position) const {
274   if (getKind() == AffineExprKind::SymbolId) {
275     return *this == mlir::getAffineSymbolExpr(position, getContext());
276   }
277   if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) {
278     return expr.getLHS().isFunctionOfSymbol(position) ||
279            expr.getRHS().isFunctionOfSymbol(position);
280   }
281   return false;
282 }
283 
284 AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr)
285     : AffineExpr(ptr) {}
286 AffineExpr AffineBinaryOpExpr::getLHS() const {
287   return static_cast<ImplType *>(expr)->lhs;
288 }
289 AffineExpr AffineBinaryOpExpr::getRHS() const {
290   return static_cast<ImplType *>(expr)->rhs;
291 }
292 
293 AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : AffineExpr(ptr) {}
294 unsigned AffineDimExpr::getPosition() const {
295   return static_cast<ImplType *>(expr)->position;
296 }
297 
298 /// Returns true if the expression is divisible by the given symbol with
299 /// position `symbolPos`. The argument `opKind` specifies here what kind of
300 /// division or mod operation called this division. It helps in implementing the
301 /// commutative property of the floordiv and ceildiv operations. If the argument
302 ///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv
303 /// operation, then the commutative property can be used otherwise, the floordiv
304 /// operation is not divisible. The same argument holds for ceildiv operation.
305 static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos,
306                                 AffineExprKind opKind) {
307   // The argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
308   assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
309           opKind == AffineExprKind::CeilDiv) &&
310          "unexpected opKind");
311   switch (expr.getKind()) {
312   case AffineExprKind::Constant:
313     if (expr.cast<AffineConstantExpr>().getValue())
314       return false;
315     return true;
316   case AffineExprKind::DimId:
317     return false;
318   case AffineExprKind::SymbolId:
319     return (expr.cast<AffineSymbolExpr>().getPosition() == symbolPos);
320   // Checks divisibility by the given symbol for both operands.
321   case AffineExprKind::Add: {
322     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
323     return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) &&
324            isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
325   }
326   // Checks divisibility by the given symbol for both operands. Consider the
327   // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`,
328   // this is a division by s1 and both the operands of modulo are divisible by
329   // s1 but it is not divisible by s1 always. The third argument is
330   // `AffineExprKind::Mod` for this reason.
331   case AffineExprKind::Mod: {
332     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
333     return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos,
334                                AffineExprKind::Mod) &&
335            isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos,
336                                AffineExprKind::Mod);
337   }
338   // Checks if any of the operand divisible by the given symbol.
339   case AffineExprKind::Mul: {
340     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
341     return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) ||
342            isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
343   }
344   // Floordiv and ceildiv are divisible by the given symbol when the first
345   // operand is divisible, and the affine expression kind of the argument expr
346   // is same as the argument `opKind`. This can be inferred from commutative
347   // property of floordiv and ceildiv operations and are as follow:
348   // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
349   // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
350   // It will fail if operations are not same. For example:
351   // (exps1 ceildiv exp2) floordiv exp3 can not be simplified.
352   case AffineExprKind::FloorDiv:
353   case AffineExprKind::CeilDiv: {
354     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
355     if (opKind != expr.getKind())
356       return false;
357     return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind());
358   }
359   }
360   llvm_unreachable("Unknown AffineExpr");
361 }
362 
363 /// Divides the given expression by the given symbol at position `symbolPos`. It
364 /// considers the divisibility condition is checked before calling itself. A
365 /// null expression is returned whenever the divisibility condition fails.
366 static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos,
367                                  AffineExprKind opKind) {
368   // THe argument `opKind` can either be Modulo, Floordiv or Ceildiv only.
369   assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
370           opKind == AffineExprKind::CeilDiv) &&
371          "unexpected opKind");
372   switch (expr.getKind()) {
373   case AffineExprKind::Constant:
374     if (expr.cast<AffineConstantExpr>().getValue() != 0)
375       return nullptr;
376     return getAffineConstantExpr(0, expr.getContext());
377   case AffineExprKind::DimId:
378     return nullptr;
379   case AffineExprKind::SymbolId:
380     return getAffineConstantExpr(1, expr.getContext());
381   // Dividing both operands by the given symbol.
382   case AffineExprKind::Add: {
383     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
384     return getAffineBinaryOpExpr(
385         expr.getKind(), symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind),
386         symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind));
387   }
388   // Dividing both operands by the given symbol.
389   case AffineExprKind::Mod: {
390     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
391     return getAffineBinaryOpExpr(
392         expr.getKind(),
393         symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
394         symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind()));
395   }
396   // Dividing any of the operand by the given symbol.
397   case AffineExprKind::Mul: {
398     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
399     if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind))
400       return binaryExpr.getLHS() *
401              symbolicDivide(binaryExpr.getRHS(), symbolPos, opKind);
402     return symbolicDivide(binaryExpr.getLHS(), symbolPos, opKind) *
403            binaryExpr.getRHS();
404   }
405   // Dividing first operand only by the given symbol.
406   case AffineExprKind::FloorDiv:
407   case AffineExprKind::CeilDiv: {
408     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
409     return getAffineBinaryOpExpr(
410         expr.getKind(),
411         symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
412         binaryExpr.getRHS());
413   }
414   }
415   llvm_unreachable("Unknown AffineExpr");
416 }
417 
418 /// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv
419 /// operations when the second operand simplifies to a symbol and the first
420 /// operand is divisible by that symbol. It can be applied to any semi-affine
421 /// expression. Returned expression can either be a semi-affine or pure affine
422 /// expression.
423 static AffineExpr simplifySemiAffine(AffineExpr expr) {
424   switch (expr.getKind()) {
425   case AffineExprKind::Constant:
426   case AffineExprKind::DimId:
427   case AffineExprKind::SymbolId:
428     return expr;
429   case AffineExprKind::Add:
430   case AffineExprKind::Mul: {
431     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
432     return getAffineBinaryOpExpr(expr.getKind(),
433                                  simplifySemiAffine(binaryExpr.getLHS()),
434                                  simplifySemiAffine(binaryExpr.getRHS()));
435   }
436   // Check if the simplification of the second operand is a symbol, and the
437   // first operand is divisible by it. If the operation is a modulo, a constant
438   // zero expression is returned. In the case of floordiv and ceildiv, the
439   // symbol from the simplification of the second operand divides the first
440   // operand. Otherwise, simplification is not possible.
441   case AffineExprKind::FloorDiv:
442   case AffineExprKind::CeilDiv:
443   case AffineExprKind::Mod: {
444     AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
445     AffineExpr sLHS = simplifySemiAffine(binaryExpr.getLHS());
446     AffineExpr sRHS = simplifySemiAffine(binaryExpr.getRHS());
447     AffineSymbolExpr symbolExpr =
448         simplifySemiAffine(binaryExpr.getRHS()).dyn_cast<AffineSymbolExpr>();
449     if (!symbolExpr)
450       return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
451     unsigned symbolPos = symbolExpr.getPosition();
452     if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind()))
453       return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
454     if (expr.getKind() == AffineExprKind::Mod)
455       return getAffineConstantExpr(0, expr.getContext());
456     return symbolicDivide(sLHS, symbolPos, expr.getKind());
457   }
458   }
459   llvm_unreachable("Unknown AffineExpr");
460 }
461 
462 static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
463                                        MLIRContext *context) {
464   auto assignCtx = [context](AffineDimExprStorage *storage) {
465     storage->context = context;
466   };
467 
468   StorageUniquer &uniquer = context->getAffineUniquer();
469   return uniquer.get<AffineDimExprStorage>(
470       assignCtx, static_cast<unsigned>(kind), position);
471 }
472 
473 AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) {
474   return getAffineDimOrSymbol(AffineExprKind::DimId, position, context);
475 }
476 
477 AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr)
478     : AffineExpr(ptr) {}
479 unsigned AffineSymbolExpr::getPosition() const {
480   return static_cast<ImplType *>(expr)->position;
481 }
482 
483 AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) {
484   return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context);
485   ;
486 }
487 
488 AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr)
489     : AffineExpr(ptr) {}
490 int64_t AffineConstantExpr::getValue() const {
491   return static_cast<ImplType *>(expr)->constant;
492 }
493 
494 bool AffineExpr::operator==(int64_t v) const {
495   return *this == getAffineConstantExpr(v, getContext());
496 }
497 
498 AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) {
499   auto assignCtx = [context](AffineConstantExprStorage *storage) {
500     storage->context = context;
501   };
502 
503   StorageUniquer &uniquer = context->getAffineUniquer();
504   return uniquer.get<AffineConstantExprStorage>(assignCtx, constant);
505 }
506 
507 /// Simplify add expression. Return nullptr if it can't be simplified.
508 static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
509   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
510   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
511   // Fold if both LHS, RHS are a constant.
512   if (lhsConst && rhsConst)
513     return getAffineConstantExpr(lhsConst.getValue() + rhsConst.getValue(),
514                                  lhs.getContext());
515 
516   // Canonicalize so that only the RHS is a constant. (4 + d0 becomes d0 + 4).
517   // If only one of them is a symbolic expressions, make it the RHS.
518   if (lhs.isa<AffineConstantExpr>() ||
519       (lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())) {
520     return rhs + lhs;
521   }
522 
523   // At this point, if there was a constant, it would be on the right.
524 
525   // Addition with a zero is a noop, return the other input.
526   if (rhsConst) {
527     if (rhsConst.getValue() == 0)
528       return lhs;
529   }
530   // Fold successive additions like (d0 + 2) + 3 into d0 + 5.
531   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
532   if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Add) {
533     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>())
534       return lBin.getLHS() + (lrhs.getValue() + rhsConst.getValue());
535   }
536 
537   // Detect "c1 * expr + c_2 * expr" as "(c1 + c2) * expr".
538   // c1 is rRhsConst, c2 is rLhsConst; firstExpr, secondExpr are their
539   // respective multiplicands.
540   Optional<int64_t> rLhsConst, rRhsConst;
541   AffineExpr firstExpr, secondExpr;
542   AffineConstantExpr rLhsConstExpr;
543   auto lBinOpExpr = lhs.dyn_cast<AffineBinaryOpExpr>();
544   if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Mul &&
545       (rLhsConstExpr = lBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) {
546     rLhsConst = rLhsConstExpr.getValue();
547     firstExpr = lBinOpExpr.getLHS();
548   } else {
549     rLhsConst = 1;
550     firstExpr = lhs;
551   }
552 
553   auto rBinOpExpr = rhs.dyn_cast<AffineBinaryOpExpr>();
554   AffineConstantExpr rRhsConstExpr;
555   if (rBinOpExpr && rBinOpExpr.getKind() == AffineExprKind::Mul &&
556       (rRhsConstExpr = rBinOpExpr.getRHS().dyn_cast<AffineConstantExpr>())) {
557     rRhsConst = rRhsConstExpr.getValue();
558     secondExpr = rBinOpExpr.getLHS();
559   } else {
560     rRhsConst = 1;
561     secondExpr = rhs;
562   }
563 
564   if (rLhsConst && rRhsConst && firstExpr == secondExpr)
565     return getAffineBinaryOpExpr(
566         AffineExprKind::Mul, firstExpr,
567         getAffineConstantExpr(rLhsConst.getValue() + rRhsConst.getValue(),
568                               lhs.getContext()));
569 
570   // When doing successive additions, bring constant to the right: turn (d0 + 2)
571   // + d1 into (d0 + d1) + 2.
572   if (lBin && lBin.getKind() == AffineExprKind::Add) {
573     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
574       return lBin.getLHS() + rhs + lrhs;
575     }
576   }
577 
578   // Detect and transform "expr - c * (expr floordiv c)" to "expr mod c". This
579   // leads to a much more efficient form when 'c' is a power of two, and in
580   // general a more compact and readable form.
581 
582   // Process '(expr floordiv c) * (-c)'.
583   if (!rBinOpExpr)
584     return nullptr;
585 
586   auto lrhs = rBinOpExpr.getLHS();
587   auto rrhs = rBinOpExpr.getRHS();
588 
589   // Process lrhs, which is 'expr floordiv c'.
590   AffineBinaryOpExpr lrBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>();
591   if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
592     return nullptr;
593 
594   auto llrhs = lrBinOpExpr.getLHS();
595   auto rlrhs = lrBinOpExpr.getRHS();
596 
597   if (lhs == llrhs && rlrhs == -rrhs) {
598     return lhs % rlrhs;
599   }
600   return nullptr;
601 }
602 
603 AffineExpr AffineExpr::operator+(int64_t v) const {
604   return *this + getAffineConstantExpr(v, getContext());
605 }
606 AffineExpr AffineExpr::operator+(AffineExpr other) const {
607   if (auto simplified = simplifyAdd(*this, other))
608     return simplified;
609 
610   StorageUniquer &uniquer = getContext()->getAffineUniquer();
611   return uniquer.get<AffineBinaryOpExprStorage>(
612       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
613 }
614 
615 /// Simplify a multiply expression. Return nullptr if it can't be simplified.
616 static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) {
617   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
618   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
619 
620   if (lhsConst && rhsConst)
621     return getAffineConstantExpr(lhsConst.getValue() * rhsConst.getValue(),
622                                  lhs.getContext());
623 
624   assert(lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant());
625 
626   // Canonicalize the mul expression so that the constant/symbolic term is the
627   // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a
628   // constant. (Note that a constant is trivially symbolic).
629   if (!rhs.isSymbolicOrConstant() || lhs.isa<AffineConstantExpr>()) {
630     // At least one of them has to be symbolic.
631     return rhs * lhs;
632   }
633 
634   // At this point, if there was a constant, it would be on the right.
635 
636   // Multiplication with a one is a noop, return the other input.
637   if (rhsConst) {
638     if (rhsConst.getValue() == 1)
639       return lhs;
640     // Multiplication with zero.
641     if (rhsConst.getValue() == 0)
642       return rhsConst;
643   }
644 
645   // Fold successive multiplications: eg: (d0 * 2) * 3 into d0 * 6.
646   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
647   if (lBin && rhsConst && lBin.getKind() == AffineExprKind::Mul) {
648     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>())
649       return lBin.getLHS() * (lrhs.getValue() * rhsConst.getValue());
650   }
651 
652   // When doing successive multiplication, bring constant to the right: turn (d0
653   // * 2) * d1 into (d0 * d1) * 2.
654   if (lBin && lBin.getKind() == AffineExprKind::Mul) {
655     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
656       return (lBin.getLHS() * rhs) * lrhs;
657     }
658   }
659 
660   return nullptr;
661 }
662 
663 AffineExpr AffineExpr::operator*(int64_t v) const {
664   return *this * getAffineConstantExpr(v, getContext());
665 }
666 AffineExpr AffineExpr::operator*(AffineExpr other) const {
667   if (auto simplified = simplifyMul(*this, other))
668     return simplified;
669 
670   StorageUniquer &uniquer = getContext()->getAffineUniquer();
671   return uniquer.get<AffineBinaryOpExprStorage>(
672       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
673 }
674 
675 // Unary minus, delegate to operator*.
676 AffineExpr AffineExpr::operator-() const {
677   return *this * getAffineConstantExpr(-1, getContext());
678 }
679 
680 // Delegate to operator+.
681 AffineExpr AffineExpr::operator-(int64_t v) const { return *this + (-v); }
682 AffineExpr AffineExpr::operator-(AffineExpr other) const {
683   return *this + (-other);
684 }
685 
686 static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) {
687   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
688   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
689 
690   // mlir floordiv by zero or negative numbers is undefined and preserved as is.
691   if (!rhsConst || rhsConst.getValue() < 1)
692     return nullptr;
693 
694   if (lhsConst)
695     return getAffineConstantExpr(
696         floorDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext());
697 
698   // Fold floordiv of a multiply with a constant that is a multiple of the
699   // divisor. Eg: (i * 128) floordiv 64 = i * 2.
700   if (rhsConst == 1)
701     return lhs;
702 
703   // Simplify (expr * const) floordiv divConst when expr is known to be a
704   // multiple of divConst.
705   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
706   if (lBin && lBin.getKind() == AffineExprKind::Mul) {
707     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
708       // rhsConst is known to be a positive constant.
709       if (lrhs.getValue() % rhsConst.getValue() == 0)
710         return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
711     }
712   }
713 
714   // Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is
715   // known to be a multiple of divConst.
716   if (lBin && lBin.getKind() == AffineExprKind::Add) {
717     int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
718     int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
719     // rhsConst is known to be a positive constant.
720     if (llhsDiv % rhsConst.getValue() == 0 ||
721         lrhsDiv % rhsConst.getValue() == 0)
722       return lBin.getLHS().floorDiv(rhsConst.getValue()) +
723              lBin.getRHS().floorDiv(rhsConst.getValue());
724   }
725 
726   return nullptr;
727 }
728 
729 AffineExpr AffineExpr::floorDiv(uint64_t v) const {
730   return floorDiv(getAffineConstantExpr(v, getContext()));
731 }
732 AffineExpr AffineExpr::floorDiv(AffineExpr other) const {
733   if (auto simplified = simplifyFloorDiv(*this, other))
734     return simplified;
735 
736   StorageUniquer &uniquer = getContext()->getAffineUniquer();
737   return uniquer.get<AffineBinaryOpExprStorage>(
738       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::FloorDiv), *this,
739       other);
740 }
741 
742 static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) {
743   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
744   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
745 
746   if (!rhsConst || rhsConst.getValue() < 1)
747     return nullptr;
748 
749   if (lhsConst)
750     return getAffineConstantExpr(
751         ceilDiv(lhsConst.getValue(), rhsConst.getValue()), lhs.getContext());
752 
753   // Fold ceildiv of a multiply with a constant that is a multiple of the
754   // divisor. Eg: (i * 128) ceildiv 64 = i * 2.
755   if (rhsConst.getValue() == 1)
756     return lhs;
757 
758   // Simplify (expr * const) ceildiv divConst when const is known to be a
759   // multiple of divConst.
760   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
761   if (lBin && lBin.getKind() == AffineExprKind::Mul) {
762     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
763       // rhsConst is known to be a positive constant.
764       if (lrhs.getValue() % rhsConst.getValue() == 0)
765         return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
766     }
767   }
768 
769   return nullptr;
770 }
771 
772 AffineExpr AffineExpr::ceilDiv(uint64_t v) const {
773   return ceilDiv(getAffineConstantExpr(v, getContext()));
774 }
775 AffineExpr AffineExpr::ceilDiv(AffineExpr other) const {
776   if (auto simplified = simplifyCeilDiv(*this, other))
777     return simplified;
778 
779   StorageUniquer &uniquer = getContext()->getAffineUniquer();
780   return uniquer.get<AffineBinaryOpExprStorage>(
781       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::CeilDiv), *this,
782       other);
783 }
784 
785 static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) {
786   auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
787   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
788 
789   // mod w.r.t zero or negative numbers is undefined and preserved as is.
790   if (!rhsConst || rhsConst.getValue() < 1)
791     return nullptr;
792 
793   if (lhsConst)
794     return getAffineConstantExpr(mod(lhsConst.getValue(), rhsConst.getValue()),
795                                  lhs.getContext());
796 
797   // Fold modulo of an expression that is known to be a multiple of a constant
798   // to zero if that constant is a multiple of the modulo factor. Eg: (i * 128)
799   // mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0.
800   if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0)
801     return getAffineConstantExpr(0, lhs.getContext());
802 
803   // Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is
804   // known to be a multiple of divConst.
805   auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
806   if (lBin && lBin.getKind() == AffineExprKind::Add) {
807     int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
808     int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
809     // rhsConst is known to be a positive constant.
810     if (llhsDiv % rhsConst.getValue() == 0)
811       return lBin.getRHS() % rhsConst.getValue();
812     if (lrhsDiv % rhsConst.getValue() == 0)
813       return lBin.getLHS() % rhsConst.getValue();
814   }
815 
816   return nullptr;
817 }
818 
819 AffineExpr AffineExpr::operator%(uint64_t v) const {
820   return *this % getAffineConstantExpr(v, getContext());
821 }
822 AffineExpr AffineExpr::operator%(AffineExpr other) const {
823   if (auto simplified = simplifyMod(*this, other))
824     return simplified;
825 
826   StorageUniquer &uniquer = getContext()->getAffineUniquer();
827   return uniquer.get<AffineBinaryOpExprStorage>(
828       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
829 }
830 
831 AffineExpr AffineExpr::compose(AffineMap map) const {
832   SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(),
833                                              map.getResults().end());
834   return replaceDimsAndSymbols(dimReplacements, {});
835 }
836 raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) {
837   expr.print(os);
838   return os;
839 }
840 
841 /// Constructs an affine expression from a flat ArrayRef. If there are local
842 /// identifiers (neither dimensional nor symbolic) that appear in the sum of
843 /// products expression, `localExprs` is expected to have the AffineExpr
844 /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be
845 /// in the format [dims, symbols, locals, constant term].
846 AffineExpr mlir::getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
847                                            unsigned numDims,
848                                            unsigned numSymbols,
849                                            ArrayRef<AffineExpr> localExprs,
850                                            MLIRContext *context) {
851   // Assert expected numLocals = flatExprs.size() - numDims - numSymbols - 1.
852   assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
853          "unexpected number of local expressions");
854 
855   auto expr = getAffineConstantExpr(0, context);
856   // Dimensions and symbols.
857   for (unsigned j = 0; j < numDims + numSymbols; j++) {
858     if (flatExprs[j] == 0)
859       continue;
860     auto id = j < numDims ? getAffineDimExpr(j, context)
861                           : getAffineSymbolExpr(j - numDims, context);
862     expr = expr + id * flatExprs[j];
863   }
864 
865   // Local identifiers.
866   for (unsigned j = numDims + numSymbols, e = flatExprs.size() - 1; j < e;
867        j++) {
868     if (flatExprs[j] == 0)
869       continue;
870     auto term = localExprs[j - numDims - numSymbols] * flatExprs[j];
871     expr = expr + term;
872   }
873 
874   // Constant term.
875   int64_t constTerm = flatExprs[flatExprs.size() - 1];
876   if (constTerm != 0)
877     expr = expr + constTerm;
878   return expr;
879 }
880 
881 SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims,
882                                                      unsigned numSymbols)
883     : numDims(numDims), numSymbols(numSymbols), numLocals(0) {
884   operandExprStack.reserve(8);
885 }
886 
887 void SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
888   assert(operandExprStack.size() >= 2);
889   // This is a pure affine expr; the RHS will be a constant.
890   assert(expr.getRHS().isa<AffineConstantExpr>());
891   // Get the RHS constant.
892   auto rhsConst = operandExprStack.back()[getConstantIndex()];
893   operandExprStack.pop_back();
894   // Update the LHS in place instead of pop and push.
895   auto &lhs = operandExprStack.back();
896   for (unsigned i = 0, e = lhs.size(); i < e; i++) {
897     lhs[i] *= rhsConst;
898   }
899 }
900 
901 void SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) {
902   assert(operandExprStack.size() >= 2);
903   const auto &rhs = operandExprStack.back();
904   auto &lhs = operandExprStack[operandExprStack.size() - 2];
905   assert(lhs.size() == rhs.size());
906   // Update the LHS in place.
907   for (unsigned i = 0, e = rhs.size(); i < e; i++) {
908     lhs[i] += rhs[i];
909   }
910   // Pop off the RHS.
911   operandExprStack.pop_back();
912 }
913 
914 //
915 // t = expr mod c   <=>  t = expr - c*q and c*q <= expr <= c*q + c - 1
916 //
917 // A mod expression "expr mod c" is thus flattened by introducing a new local
918 // variable q (= expr floordiv c), such that expr mod c is replaced with
919 // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
920 void SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
921   assert(operandExprStack.size() >= 2);
922   // This is a pure affine expr; the RHS will be a constant.
923   assert(expr.getRHS().isa<AffineConstantExpr>());
924   auto rhsConst = operandExprStack.back()[getConstantIndex()];
925   operandExprStack.pop_back();
926   auto &lhs = operandExprStack.back();
927   // TODO: handle modulo by zero case when this issue is fixed
928   // at the other places in the IR.
929   assert(rhsConst > 0 && "RHS constant has to be positive");
930 
931   // Check if the LHS expression is a multiple of modulo factor.
932   unsigned i, e;
933   for (i = 0, e = lhs.size(); i < e; i++)
934     if (lhs[i] % rhsConst != 0)
935       break;
936   // If yes, modulo expression here simplifies to zero.
937   if (i == lhs.size()) {
938     std::fill(lhs.begin(), lhs.end(), 0);
939     return;
940   }
941 
942   // Add a local variable for the quotient, i.e., expr % c is replaced by
943   // (expr - q * c) where q = expr floordiv c. Do this while canceling out
944   // the GCD of expr and c.
945   SmallVector<int64_t, 8> floorDividend(lhs);
946   uint64_t gcd = rhsConst;
947   for (unsigned i = 0, e = lhs.size(); i < e; i++)
948     gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(lhs[i]));
949   // Simplify the numerator and the denominator.
950   if (gcd != 1) {
951     for (unsigned i = 0, e = floorDividend.size(); i < e; i++)
952       floorDividend[i] = floorDividend[i] / static_cast<int64_t>(gcd);
953   }
954   int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd);
955 
956   // Construct the AffineExpr form of the floordiv to store in localExprs.
957   MLIRContext *context = expr.getContext();
958   auto dividendExpr = getAffineExprFromFlatForm(
959       floorDividend, numDims, numSymbols, localExprs, context);
960   auto divisorExpr = getAffineConstantExpr(floorDivisor, context);
961   auto floorDivExpr = dividendExpr.floorDiv(divisorExpr);
962   int loc;
963   if ((loc = findLocalId(floorDivExpr)) == -1) {
964     addLocalFloorDivId(floorDividend, floorDivisor, floorDivExpr);
965     // Set result at top of stack to "lhs - rhsConst * q".
966     lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst;
967   } else {
968     // Reuse the existing local id.
969     lhs[getLocalVarStartIndex() + loc] = -rhsConst;
970   }
971 }
972 
973 void SimpleAffineExprFlattener::visitCeilDivExpr(AffineBinaryOpExpr expr) {
974   visitDivExpr(expr, /*isCeil=*/true);
975 }
976 void SimpleAffineExprFlattener::visitFloorDivExpr(AffineBinaryOpExpr expr) {
977   visitDivExpr(expr, /*isCeil=*/false);
978 }
979 
980 void SimpleAffineExprFlattener::visitDimExpr(AffineDimExpr expr) {
981   operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
982   auto &eq = operandExprStack.back();
983   assert(expr.getPosition() < numDims && "Inconsistent number of dims");
984   eq[getDimStartIndex() + expr.getPosition()] = 1;
985 }
986 
987 void SimpleAffineExprFlattener::visitSymbolExpr(AffineSymbolExpr expr) {
988   operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
989   auto &eq = operandExprStack.back();
990   assert(expr.getPosition() < numSymbols && "inconsistent number of symbols");
991   eq[getSymbolStartIndex() + expr.getPosition()] = 1;
992 }
993 
994 void SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) {
995   operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
996   auto &eq = operandExprStack.back();
997   eq[getConstantIndex()] = expr.getValue();
998 }
999 
1000 // t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
1001 // A floordiv is thus flattened by introducing a new local variable q, and
1002 // replacing that expression with 'q' while adding the constraints
1003 // c * q <= expr <= c * q + c - 1 to localVarCst (done by
1004 // FlatAffineConstraints::addLocalFloorDiv).
1005 //
1006 // A ceildiv is similarly flattened:
1007 // t = expr ceildiv c   <=> t =  (expr + c - 1) floordiv c
1008 void SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
1009                                              bool isCeil) {
1010   assert(operandExprStack.size() >= 2);
1011   assert(expr.getRHS().isa<AffineConstantExpr>());
1012 
1013   // This is a pure affine expr; the RHS is a positive constant.
1014   int64_t rhsConst = operandExprStack.back()[getConstantIndex()];
1015   // TODO: handle division by zero at the same time the issue is
1016   // fixed at other places.
1017   assert(rhsConst > 0 && "RHS constant has to be positive");
1018   operandExprStack.pop_back();
1019   auto &lhs = operandExprStack.back();
1020 
1021   // Simplify the floordiv, ceildiv if possible by canceling out the greatest
1022   // common divisors of the numerator and denominator.
1023   uint64_t gcd = std::abs(rhsConst);
1024   for (unsigned i = 0, e = lhs.size(); i < e; i++)
1025     gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(lhs[i]));
1026   // Simplify the numerator and the denominator.
1027   if (gcd != 1) {
1028     for (unsigned i = 0, e = lhs.size(); i < e; i++)
1029       lhs[i] = lhs[i] / static_cast<int64_t>(gcd);
1030   }
1031   int64_t divisor = rhsConst / static_cast<int64_t>(gcd);
1032   // If the divisor becomes 1, the updated LHS is the result. (The
1033   // divisor can't be negative since rhsConst is positive).
1034   if (divisor == 1)
1035     return;
1036 
1037   // If the divisor cannot be simplified to one, we will have to retain
1038   // the ceil/floor expr (simplified up until here). Add an existential
1039   // quantifier to express its result, i.e., expr1 div expr2 is replaced
1040   // by a new identifier, q.
1041   MLIRContext *context = expr.getContext();
1042   auto a =
1043       getAffineExprFromFlatForm(lhs, numDims, numSymbols, localExprs, context);
1044   auto b = getAffineConstantExpr(divisor, context);
1045 
1046   int loc;
1047   auto divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
1048   if ((loc = findLocalId(divExpr)) == -1) {
1049     if (!isCeil) {
1050       SmallVector<int64_t, 8> dividend(lhs);
1051       addLocalFloorDivId(dividend, divisor, divExpr);
1052     } else {
1053       // lhs ceildiv c <=>  (lhs + c - 1) floordiv c
1054       SmallVector<int64_t, 8> dividend(lhs);
1055       dividend.back() += divisor - 1;
1056       addLocalFloorDivId(dividend, divisor, divExpr);
1057     }
1058   }
1059   // Set the expression on stack to the local var introduced to capture the
1060   // result of the division (floor or ceil).
1061   std::fill(lhs.begin(), lhs.end(), 0);
1062   if (loc == -1)
1063     lhs[getLocalVarStartIndex() + numLocals - 1] = 1;
1064   else
1065     lhs[getLocalVarStartIndex() + loc] = 1;
1066 }
1067 
1068 // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
1069 // The local identifier added is always a floordiv of a pure add/mul affine
1070 // function of other identifiers, coefficients of which are specified in
1071 // dividend and with respect to a positive constant divisor. localExpr is the
1072 // simplified tree expression (AffineExpr) corresponding to the quantifier.
1073 void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
1074                                                    int64_t divisor,
1075                                                    AffineExpr localExpr) {
1076   assert(divisor > 0 && "positive constant divisor expected");
1077   for (auto &subExpr : operandExprStack)
1078     subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
1079   localExprs.push_back(localExpr);
1080   numLocals++;
1081   // dividend and divisor are not used here; an override of this method uses it.
1082 }
1083 
1084 int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
1085   SmallVectorImpl<AffineExpr>::iterator it;
1086   if ((it = llvm::find(localExprs, localExpr)) == localExprs.end())
1087     return -1;
1088   return it - localExprs.begin();
1089 }
1090 
1091 /// Simplify the affine expression by flattening it and reconstructing it.
1092 AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
1093                                     unsigned numSymbols) {
1094   // Simplify semi-affine expressions separately.
1095   if (!expr.isPureAffine())
1096     expr = simplifySemiAffine(expr);
1097   if (!expr.isPureAffine())
1098     return expr;
1099 
1100   SimpleAffineExprFlattener flattener(numDims, numSymbols);
1101   flattener.walkPostOrder(expr);
1102   ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
1103   auto simplifiedExpr =
1104       getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
1105                                 flattener.localExprs, expr.getContext());
1106   flattener.operandExprStack.pop_back();
1107   assert(flattener.operandExprStack.empty());
1108 
1109   return simplifiedExpr;
1110 }
1111