[mlir][analysis][NFC] Move functions from FlatAffineValueConstraints to IntegerPolyhedron

[springerm]

Summary

Move addBound overloads with AffineMaps, getSliceBounds and getLowerAndUpperBound from FlatAffineValueConstraints to IntegerPolyhedron. This functionality can now be used without depending on the Affine dialect.

This is in preparation of D145681, which adds ValueBoundsOpInterface to mlir/Interfaces.

Depends On: D145935

[Groverkss]

This introduces a dependency of the Presburger library on IR constructs like AffineMap, and IntegerSet, which was discussed before as a dependency we do not want to put. This is the reason FlatAffineConstraints contains these functions, and they were not moved.

[Groverkss]

The Presburger equivalent of an AffineMap is presburger::MultiAffineFunction (aka MAF). Ideally, a conversion between AffineMap to MAF should be added, and these methods should use MAF instead.

[springerm]

In D146029#4192398, @Groverkss wrote:

This introduces a dependency of the Presburger library on IR constructs like AffineMap, and IntegerSet, which was discussed before as a dependency we do not want to put. This is the reason FlatAffineConstraints contains these functions, and they were not moved.

How about adding another class in the hierarchy:

IntegerRelation
   ^
   |
IntegerPolyhedron
   ^
   |
FlatAffineConstraints (or another name)
   ^
   |
FlatAffineValueConstraints

FlatAffineConstraints would be in in mlir/Analysis. I can be in a build target/directory different from Presburger.

Context: We have a use case for FlatAffineValueConstraints in an op interface (mlir/Interfaces). Build targets in mlir/Interfaces cannot depend on any dialect, so we cannot use FlatAffineValueConstraints.
Discussion on Discourse: https://discourse.llvm.org/t/rfc-valueboundsopinterface-compute-bounds-of-index-values-dyn-dim-sizes/69174

[springerm]

In D146029#4192402, @Groverkss wrote:

The Presburger equivalent of an AffineMap is presburger::MultiAffineFunction (aka MAF). Ideally, a conversion between AffineMap to MAF should be added, and these methods should use MAF instead.

If AffineMap and MultiAffineFunction are equivalent, why do we have MultiAffineFunction? (AffineExpr and AffineMap are in the IR build target, so they can be used without depending on any dialect.)

[springerm]

In D146029#4192429, @springerm wrote:

In D146029#4192398, @Groverkss wrote:

This introduces a dependency of the Presburger library on IR constructs like AffineMap, and IntegerSet, which was discussed before as a dependency we do not want to put. This is the reason FlatAffineConstraints contains these functions, and they were not moved.

How about adding another class in the hierarchy:

Basically, I'm trying to see if we can extract a subset of FlatAffineValueConstraints that does not depend on any affine ops and put it in mlir/Analysis. A better name would be FlatValueConstraints. It would not have functions such as addAffineIfOpDomain etc. FlatAffineValueConstraints (as a subclass) would add provide those functions. Would this work?

[Groverkss]

In D146029#4192443, @springerm wrote:

In D146029#4192402, @Groverkss wrote:

The Presburger equivalent of an AffineMap is presburger::MultiAffineFunction (aka MAF). Ideally, a conversion between AffineMap to MAF should be added, and these methods should use MAF instead.

If AffineMap and MultiAffineFunction are equivalent, why do we have MultiAffineFunction? (AffineExpr and AffineMap are in the IR build target, so they can be used without depending on any dialect.)

Sorry for the confusion, when I said they are "equivalent", I meant they represent the same mathematical function. Their storage is very different. MAF stores a flattened version of the function, which is useful for analysis, while AffineMap, and AffineExpr store an AST representation of the function.

[Groverkss]

In D146029#4192429, @springerm wrote:

In D146029#4192398, @Groverkss wrote:

This introduces a dependency of the Presburger library on IR constructs like AffineMap, and IntegerSet, which was discussed before as a dependency we do not want to put. This is the reason FlatAffineConstraints contains these functions, and they were not moved.

How about adding another class in the hierarchy:

IntegerRelation
   ^
   |
IntegerPolyhedron
   ^
   |
FlatAffineConstraints (or another name)
   ^
   |
FlatAffineValueConstraints

FlatAffineConstraints would be in in mlir/Analysis. I can be in a build target/directory different from Presburger.

Context: We have a use case for FlatAffineValueConstraints in an op interface (mlir/Interfaces). Build targets in mlir/Interfaces cannot depend on any dialect, so we cannot use FlatAffineValueConstraints.
Discussion on Discourse: https://discourse.llvm.org/t/rfc-valueboundsopinterface-compute-bounds-of-index-values-dyn-dim-sizes/69174

I would personally prefer the conversion to MAF and using MAF for adding bounds solution and would also be up for helping by sending patches for it.

For the changing hierarchy solution, I think it is fine, but I would wait to see what @ftynse and @bondhugula think.

[springerm]

Splitting FlatAffineValueConstraints: D146201

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 #include <optional>
 
 namespace mlir {
+class AffineExpr;
+class AffineMap;
+class IntegerSet;
+class MLIRContext;
+
 namespace presburger {
 
 class IntegerRelation;
 class IntegerPolyhedron;
 class PresburgerSet;
 class PresburgerRelation;
 struct SymbolicLexMin;
 
public:
 
   /// Return the kind of this IntegerRelation.
   Kind getKind() const override { return Kind::IntegerPolyhedron; }
 
   static bool classof(const IntegerRelation *cst) {
     return cst->getKind() == Kind::IntegerPolyhedron;
   }
 
+  /// Adds a bound for the variable at the specified position with constraints
+  /// being drawn from the specified bound map. In case of an EQ bound, the
+  /// bound map is expected to have exactly one result. In case of a LB/UB, the
+  /// bound map may have more than one result, for each of which an inequality
+  /// is added.
+  ///
+  /// The bound can be added as open or closed by specifying isClosedBound. In
+  /// case of a LB/UB, isClosedBound = false means the bound is added internally
+  /// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can
+  /// only be added as a closed bound.
+  ///
+  /// Note: The dimensions/symbols of this IntegerRelation must match the
+  /// dimensions/symbols of the affine map.
+  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
+                         bool isClosedBound);
+
+  /// Adds a bound for the variable at the specified position with constraints
+  /// being drawn from the specified bound map. In case of an EQ bound, the
+  /// bound map is expected to have exactly one result. In case of a LB/UB, the
+  /// bound map may have more than one result, for each of which an inequality
+  /// is added.
+  ///
+  /// Note: The dimensions/symbols of this IntegerRelation must match the
+  /// dimensions/symbols of the affine map. A lower bound is assumed to be
+  /// closed and an upper bound is assumed to be open.
+  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap);
+
+  /// The `addBound` overload above hides the inherited overloads by default, so
+  /// we explicitly introduce them here.
+  using IntegerRelation::addBound;
+
   // Clones this object.
   std::unique_ptr<IntegerPolyhedron> clone() const;
 
+  /// Check if the pos^th variable can be expressed as a floordiv of an affine
+  /// function of other variables (where the divisor is a positive constant)
+  /// given the initial set of expressions in `exprs`. If it can be, the
+  /// corresponding position in `exprs` is set as the detected affine expr. For
+  /// eg: 4q <= i + j <= 4q + 3   <=>   q = (i + j) floordiv 4. An equality can
+  /// also yield a floordiv: eg.  4q = i + j <=> q = (i + j) floordiv 4. 32q +
+  /// 28
+  /// <= i <= 32q + 31 => q = i floordiv 32.
+  bool detectAsFloorDiv(unsigned pos, MLIRContext *context,
+                        SmallVectorImpl<AffineExpr> &exprs) const;
+
+  /// Determine whether the variable at 'pos' (say var_r) can be expressed as
+  /// modulo of another known variable (say var_n) w.r.t a constant. For
+  /// example, if the following constraints hold true:
+  /// ```
+  /// 0 <= var_r <= divisor - 1
+  /// var_n - (divisor * q_expr) = var_r
+  /// ```
+  /// where `var_n` is a known variable (called dividend), and `q_expr` is an
+  /// `AffineExpr` (called the quotient expression), `var_r` can be written as:
+  ///
+  /// `var_r = var_n mod divisor`.
+  ///
+  /// Additionally, in a special case of the above constaints where `q_expr` is
+  /// an variable itself that is not yet known (say `var_q`), it can be written
+  /// as a floordiv in the following way:
+  ///
+  /// `var_q = var_n floordiv divisor`.
+  ///
+  /// Returns true if the above mod or floordiv are detected, updating 'memo'
+  /// with these new expressions. Returns false otherwise.
+  bool detectAsMod(unsigned pos, int64_t lbConst, int64_t ubConst,
+                   SmallVectorImpl<AffineExpr> &memo,
+                   MLIRContext *context) const;
+
+  /// Gets the lower and upper bound of the `offset` + `pos`th variable
+  /// treating [0, offset) U [offset + num, symStartPos) as dimensions and
+  /// [symStartPos, getNumDimAndSymbolVars) as symbols, and `pos` lies in
+  /// [0, num). The multi-dimensional maps in the returned pair represent the
+  /// max and min of potentially multiple affine expressions. `localExprs` holds
+  /// pre-computed AffineExpr's for all local variables in the system.
+  ///
+  /// By default the returned lower bounds are closed and upper bounds are open.
+  /// If `closedUb` is true, the upper bound is closed.
+  std::pair<AffineMap, AffineMap>
+  getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num,
+                        unsigned symStartPos, ArrayRef<AffineExpr> localExprs,
+                        MLIRContext *context, bool closedUB = false) const;
+
+  /// Computes the lower and upper bounds of the first `num` dimensional
+  /// variables (starting at `offset`) as an affine map of the remaining
+  /// variables (dimensional and symbolic). This method is able to detect
+  /// variables as floordiv's and mod's of affine expressions of other
+  /// variables with respect to (positive) constants. Sets bound map to a
+  /// null AffineMap if such a bound can't be found (or yet unimplemented).
+  ///
+  /// By default the returned lower bounds are closed and upper bounds are open.
+  /// If `closedUb` is true, the upper bound is closed.
+  void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
+                      SmallVectorImpl<AffineMap> *lbMaps,
+                      SmallVectorImpl<AffineMap> *ubMaps,
+                      bool closedUB = false);
+
   /// Insert `num` variables of the specified kind at position `pos`.
   /// Positions are relative to the kind of variable. Return the absolute
   /// column position (i.e., not relative to the kind of variable) of the
   /// first added variable.
   unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Return the intersection of the two relations.
   /// If there are locals, they will be merged.
   IntegerPolyhedron intersect(const IntegerPolyhedron &other) const;
 
   /// Return the set difference of this set and the given set, i.e.,
   /// return `this \ set`.
   PresburgerSet subtract(const PresburgerSet &other) const;
+
+protected:
+  /// Given an affine map that is aligned with this constraint system:
+  /// * Flatten the map.
+  /// * Add newly introduced local columns at the beginning of this constraint
+  ///   system (local column pos 0).
+  /// * Add equalities that define the new local columns to this constraint
+  ///   system.
+  /// * Return the flattened expressions via `flattenedExprs`.
+  ///
+  /// Note: This is a shared helper function of `addLowerOrUpperBound` and
+  ///       `composeMatchingMap`.
+  LogicalResult flattenAlignedMapAndMergeLocals(
+      AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
 };
 
 } // namespace presburger
+
+/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
+/// dimensions, symbols, and additional variables that represent floor divisions
+/// of dimensions, symbols, and in turn other floor divisions.  Returns failure
+/// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
+/// 'cst' contains constraints that connect newly introduced local variables
+/// to existing dimensional and symbolic variables. See documentation for
+/// AffineExprFlattener on how mod's and div's are flattened.
+LogicalResult
+getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
+                       SmallVectorImpl<int64_t> *flattenedExpr,
+                       presburger::IntegerPolyhedron *cst = nullptr);
+
+/// Flattens the result expressions of the map to their corresponding flattened
+/// forms and set in 'flattenedExprs'. Returns failure if any expression in the
+/// map could not be flattened (i.e., semi-affine is not yet handled). 'cst'
+/// contains constraints that connect newly introduced local variables to
+/// existing dimensional and / symbolic variables. See documentation for
+/// AffineExprFlattener on how mod's and div's are flattened. For all affine
+/// expressions that share the same operands (like those of an affine map), this
+/// method should be used instead of repeatedly calling getFlattenedAffineExpr
+/// since local variables added to deal with div's and mod's will be reused
+/// across expressions.
+LogicalResult
+getFlattenedAffineExprs(AffineMap map,
+                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+                        presburger::IntegerPolyhedron *cst = nullptr);
+LogicalResult
+getFlattenedAffineExprs(IntegerSet set,
+                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+                        presburger::IntegerPolyhedron *cst = nullptr);
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERRELATION_H

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

public:
   /// instance argument (`ifOp`). It is asserted that:
   /// 1) The IntegerSet of the given `affine.if` instance should not contain
   /// semi-affine expressions,
   /// 2) The columns of the constraint system created from `ifOp` should match
   /// the columns in the current one regarding numbers and values.
   void addAffineIfOpDomain(AffineIfOp ifOp);
 
   /// Adds a bound for the variable at the specified position with constraints
-  /// being drawn from the specified bound map. In case of an EQ bound, the
-  /// bound map is expected to have exactly one result. In case of a LB/UB, the
-  /// bound map may have more than one result, for each of which an inequality
-  /// is added.
-  ///
-  /// The bound can be added as open or closed by specifying isClosedBound. In
-  /// case of a LB/UB, isClosedBound = false means the bound is added internally
-  /// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can
-  /// only be added as a closed bound.
-  ///
-  /// Note: The dimensions/symbols of this FlatAffineConstraints must match the
-  /// dimensions/symbols of the affine map.
-  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
-                         bool isClosedBound);
-
-  /// Adds a bound for the variable at the specified position with constraints
-  /// being drawn from the specified bound map. In case of an EQ bound, the
-  /// bound map is expected to have exactly one result. In case of a LB/UB, the
-  /// bound map may have more than one result, for each of which an inequality
-  /// is added.
-  /// Note: The dimensions/symbols of this FlatAffineConstraints must match the
-  /// dimensions/symbols of the affine map. By default the lower bound is closed
-  /// and the upper bound is open.
-  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap);
-
-  /// Adds a bound for the variable at the specified position with constraints
   /// being drawn from the specified bound map and operands. In case of an
   /// EQ bound, the  bound map is expected to have exactly one result. In case
   /// of a LB/UB, the bound map may have more than one result, for each of which
   /// an inequality is added.
   LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
                          ValueRange operands);
 
   /// Adds a constant bound for the variable associated with the given Value.
   void addBound(BoundType type, Value val, int64_t value);
 
   /// The `addBound` overload above hides the inherited overloads by default, so
   /// we explicitly introduce them here.
   using IntegerPolyhedron::addBound;
 
   /// Returns the constraint system as an integer set. Returns a null integer
   /// set if the system has no constraints, or if an integer set couldn't be
   /// constructed as a result of a local variable's explicit representation not
   /// being known and such a local variable appearing in any of the constraints.
   IntegerSet getAsIntegerSet(MLIRContext *context) const;
 
-  /// Computes the lower and upper bounds of the first `num` dimensional
-  /// variables (starting at `offset`) as an affine map of the remaining
-  /// variables (dimensional and symbolic). This method is able to detect
-  /// variables as floordiv's and mod's of affine expressions of other
-  /// variables with respect to (positive) constants. Sets bound map to a
-  /// null AffineMap if such a bound can't be found (or yet unimplemented).
-  ///
-  /// By default the returned lower bounds are closed and upper bounds are open.
-  /// If `closedUb` is true, the upper bound is closed.
-  void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
-                      SmallVectorImpl<AffineMap> *lbMaps,
-                      SmallVectorImpl<AffineMap> *ubMaps,
-                      bool closedUB = false);
-
   /// Composes an affine map whose dimensions and symbols match one to one with
   /// the dimensions and symbols of this FlatAffineConstraints. The results of
   /// the map `other` are added as the leading dimensions of this constraint
   /// system. Returns failure if `other` is a semi-affine map.
   LogicalResult composeMatchingMap(AffineMap other);
 
-  /// Gets the lower and upper bound of the `offset` + `pos`th variable
-  /// treating [0, offset) U [offset + num, symStartPos) as dimensions and
-  /// [symStartPos, getNumDimAndSymbolVars) as symbols, and `pos` lies in
-  /// [0, num). The multi-dimensional maps in the returned pair represent the
-  /// max and min of potentially multiple affine expressions. `localExprs` holds
-  /// pre-computed AffineExpr's for all local variables in the system.
-  ///
-  /// By default the returned lower bounds are closed and upper bounds are open.
-  /// If `closedUb` is true, the upper bound is closed.
-  std::pair<AffineMap, AffineMap>
-  getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num,
-                        unsigned symStartPos, ArrayRef<AffineExpr> localExprs,
-                        MLIRContext *context, bool closedUB = false) const;
-
   /// Returns the bound for the variable at `pos` from the inequality at
   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
   /// affine value map can either be a lower bound or an upper bound depending
   /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does
   /// not involve the `pos`th variable.
   void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos,
                                AffineValueMap &vmap,
                                MLIRContext *context) const;

 protected:
   using VarKind = presburger::VarKind;
 
   /// Returns false if the fields corresponding to various variable counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
   bool hasConsistentState() const override;
 
-  /// Given an affine map that is aligned with this constraint system:
-  /// * Flatten the map.
-  /// * Add newly introduced local columns at the beginning of this constraint
-  ///   system (local column pos 0).
-  /// * Add equalities that define the new local columns to this constraint
-  ///   system.
-  /// * Return the flattened expressions via `flattenedExprs`.
-  ///
-  /// Note: This is a shared helper function of `addLowerOrUpperBound` and
-  ///       `composeMatchingMap`.
-  LogicalResult flattenAlignedMapAndMergeLocals(
-      AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
-
   /// Eliminates the variable at the specified position using Fourier-Motzkin
   /// variable elimination, but uses Gaussian elimination if there is an
   /// equality involving that variable. If the result of the elimination is
   /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is
   /// set to true, a potential under approximation (subset) of the rational
   /// shadow / exact integer shadow is computed.
   // See implementation comments for more details.
   void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,

 protected:
   // Number of dimension variables corresponding to domain variables.
   unsigned numDomainDims;
 
   // Number of dimension variables corresponding to range variables.
   unsigned numRangeDims;
 };
 
-/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
-/// dimensions, symbols, and additional variables that represent floor divisions
-/// of dimensions, symbols, and in turn other floor divisions.  Returns failure
-/// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
-/// 'cst' contains constraints that connect newly introduced local variables
-/// to existing dimensional and symbolic variables. See documentation for
-/// AffineExprFlattener on how mod's and div's are flattened.
-LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
-                                     unsigned numSymbols,
-                                     SmallVectorImpl<int64_t> *flattenedExpr,
-                                     FlatAffineValueConstraints *cst = nullptr);
-
-/// Flattens the result expressions of the map to their corresponding flattened
-/// forms and set in 'flattenedExprs'. Returns failure if any expression in the
-/// map could not be flattened (i.e., semi-affine is not yet handled). 'cst'
-/// contains constraints that connect newly introduced local variables to
-/// existing dimensional and / symbolic variables. See documentation for
-/// AffineExprFlattener on how mod's and div's are flattened. For all affine
-/// expressions that share the same operands (like those of an affine map), this
-/// method should be used instead of repeatedly calling getFlattenedAffineExpr
-/// since local variables added to deal with div's and mod's will be reused
-/// across expressions.
-LogicalResult
-getFlattenedAffineExprs(AffineMap map,
-                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-                        FlatAffineValueConstraints *cst = nullptr);
-LogicalResult
-getFlattenedAffineExprs(IntegerSet set,
-                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-                        FlatAffineValueConstraints *cst = nullptr);
-
 LogicalResult
 getMultiAffineFunctionFromMap(AffineMap map,
                               presburger::MultiAffineFunction &multiAff);
 
 /// Re-indexes the dimensions and symbols of an affine map with given `operands`
 /// values to align with `dims` and `syms` values.
 ///
 /// Each dimension/symbol of the map, bound to an operand `o`, is replaced with

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/LinearTransform.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
+#include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/IntegerSet.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Debug.h"
 #include <numeric>
 #include <optional>
 
 #define DEBUG_TYPE "presburger"
 
void IntegerRelation::swapVar(unsigned posA, unsigned posB) {
   equalities.swapColumns(posA, posB);
 }
 
 void IntegerRelation::clearConstraints() {
   equalities.resizeVertically(0);
   inequalities.resizeVertically(0);
 }
 
+bool IntegerPolyhedron::detectAsMod(unsigned pos, int64_t lbConst,
+                                    int64_t ubConst,
+                                    SmallVectorImpl<AffineExpr> &memo,
+                                    MLIRContext *context) const {
+  assert(pos < getNumVars() && "invalid position");
+
+  // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can
+  // be determined.
+  if (lbConst != 0 || ubConst < 1)
+    return false;
+  int64_t divisor = ubConst + 1;
+
+  // Check for the aforementioned conditions in each equality.
+  for (unsigned curEquality = 0, numEqualities = getNumEqualities();
+       curEquality < numEqualities; curEquality++) {
+    int64_t coefficientAtPos = atEq64(curEquality, pos);
+    // If current equality does not involve `var_r`, continue to the next
+    // equality.
+    if (coefficientAtPos == 0)
+      continue;
+
+    // Constant term should be 0 in this equality.
+    if (atEq64(curEquality, getNumCols() - 1) != 0)
+      continue;
+
+    // Traverse through the equality and construct the dividend expression
+    // `dividendExpr`, to contain all the variables which are known and are
+    // not divisible by `(coefficientAtPos * divisor)`. Hope here is that the
+    // `dividendExpr` gets simplified into a single variable `var_n` discussed
+    // above.
+    auto dividendExpr = getAffineConstantExpr(0, context);
+
+    // Track the terms that go into quotient expression, later used to detect
+    // additional floordiv.
+    unsigned quotientCount = 0;
+    int quotientPosition = -1;
+    int quotientSign = 1;
+
+    // Consider each term in the current equality.
+    unsigned curVar, e;
+    for (curVar = 0, e = getNumDimAndSymbolVars(); curVar < e; ++curVar) {
+      // Ignore var_r.
+      if (curVar == pos)
+        continue;
+      int64_t coefficientOfCurVar = atEq64(curEquality, curVar);
+      // Ignore vars that do not contribute to the current equality.
+      if (coefficientOfCurVar == 0)
+        continue;
+      // Check if the current var goes into the quotient expression.
+      if (coefficientOfCurVar % (divisor * coefficientAtPos) == 0) {
+        quotientCount++;
+        quotientPosition = curVar;
+        quotientSign = (coefficientOfCurVar * coefficientAtPos) > 0 ? 1 : -1;
+        continue;
+      }
+      // Variables that are part of dividendExpr should be known.
+      if (!memo[curVar])
+        break;
+      // Append the current variable to the dividend expression.
+      dividendExpr = dividendExpr + memo[curVar] * coefficientOfCurVar;
+    }
+
+    // Can't construct expression as it depends on a yet uncomputed var.
+    if (curVar < e)
+      continue;
+
+    // Express `var_r` in terms of the other vars collected so far.
+    if (coefficientAtPos > 0)
+      dividendExpr = (-dividendExpr).floorDiv(coefficientAtPos);
+    else
+      dividendExpr = dividendExpr.floorDiv(-coefficientAtPos);
+
+    // Simplify the expression.
+    dividendExpr =
+        simplifyAffineExpr(dividendExpr, getNumDimVars(), getNumSymbolVars());
+    // Only if the final dividend expression is just a single var (which we call
+    // `var_n`), we can proceed.
+    // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it
+    // to dims themselves.
+    auto dimExpr = dividendExpr.dyn_cast<AffineDimExpr>();
+    if (!dimExpr)
+      continue;
+
+    // Express `var_r` as `var_n % divisor` and store the expression in `memo`.
+    if (quotientCount >= 1) {
+      auto ub = getConstantBound64(BoundType::UB, dimExpr.getPosition());
+      // If `var_n` has an upperbound that is less than the divisor, mod can be
+      // eliminated altogether.
+      if (ub && *ub < divisor)
+        memo[pos] = dimExpr;
+      else
+        memo[pos] = dimExpr % divisor;
+      // If a unique quotient `var_q` was seen, it can be expressed as
+      // `var_n floordiv divisor`.
+      if (quotientCount == 1 && !memo[quotientPosition])
+        memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign;
+
+      return true;
+    }
+  }
+  return false;
+}
+
+bool IntegerPolyhedron::detectAsFloorDiv(
+    unsigned pos, MLIRContext *context,
+    SmallVectorImpl<AffineExpr> &exprs) const {
+  assert(pos < getNumVars() && "invalid position");
+
+  // Get upper-lower bound pair for this variable.
+  SmallVector<bool, 8> foundRepr(getNumVars(), false);
+  for (unsigned i = 0, e = getNumVars(); i < e; ++i)
+    if (exprs[i])
+      foundRepr[i] = true;
+
+  SmallVector<int64_t, 8> dividend(getNumCols());
+  unsigned divisor;
+  auto ulPair = computeSingleVarRepr(*this, foundRepr, pos, dividend, divisor);
+
+  // No upper-lower bound pair found for this var.
+  if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality)
+    return false;
+
+  // Construct the dividend expression.
+  auto dividendExpr = getAffineConstantExpr(dividend.back(), context);
+  for (unsigned c = 0, f = getNumVars(); c < f; c++)
+    if (dividend[c] != 0)
+      dividendExpr = dividendExpr + dividend[c] * exprs[c];
+
+  // Successfully detected the floordiv.
+  exprs[pos] = dividendExpr.floorDiv(divisor);
+  return true;
+}
+
+std::pair<AffineMap, AffineMap> IntegerPolyhedron::getLowerAndUpperBound(
+    unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
+    ArrayRef<AffineExpr> localExprs, MLIRContext *context,
+    bool closedUB) const {
+  assert(pos + offset < getNumDimVars() && "invalid dim start pos");
+  assert(symStartPos >= (pos + offset) && "invalid sym start pos");
+  assert(getNumLocalVars() == localExprs.size() &&
+         "incorrect local exprs count");
+
+  SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
+  getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices,
+                               offset, num);
+
+  /// Add to 'b' from 'a' in set [0, offset) U [offset + num, symbStartPos).
+  auto addCoeffs = [&](ArrayRef<int64_t> a, SmallVectorImpl<int64_t> &b) {
+    b.clear();
+    for (unsigned i = 0, e = a.size(); i < e; ++i) {
+      if (i < offset || i >= offset + num)
+        b.push_back(a[i]);
+    }
+  };
+
+  SmallVector<int64_t, 8> lb, ub;
+  SmallVector<AffineExpr, 4> lbExprs;
+  unsigned dimCount = symStartPos - num;
+  unsigned symCount = getNumDimAndSymbolVars() - symStartPos;
+  lbExprs.reserve(lbIndices.size() + eqIndices.size());
+  // Lower bound expressions.
+  for (auto idx : lbIndices) {
+    auto ineq = getInequality64(idx);
+    // Extract the lower bound (in terms of other coeff's + const), i.e., if
+    // i - j + 1 >= 0 is the constraint, 'pos' is for i the lower bound is j
+    // - 1.
+    addCoeffs(ineq, lb);
+    std::transform(lb.begin(), lb.end(), lb.begin(), std::negate<int64_t>());
+    auto expr =
+        getAffineExprFromFlatForm(lb, dimCount, symCount, localExprs, context);
+    // expr ceildiv divisor is (expr + divisor - 1) floordiv divisor
+    int64_t divisor = std::abs(ineq[pos + offset]);
+    expr = (expr + divisor - 1).floorDiv(divisor);
+    lbExprs.push_back(expr);
+  }
+
+  SmallVector<AffineExpr, 4> ubExprs;
+  ubExprs.reserve(ubIndices.size() + eqIndices.size());
+  // Upper bound expressions.
+  for (auto idx : ubIndices) {
+    auto ineq = getInequality64(idx);
+    // Extract the upper bound (in terms of other coeff's + const).
+    addCoeffs(ineq, ub);
+    auto expr =
+        getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context);
+    expr = expr.floorDiv(std::abs(ineq[pos + offset]));
+    int64_t ubAdjustment = closedUB ? 0 : 1;
+    ubExprs.push_back(expr + ubAdjustment);
+  }
+
+  // Equalities. It's both a lower and a upper bound.
+  SmallVector<int64_t, 4> b;
+  for (auto idx : eqIndices) {
+    auto eq = getEquality64(idx);
+    addCoeffs(eq, b);
+    if (eq[pos + offset] > 0)
+      std::transform(b.begin(), b.end(), b.begin(), std::negate<int64_t>());
+
+    // Extract the upper bound (in terms of other coeff's + const).
+    auto expr =
+        getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
+    expr = expr.floorDiv(std::abs(eq[pos + offset]));
+    // Upper bound is exclusive.
+    ubExprs.push_back(expr + 1);
+    // Lower bound.
+    expr =
+        getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
+    expr = expr.ceilDiv(std::abs(eq[pos + offset]));
+    lbExprs.push_back(expr);
+  }
+
+  auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context);
+  auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context);
+
+  return {lbMap, ubMap};
+}
+
 /// Gather all lower and upper bounds of the variable at `pos`, and
 /// optionally any equalities on it. In addition, the bounds are to be
 /// independent of variables in position range [`offset`, `offset` + `num`).
 void IntegerRelation::getLowerAndUpperBoundIndices(
     unsigned pos, SmallVectorImpl<unsigned> *lbIndices,
     SmallVectorImpl<unsigned> *ubIndices, SmallVectorImpl<unsigned> *eqIndices,
     unsigned offset, unsigned num) const {
   assert(pos < getNumVars() && "invalid position");
for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     }
     os << ">= 0\n";
   }
   os << '\n';
 }
 
 void IntegerRelation::dump() const { print(llvm::errs()); }
 
+LogicalResult IntegerPolyhedron::addBound(BoundType type, unsigned pos,
+                                          AffineMap boundMap,
+                                          bool isClosedBound) {
+  assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch");
+  assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
+  assert(pos < getNumDimAndSymbolVars() && "invalid position");
+  assert((type != BoundType::EQ || isClosedBound) &&
+         "EQ bound must be closed.");
+
+  // Equality follows the logic of lower bound except that we add an equality
+  // instead of an inequality.
+  assert((type != BoundType::EQ || boundMap.getNumResults() == 1) &&
+         "single result expected");
+  bool lower = type == BoundType::LB || type == BoundType::EQ;
+
+  std::vector<SmallVector<int64_t, 8>> flatExprs;
+  if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
+    return failure();
+  assert(flatExprs.size() == boundMap.getNumResults());
+
+  // Add one (in)equality for each result.
+  for (const auto &flatExpr : flatExprs) {
+    SmallVector<int64_t> ineq(getNumCols(), 0);
+    // Dims and symbols.
+    for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) {
+      ineq[j] = lower ? -flatExpr[j] : flatExpr[j];
+    }
+    // Invalid bound: pos appears in `boundMap`.
+    // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or
+    // its callers to prevent invalid bounds from being added.
+    if (ineq[pos] != 0)
+      continue;
+    ineq[pos] = lower ? 1 : -1;
+    // Local columns of `ineq` are at the beginning.
+    unsigned j = getNumDimVars() + getNumSymbolVars();
+    unsigned end = flatExpr.size() - 1;
+    for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) {
+      ineq[j] = lower ? -flatExpr[i] : flatExpr[i];
+    }
+    // Make the bound closed in if flatExpr is open. The inequality is always
+    // created in the upper bound form, so the adjustment is -1.
+    int64_t boundAdjustment = (isClosedBound || type == BoundType::EQ) ? 0 : -1;
+    // Constant term.
+    ineq[getNumCols() - 1] = (lower ? -flatExpr[flatExpr.size() - 1]
+                                    : flatExpr[flatExpr.size() - 1]) +
+                             boundAdjustment;
+    type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq);
+  }
+
+  return success();
+}
+
+LogicalResult IntegerPolyhedron::addBound(BoundType type, unsigned pos,
+                                          AffineMap boundMap) {
+  return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB);
+}
+
 unsigned IntegerPolyhedron::insertVar(VarKind kind, unsigned pos,
                                       unsigned num) {
   assert((kind != VarKind::Domain || num == 0) &&
          "Domain has to be zero in a set");
   return IntegerRelation::insertVar(kind, pos, num);
 }
 IntegerPolyhedron
 IntegerPolyhedron::intersect(const IntegerPolyhedron &other) const {
   return IntegerPolyhedron(IntegerRelation::intersect(other));
 }
 
+/// Computes the lower and upper bounds of the first 'num' dimensional
+/// variables (starting at 'offset') as affine maps of the remaining
+/// variables (dimensional and symbolic variables). Local variables are
+/// themselves explicitly computed as affine functions of other variables in
+/// this process if needed.
+void IntegerPolyhedron::getSliceBounds(unsigned offset, unsigned num,
+                                       MLIRContext *context,
+                                       SmallVectorImpl<AffineMap> *lbMaps,
+                                       SmallVectorImpl<AffineMap> *ubMaps,
+                                       bool closedUB) {
+  assert(num < getNumDimVars() && "invalid range");
+
+  // Basic simplification.
+  normalizeConstraintsByGCD();
+
+  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
+                          << " variables\n");
+  LLVM_DEBUG(dump());
+
+  // Record computed/detected variables.
+  SmallVector<AffineExpr, 8> memo(getNumVars());
+  // Initialize dimensional and symbolic variables.
+  for (unsigned i = 0, e = getNumDimVars(); i < e; i++) {
+    if (i < offset)
+      memo[i] = getAffineDimExpr(i, context);
+    else if (i >= offset + num)
+      memo[i] = getAffineDimExpr(i - num, context);
+  }
+  for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++)
+    memo[i] = getAffineSymbolExpr(i - getNumDimVars(), context);
+
+  bool changed;
+  do {
+    changed = false;
+    // Identify yet unknown variables as constants or mod's / floordiv's of
+    // other variables if possible.
+    for (unsigned pos = 0; pos < getNumVars(); pos++) {
+      if (memo[pos])
+        continue;
+
+      auto lbConst = getConstantBound64(BoundType::LB, pos);
+      auto ubConst = getConstantBound64(BoundType::UB, pos);
+      if (lbConst.has_value() && ubConst.has_value()) {
+        // Detect equality to a constant.
+        if (*lbConst == *ubConst) {
+          memo[pos] = getAffineConstantExpr(*lbConst, context);
+          changed = true;
+          continue;
+        }
+
+        // Detect an variable as modulo of another variable w.r.t a
+        // constant.
+        if (detectAsMod(pos, *lbConst, *ubConst, memo, context)) {
+          changed = true;
+          continue;
+        }
+      }
+
+      // Detect an variable as a floordiv of an affine function of other
+      // variables (divisor is a positive constant).
+      if (detectAsFloorDiv(pos, context, memo)) {
+        changed = true;
+        continue;
+      }
+
+      // Detect an variable as an expression of other variables.
+      unsigned idx;
+      if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) {
+        continue;
+      }
+
+      // Build AffineExpr solving for variable 'pos' in terms of all others.
+      auto expr = getAffineConstantExpr(0, context);
+      unsigned j, e;
+      for (j = 0, e = getNumVars(); j < e; ++j) {
+        if (j == pos)
+          continue;
+        int64_t c = atEq64(idx, j);
+        if (c == 0)
+          continue;
+        // If any of the involved IDs hasn't been found yet, we can't proceed.
+        if (!memo[j])
+          break;
+        expr = expr + memo[j] * c;
+      }
+      if (j < e)
+        // Can't construct expression as it depends on a yet uncomputed
+        // variable.
+        continue;
+
+      // Add constant term to AffineExpr.
+      expr = expr + atEq64(idx, getNumVars());
+      int64_t vPos = atEq64(idx, pos);
+      assert(vPos != 0 && "expected non-zero here");
+      if (vPos > 0)
+        expr = (-expr).floorDiv(vPos);
+      else
+        // vPos < 0.
+        expr = expr.floorDiv(-vPos);
+      // Successfully constructed expression.
+      memo[pos] = expr;
+      changed = true;
+    }
+    // This loop is guaranteed to reach a fixed point - since once an
+    // variable's explicit form is computed (in memo[pos]), it's not updated
+    // again.
+  } while (changed);
+
+  int64_t ubAdjustment = closedUB ? 0 : 1;
+
+  // Set the lower and upper bound maps for all the variables that were
+  // computed as affine expressions of the rest as the "detected expr" and
+  // "detected expr + 1" respectively; set the undetected ones to null.
+  std::optional<IntegerPolyhedron> tmpClone;
+  for (unsigned pos = 0; pos < num; pos++) {
+    unsigned numMapDims = getNumDimVars() - num;
+    unsigned numMapSymbols = getNumSymbolVars();
+    AffineExpr expr = memo[pos + offset];
+    if (expr)
+      expr = simplifyAffineExpr(expr, numMapDims, numMapSymbols);
+
+    AffineMap &lbMap = (*lbMaps)[pos];
+    AffineMap &ubMap = (*ubMaps)[pos];
+
+    if (expr) {
+      lbMap = AffineMap::get(numMapDims, numMapSymbols, expr);
+      ubMap = AffineMap::get(numMapDims, numMapSymbols, expr + ubAdjustment);
+    } else {
+      // TODO: Whenever there are local variables in the dependence
+      // constraints, we'll conservatively over-approximate, since we don't
+      // always explicitly compute them above (in the while loop).
+      if (getNumLocalVars() == 0) {
+        // Work on a copy so that we don't update this constraint system.
+        if (!tmpClone) {
+          tmpClone.emplace(IntegerPolyhedron(*this));
+          // Removing redundant inequalities is necessary so that we don't get
+          // redundant loop bounds.
+          tmpClone->removeRedundantInequalities();
+        }
+        std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound(
+            pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context,
+            closedUB);
+      }
+
+      // If the above fails, we'll just use the constant lower bound and the
+      // constant upper bound (if they exist) as the slice bounds.
+      // TODO: being conservative for the moment in cases that
+      // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
+      // fixed (b/126426796).
+      if (!lbMap || lbMap.getNumResults() > 1) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "WARNING: Potentially over-approximating slice lb\n");
+        auto lbConst = getConstantBound64(BoundType::LB, pos + offset);
+        if (lbConst.has_value()) {
+          lbMap = AffineMap::get(numMapDims, numMapSymbols,
+                                 getAffineConstantExpr(*lbConst, context));
+        }
+      }
+      if (!ubMap || ubMap.getNumResults() > 1) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "WARNING: Potentially over-approximating slice ub\n");
+        auto ubConst = getConstantBound64(BoundType::UB, pos + offset);
+        if (ubConst.has_value()) {
+          ubMap = AffineMap::get(
+              numMapDims, numMapSymbols,
+              getAffineConstantExpr(*ubConst + ubAdjustment, context));
+        }
+      }
+    }
+    LLVM_DEBUG(llvm::dbgs()
+               << "lb map for pos = " << Twine(pos + offset) << ", expr: ");
+    LLVM_DEBUG(lbMap.dump(););
+    LLVM_DEBUG(llvm::dbgs()
+               << "ub map for pos = " << Twine(pos + offset) << ", expr: ");
+    LLVM_DEBUG(ubMap.dump(););
+  }
+}
+
+LogicalResult IntegerPolyhedron::flattenAlignedMapAndMergeLocals(
+    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs) {
+  IntegerPolyhedron localCst(PresburgerSpace::getSetSpace());
+  if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "composition unimplemented for semi-affine maps\n");
+    return failure();
+  }
+
+  // Add localCst information.
+  if (localCst.getNumLocalVars() > 0) {
+    unsigned numLocalVars = getNumLocalVars();
+    // Insert local dims of localCst at the beginning.
+    insertVar(VarKind::Local, /*pos=*/0, /*num=*/localCst.getNumLocalVars());
+    // Insert local dims of `this` at the end of localCst.
+    localCst.appendVar(VarKind::Local, /*num=*/numLocalVars);
+    // Dimensions of localCst and this constraint set match. Append localCst to
+    // this constraint set.
+    append(localCst);
+  }
+
+  return success();
+}
+
 PresburgerSet IntegerPolyhedron::subtract(const PresburgerSet &other) const {
   return PresburgerSet(IntegerRelation::subtract(other));
 }
+
+namespace {
+
+// See comments for SimpleAffineExprFlattener.
+// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
+// constraint information associated with mod's, floordiv's, and ceildiv's
+// in FlatAffineValueConstraints 'localVarCst'.
+struct AffineExprFlattener : public SimpleAffineExprFlattener {
+public:
+  // Constraints connecting newly introduced local variables (for mod's and
+  // div's) to existing (dimensional and symbolic) ones. These are always
+  // inequalities.
+  IntegerPolyhedron localVarCst;
+
+  AffineExprFlattener(unsigned nDims, unsigned nSymbols)
+      : SimpleAffineExprFlattener(nDims, nSymbols),
+        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
+
+private:
+  // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
+  // The local variable added is always a floordiv of a pure add/mul affine
+  // function of other variables, coefficients of which are specified in
+  // `dividend' and with respect to the positive constant `divisor'. localExpr
+  // is the simplified tree expression (AffineExpr) corresponding to the
+  // quantifier.
+  void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
+                          AffineExpr localExpr) override {
+    SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr);
+    // Update localVarCst.
+    localVarCst.addLocalFloorDiv(dividend, divisor);
+  }
+};
+
+} // namespace
+
+// Flattens the expressions in map. Returns failure if 'expr' was unable to be
+// flattened (i.e., semi-affine expressions not handled yet).
+static LogicalResult
+getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
+                        unsigned numSymbols,
+                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+                        IntegerPolyhedron *localVarCst) {
+  if (exprs.empty()) {
+    if (localVarCst)
+      *localVarCst =
+          IntegerPolyhedron(PresburgerSpace::getSetSpace(numDims, numSymbols));
+    return success();
+  }
+
+  AffineExprFlattener flattener(numDims, numSymbols);
+  // Use the same flattener to simplify each expression successively. This way
+  // local variables / expressions are shared.
+  for (auto expr : exprs) {
+    if (!expr.isPureAffine())
+      return failure();
+
+    flattener.walkPostOrder(expr);
+  }
+
+  assert(flattener.operandExprStack.size() == exprs.size());
+  flattenedExprs->clear();
+  flattenedExprs->assign(flattener.operandExprStack.begin(),
+                         flattener.operandExprStack.end());
+
+  if (localVarCst)
+    localVarCst->clearAndCopyFrom(flattener.localVarCst);
+
+  return success();
+}
+
+// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
+// be flattened (semi-affine expressions not handled yet).
+LogicalResult mlir::getFlattenedAffineExpr(
+    AffineExpr expr, unsigned numDims, unsigned numSymbols,
+    SmallVectorImpl<int64_t> *flattenedExpr, IntegerPolyhedron *localVarCst) {
+  std::vector<SmallVector<int64_t, 8>> flattenedExprs;
+  LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols,
+                                                &flattenedExprs, localVarCst);
+  *flattenedExpr = flattenedExprs[0];
+  return ret;
+}
+
+/// Flattens the expressions in map. Returns failure if 'expr' was unable to be
+/// flattened (i.e., semi-affine expressions not handled yet).
+LogicalResult mlir::getFlattenedAffineExprs(
+    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+    IntegerPolyhedron *localVarCst) {
+  if (map.getNumResults() == 0) {
+    if (localVarCst)
+      *localVarCst = IntegerPolyhedron(
+          PresburgerSpace::getSetSpace(map.getNumDims(), map.getNumSymbols()));
+    return success();
+  }
+  return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(),
+                                   map.getNumSymbols(), flattenedExprs,
+                                   localVarCst);
+}
+
+LogicalResult mlir::getFlattenedAffineExprs(
+    IntegerSet set, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+    IntegerPolyhedron *localVarCst) {
+  if (set.getNumConstraints() == 0) {
+    if (localVarCst)
+      *localVarCst = IntegerPolyhedron(
+          PresburgerSpace::getSetSpace(set.getNumDims(), set.getNumSymbols()));
+    return success();
+  }
+  return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(),
+                                   set.getNumSymbols(), flattenedExprs,
+                                   localVarCst);
+}

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
 #define DEBUG_TYPE "affine-structures"
 
 using namespace mlir;
 using namespace presburger;
 
-namespace {
-
-// See comments for SimpleAffineExprFlattener.
-// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
-// constraint information associated with mod's, floordiv's, and ceildiv's
-// in FlatAffineValueConstraints 'localVarCst'.
-struct AffineExprFlattener : public SimpleAffineExprFlattener {
-public:
-  // Constraints connecting newly introduced local variables (for mod's and
-  // div's) to existing (dimensional and symbolic) ones. These are always
-  // inequalities.
-  IntegerPolyhedron localVarCst;
-
-  AffineExprFlattener(unsigned nDims, unsigned nSymbols)
-      : SimpleAffineExprFlattener(nDims, nSymbols),
-        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
-
-private:
-  // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
-  // The local variable added is always a floordiv of a pure add/mul affine
-  // function of other variables, coefficients of which are specified in
-  // `dividend' and with respect to the positive constant `divisor'. localExpr
-  // is the simplified tree expression (AffineExpr) corresponding to the
-  // quantifier.
-  void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
-                          AffineExpr localExpr) override {
-    SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr);
-    // Update localVarCst.
-    localVarCst.addLocalFloorDiv(dividend, divisor);
-  }
-};
-
-} // namespace
-
-// Flattens the expressions in map. Returns failure if 'expr' was unable to be
-// flattened (i.e., semi-affine expressions not handled yet).
-static LogicalResult
-getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
-                        unsigned numSymbols,
-                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-                        FlatAffineValueConstraints *localVarCst) {
-  if (exprs.empty()) {
-    if (localVarCst)
-      *localVarCst = FlatAffineValueConstraints(numDims, numSymbols);
-    return success();
-  }
-
-  AffineExprFlattener flattener(numDims, numSymbols);
-  // Use the same flattener to simplify each expression successively. This way
-  // local variables / expressions are shared.
-  for (auto expr : exprs) {
-    if (!expr.isPureAffine())
-      return failure();
-
-    flattener.walkPostOrder(expr);
-  }
-
-  assert(flattener.operandExprStack.size() == exprs.size());
-  flattenedExprs->clear();
-  flattenedExprs->assign(flattener.operandExprStack.begin(),
-                         flattener.operandExprStack.end());
-
-  if (localVarCst)
-    localVarCst->clearAndCopyFrom(flattener.localVarCst);
-
-  return success();
-}
-
-// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
-// be flattened (semi-affine expressions not handled yet).
-LogicalResult
-mlir::getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
-                             unsigned numSymbols,
-                             SmallVectorImpl<int64_t> *flattenedExpr,
-                             FlatAffineValueConstraints *localVarCst) {
-  std::vector<SmallVector<int64_t, 8>> flattenedExprs;
-  LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols,
-                                                &flattenedExprs, localVarCst);
-  *flattenedExpr = flattenedExprs[0];
-  return ret;
-}
-
-/// Flattens the expressions in map. Returns failure if 'expr' was unable to be
-/// flattened (i.e., semi-affine expressions not handled yet).
-LogicalResult mlir::getFlattenedAffineExprs(
-    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-    FlatAffineValueConstraints *localVarCst) {
-  if (map.getNumResults() == 0) {
-    if (localVarCst)
-      *localVarCst =
-          FlatAffineValueConstraints(map.getNumDims(), map.getNumSymbols());
-    return success();
-  }
-  return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(),
-                                   map.getNumSymbols(), flattenedExprs,
-                                   localVarCst);
-}
-
-LogicalResult mlir::getFlattenedAffineExprs(
-    IntegerSet set, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
-    FlatAffineValueConstraints *localVarCst) {
-  if (set.getNumConstraints() == 0) {
-    if (localVarCst)
-      *localVarCst =
-          FlatAffineValueConstraints(set.getNumDims(), set.getNumSymbols());
-    return success();
-  }
-  return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(),
-                                   set.getNumSymbols(), flattenedExprs,
-                                   localVarCst);
-}
-
 //===----------------------------------------------------------------------===//
-// FlatAffineConstraints / FlatAffineValueConstraints.
+// FlatAffineValueConstraints.
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<FlatAffineValueConstraints>
 FlatAffineValueConstraints::clone() const {
   return std::make_unique<FlatAffineValueConstraints>(*this);
 }
 
 // Construct from an IntegerSet.
void FlatAffineValueConstraints::removeVarRange(VarKind kind, unsigned varStart,
   unsigned offset = getVarKindOffset(kind);
 
   if (kind != VarKind::Local) {
     values.erase(values.begin() + varStart + offset,
                  values.begin() + varLimit + offset);
   }
 }
 
-// Determine whether the variable at 'pos' (say var_r) can be expressed as
-// modulo of another known variable (say var_n) w.r.t a constant. For example,
-// if the following constraints hold true:
-// ```
-// 0 <= var_r <= divisor - 1
-// var_n - (divisor * q_expr) = var_r
-// ```
-// where `var_n` is a known variable (called dividend), and `q_expr` is an
-// `AffineExpr` (called the quotient expression), `var_r` can be written as:
-//
-// `var_r = var_n mod divisor`.
-//
-// Additionally, in a special case of the above constaints where `q_expr` is an
-// variable itself that is not yet known (say `var_q`), it can be written as a
-// floordiv in the following way:
-//
-// `var_q = var_n floordiv divisor`.
-//
-// Returns true if the above mod or floordiv are detected, updating 'memo' with
-// these new expressions. Returns false otherwise.
-static bool detectAsMod(const FlatAffineValueConstraints &cst, unsigned pos,
-                        int64_t lbConst, int64_t ubConst,
-                        SmallVectorImpl<AffineExpr> &memo,
-                        MLIRContext *context) {
-  assert(pos < cst.getNumVars() && "invalid position");
-
-  // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can
-  // be determined.
-  if (lbConst != 0 || ubConst < 1)
-    return false;
-  int64_t divisor = ubConst + 1;
-
-  // Check for the aforementioned conditions in each equality.
-  for (unsigned curEquality = 0, numEqualities = cst.getNumEqualities();
-       curEquality < numEqualities; curEquality++) {
-    int64_t coefficientAtPos = cst.atEq64(curEquality, pos);
-    // If current equality does not involve `var_r`, continue to the next
-    // equality.
-    if (coefficientAtPos == 0)
-      continue;
-
-    // Constant term should be 0 in this equality.
-    if (cst.atEq64(curEquality, cst.getNumCols() - 1) != 0)
-      continue;
-
-    // Traverse through the equality and construct the dividend expression
-    // `dividendExpr`, to contain all the variables which are known and are
-    // not divisible by `(coefficientAtPos * divisor)`. Hope here is that the
-    // `dividendExpr` gets simplified into a single variable `var_n` discussed
-    // above.
-    auto dividendExpr = getAffineConstantExpr(0, context);
-
-    // Track the terms that go into quotient expression, later used to detect
-    // additional floordiv.
-    unsigned quotientCount = 0;
-    int quotientPosition = -1;
-    int quotientSign = 1;
-
-    // Consider each term in the current equality.
-    unsigned curVar, e;
-    for (curVar = 0, e = cst.getNumDimAndSymbolVars(); curVar < e; ++curVar) {
-      // Ignore var_r.
-      if (curVar == pos)
-        continue;
-      int64_t coefficientOfCurVar = cst.atEq64(curEquality, curVar);
-      // Ignore vars that do not contribute to the current equality.
-      if (coefficientOfCurVar == 0)
-        continue;
-      // Check if the current var goes into the quotient expression.
-      if (coefficientOfCurVar % (divisor * coefficientAtPos) == 0) {
-        quotientCount++;
-        quotientPosition = curVar;
-        quotientSign = (coefficientOfCurVar * coefficientAtPos) > 0 ? 1 : -1;
-        continue;
-      }
-      // Variables that are part of dividendExpr should be known.
-      if (!memo[curVar])
-        break;
-      // Append the current variable to the dividend expression.
-      dividendExpr = dividendExpr + memo[curVar] * coefficientOfCurVar;
-    }
-
-    // Can't construct expression as it depends on a yet uncomputed var.
-    if (curVar < e)
-      continue;
-
-    // Express `var_r` in terms of the other vars collected so far.
-    if (coefficientAtPos > 0)
-      dividendExpr = (-dividendExpr).floorDiv(coefficientAtPos);
-    else
-      dividendExpr = dividendExpr.floorDiv(-coefficientAtPos);
-
-    // Simplify the expression.
-    dividendExpr = simplifyAffineExpr(dividendExpr, cst.getNumDimVars(),
-                                      cst.getNumSymbolVars());
-    // Only if the final dividend expression is just a single var (which we call
-    // `var_n`), we can proceed.
-    // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it
-    // to dims themselves.
-    auto dimExpr = dividendExpr.dyn_cast<AffineDimExpr>();
-    if (!dimExpr)
-      continue;
-
-    // Express `var_r` as `var_n % divisor` and store the expression in `memo`.
-    if (quotientCount >= 1) {
-      auto ub = cst.getConstantBound64(
-          FlatAffineValueConstraints::BoundType::UB, dimExpr.getPosition());
-      // If `var_n` has an upperbound that is less than the divisor, mod can be
-      // eliminated altogether.
-      if (ub && *ub < divisor)
-        memo[pos] = dimExpr;
-      else
-        memo[pos] = dimExpr % divisor;
-      // If a unique quotient `var_q` was seen, it can be expressed as
-      // `var_n floordiv divisor`.
-      if (quotientCount == 1 && !memo[quotientPosition])
-        memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign;
-
-      return true;
-    }
-  }
-  return false;
-}
-
-/// Check if the pos^th variable can be expressed as a floordiv of an affine
-/// function of other variables (where the divisor is a positive constant)
-/// given the initial set of expressions in `exprs`. If it can be, the
-/// corresponding position in `exprs` is set as the detected affine expr. For
-/// eg: 4q <= i + j <= 4q + 3   <=>   q = (i + j) floordiv 4. An equality can
-/// also yield a floordiv: eg.  4q = i + j <=> q = (i + j) floordiv 4. 32q + 28
-/// <= i <= 32q + 31 => q = i floordiv 32.
-static bool detectAsFloorDiv(const FlatAffineValueConstraints &cst,
-                             unsigned pos, MLIRContext *context,
-                             SmallVectorImpl<AffineExpr> &exprs) {
-  assert(pos < cst.getNumVars() && "invalid position");
-
-  // Get upper-lower bound pair for this variable.
-  SmallVector<bool, 8> foundRepr(cst.getNumVars(), false);
-  for (unsigned i = 0, e = cst.getNumVars(); i < e; ++i)
-    if (exprs[i])
-      foundRepr[i] = true;
-
-  SmallVector<int64_t, 8> dividend(cst.getNumCols());
-  unsigned divisor;
-  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor);
-
-  // No upper-lower bound pair found for this var.
-  if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality)
-    return false;
-
-  // Construct the dividend expression.
-  auto dividendExpr = getAffineConstantExpr(dividend.back(), context);
-  for (unsigned c = 0, f = cst.getNumVars(); c < f; c++)
-    if (dividend[c] != 0)
-      dividendExpr = dividendExpr + dividend[c] * exprs[c];
-
-  // Successfully detected the floordiv.
-  exprs[pos] = dividendExpr.floorDiv(divisor);
-  return true;
-}
-
-std::pair<AffineMap, AffineMap>
-FlatAffineValueConstraints::getLowerAndUpperBound(
-    unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
-    ArrayRef<AffineExpr> localExprs, MLIRContext *context,
-    bool closedUB) const {
-  assert(pos + offset < getNumDimVars() && "invalid dim start pos");
-  assert(symStartPos >= (pos + offset) && "invalid sym start pos");
-  assert(getNumLocalVars() == localExprs.size() &&
-         "incorrect local exprs count");
-
-  SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
-  getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices,
-                               offset, num);
-
-  /// Add to 'b' from 'a' in set [0, offset) U [offset + num, symbStartPos).
-  auto addCoeffs = [&](ArrayRef<int64_t> a, SmallVectorImpl<int64_t> &b) {
-    b.clear();
-    for (unsigned i = 0, e = a.size(); i < e; ++i) {
-      if (i < offset || i >= offset + num)
-        b.push_back(a[i]);
-    }
-  };
-
-  SmallVector<int64_t, 8> lb, ub;
-  SmallVector<AffineExpr, 4> lbExprs;
-  unsigned dimCount = symStartPos - num;
-  unsigned symCount = getNumDimAndSymbolVars() - symStartPos;
-  lbExprs.reserve(lbIndices.size() + eqIndices.size());
-  // Lower bound expressions.
-  for (auto idx : lbIndices) {
-    auto ineq = getInequality64(idx);
-    // Extract the lower bound (in terms of other coeff's + const), i.e., if
-    // i - j + 1 >= 0 is the constraint, 'pos' is for i the lower bound is j
-    // - 1.
-    addCoeffs(ineq, lb);
-    std::transform(lb.begin(), lb.end(), lb.begin(), std::negate<int64_t>());
-    auto expr =
-        getAffineExprFromFlatForm(lb, dimCount, symCount, localExprs, context);
-    // expr ceildiv divisor is (expr + divisor - 1) floordiv divisor
-    int64_t divisor = std::abs(ineq[pos + offset]);
-    expr = (expr + divisor - 1).floorDiv(divisor);
-    lbExprs.push_back(expr);
-  }
-
-  SmallVector<AffineExpr, 4> ubExprs;
-  ubExprs.reserve(ubIndices.size() + eqIndices.size());
-  // Upper bound expressions.
-  for (auto idx : ubIndices) {
-    auto ineq = getInequality64(idx);
-    // Extract the upper bound (in terms of other coeff's + const).
-    addCoeffs(ineq, ub);
-    auto expr =
-        getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context);
-    expr = expr.floorDiv(std::abs(ineq[pos + offset]));
-    int64_t ubAdjustment = closedUB ? 0 : 1;
-    ubExprs.push_back(expr + ubAdjustment);
-  }
-
-  // Equalities. It's both a lower and a upper bound.
-  SmallVector<int64_t, 4> b;
-  for (auto idx : eqIndices) {
-    auto eq = getEquality64(idx);
-    addCoeffs(eq, b);
-    if (eq[pos + offset] > 0)
-      std::transform(b.begin(), b.end(), b.begin(), std::negate<int64_t>());
-
-    // Extract the upper bound (in terms of other coeff's + const).
-    auto expr =
-        getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
-    expr = expr.floorDiv(std::abs(eq[pos + offset]));
-    // Upper bound is exclusive.
-    ubExprs.push_back(expr + 1);
-    // Lower bound.
-    expr =
-        getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
-    expr = expr.ceilDiv(std::abs(eq[pos + offset]));
-    lbExprs.push_back(expr);
-  }
-
-  auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context);
-  auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context);
-
-  return {lbMap, ubMap};
-}
-
-/// Computes the lower and upper bounds of the first 'num' dimensional
-/// variables (starting at 'offset') as affine maps of the remaining
-/// variables (dimensional and symbolic variables). Local variables are
-/// themselves explicitly computed as affine functions of other variables in
-/// this process if needed.
-void FlatAffineValueConstraints::getSliceBounds(
-    unsigned offset, unsigned num, MLIRContext *context,
-    SmallVectorImpl<AffineMap> *lbMaps, SmallVectorImpl<AffineMap> *ubMaps,
-    bool closedUB) {
-  assert(num < getNumDimVars() && "invalid range");
-
-  // Basic simplification.
-  normalizeConstraintsByGCD();
-
-  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
-                          << " variables\n");
-  LLVM_DEBUG(dump());
-
-  // Record computed/detected variables.
-  SmallVector<AffineExpr, 8> memo(getNumVars());
-  // Initialize dimensional and symbolic variables.
-  for (unsigned i = 0, e = getNumDimVars(); i < e; i++) {
-    if (i < offset)
-      memo[i] = getAffineDimExpr(i, context);
-    else if (i >= offset + num)
-      memo[i] = getAffineDimExpr(i - num, context);
-  }
-  for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++)
-    memo[i] = getAffineSymbolExpr(i - getNumDimVars(), context);
-
-  bool changed;
-  do {
-    changed = false;
-    // Identify yet unknown variables as constants or mod's / floordiv's of
-    // other variables if possible.
-    for (unsigned pos = 0; pos < getNumVars(); pos++) {
-      if (memo[pos])
-        continue;
-
-      auto lbConst = getConstantBound64(BoundType::LB, pos);
-      auto ubConst = getConstantBound64(BoundType::UB, pos);
-      if (lbConst.has_value() && ubConst.has_value()) {
-        // Detect equality to a constant.
-        if (*lbConst == *ubConst) {
-          memo[pos] = getAffineConstantExpr(*lbConst, context);
-          changed = true;
-          continue;
-        }
-
-        // Detect an variable as modulo of another variable w.r.t a
-        // constant.
-        if (detectAsMod(*this, pos, *lbConst, *ubConst, memo, context)) {
-          changed = true;
-          continue;
-        }
-      }
-
-      // Detect an variable as a floordiv of an affine function of other
-      // variables (divisor is a positive constant).
-      if (detectAsFloorDiv(*this, pos, context, memo)) {
-        changed = true;
-        continue;
-      }
-
-      // Detect an variable as an expression of other variables.
-      unsigned idx;
-      if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) {
-        continue;
-      }
-
-      // Build AffineExpr solving for variable 'pos' in terms of all others.
-      auto expr = getAffineConstantExpr(0, context);
-      unsigned j, e;
-      for (j = 0, e = getNumVars(); j < e; ++j) {
-        if (j == pos)
-          continue;
-        int64_t c = atEq64(idx, j);
-        if (c == 0)
-          continue;
-        // If any of the involved IDs hasn't been found yet, we can't proceed.
-        if (!memo[j])
-          break;
-        expr = expr + memo[j] * c;
-      }
-      if (j < e)
-        // Can't construct expression as it depends on a yet uncomputed
-        // variable.
-        continue;
-
-      // Add constant term to AffineExpr.
-      expr = expr + atEq64(idx, getNumVars());
-      int64_t vPos = atEq64(idx, pos);
-      assert(vPos != 0 && "expected non-zero here");
-      if (vPos > 0)
-        expr = (-expr).floorDiv(vPos);
-      else
-        // vPos < 0.
-        expr = expr.floorDiv(-vPos);
-      // Successfully constructed expression.
-      memo[pos] = expr;
-      changed = true;
-    }
-    // This loop is guaranteed to reach a fixed point - since once an
-    // variable's explicit form is computed (in memo[pos]), it's not updated
-    // again.
-  } while (changed);
-
-  int64_t ubAdjustment = closedUB ? 0 : 1;
-
-  // Set the lower and upper bound maps for all the variables that were
-  // computed as affine expressions of the rest as the "detected expr" and
-  // "detected expr + 1" respectively; set the undetected ones to null.
-  std::optional<FlatAffineValueConstraints> tmpClone;
-  for (unsigned pos = 0; pos < num; pos++) {
-    unsigned numMapDims = getNumDimVars() - num;
-    unsigned numMapSymbols = getNumSymbolVars();
-    AffineExpr expr = memo[pos + offset];
-    if (expr)
-      expr = simplifyAffineExpr(expr, numMapDims, numMapSymbols);
-
-    AffineMap &lbMap = (*lbMaps)[pos];
-    AffineMap &ubMap = (*ubMaps)[pos];
-
-    if (expr) {
-      lbMap = AffineMap::get(numMapDims, numMapSymbols, expr);
-      ubMap = AffineMap::get(numMapDims, numMapSymbols, expr + ubAdjustment);
-    } else {
-      // TODO: Whenever there are local variables in the dependence
-      // constraints, we'll conservatively over-approximate, since we don't
-      // always explicitly compute them above (in the while loop).
-      if (getNumLocalVars() == 0) {
-        // Work on a copy so that we don't update this constraint system.
-        if (!tmpClone) {
-          tmpClone.emplace(FlatAffineValueConstraints(*this));
-          // Removing redundant inequalities is necessary so that we don't get
-          // redundant loop bounds.
-          tmpClone->removeRedundantInequalities();
-        }
-        std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound(
-            pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context,
-            closedUB);
-      }
-
-      // If the above fails, we'll just use the constant lower bound and the
-      // constant upper bound (if they exist) as the slice bounds.
-      // TODO: being conservative for the moment in cases that
-      // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
-      // fixed (b/126426796).
-      if (!lbMap || lbMap.getNumResults() > 1) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "WARNING: Potentially over-approximating slice lb\n");
-        auto lbConst = getConstantBound64(BoundType::LB, pos + offset);
-        if (lbConst.has_value()) {
-          lbMap = AffineMap::get(numMapDims, numMapSymbols,
-                                 getAffineConstantExpr(*lbConst, context));
-        }
-      }
-      if (!ubMap || ubMap.getNumResults() > 1) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "WARNING: Potentially over-approximating slice ub\n");
-        auto ubConst = getConstantBound64(BoundType::UB, pos + offset);
-        if (ubConst.has_value()) {
-          ubMap = AffineMap::get(
-              numMapDims, numMapSymbols,
-              getAffineConstantExpr(*ubConst + ubAdjustment, context));
-        }
-      }
-    }
-    LLVM_DEBUG(llvm::dbgs()
-               << "lb map for pos = " << Twine(pos + offset) << ", expr: ");
-    LLVM_DEBUG(lbMap.dump(););
-    LLVM_DEBUG(llvm::dbgs()
-               << "ub map for pos = " << Twine(pos + offset) << ", expr: ");
-    LLVM_DEBUG(ubMap.dump(););
-  }
-}
-
-LogicalResult FlatAffineValueConstraints::flattenAlignedMapAndMergeLocals(
-    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs) {
-  FlatAffineValueConstraints localCst;
-  if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) {
-    LLVM_DEBUG(llvm::dbgs()
-               << "composition unimplemented for semi-affine maps\n");
-    return failure();
-  }
-
-  // Add localCst information.
-  if (localCst.getNumLocalVars() > 0) {
-    unsigned numLocalVars = getNumLocalVars();
-    // Insert local dims of localCst at the beginning.
-    insertLocalVar(/*pos=*/0, /*num=*/localCst.getNumLocalVars());
-    // Insert local dims of `this` at the end of localCst.
-    localCst.appendLocalVar(/*num=*/numLocalVars);
-    // Dimensions of localCst and this constraint set match. Append localCst to
-    // this constraint set.
-    append(localCst);
-  }
-
-  return success();
-}
-
-LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
-                                                   AffineMap boundMap,
-                                                   bool isClosedBound) {
-  assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch");
-  assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
-  assert(pos < getNumDimAndSymbolVars() && "invalid position");
-  assert((type != BoundType::EQ || isClosedBound) &&
-         "EQ bound must be closed.");
-
-  // Equality follows the logic of lower bound except that we add an equality
-  // instead of an inequality.
-  assert((type != BoundType::EQ || boundMap.getNumResults() == 1) &&
-         "single result expected");
-  bool lower = type == BoundType::LB || type == BoundType::EQ;
-
-  std::vector<SmallVector<int64_t, 8>> flatExprs;
-  if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
-    return failure();
-  assert(flatExprs.size() == boundMap.getNumResults());
-
-  // Add one (in)equality for each result.
-  for (const auto &flatExpr : flatExprs) {
-    SmallVector<int64_t> ineq(getNumCols(), 0);
-    // Dims and symbols.
-    for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) {
-      ineq[j] = lower ? -flatExpr[j] : flatExpr[j];
-    }
-    // Invalid bound: pos appears in `boundMap`.
-    // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or
-    // its callers to prevent invalid bounds from being added.
-    if (ineq[pos] != 0)
-      continue;
-    ineq[pos] = lower ? 1 : -1;
-    // Local columns of `ineq` are at the beginning.
-    unsigned j = getNumDimVars() + getNumSymbolVars();
-    unsigned end = flatExpr.size() - 1;
-    for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) {
-      ineq[j] = lower ? -flatExpr[i] : flatExpr[i];
-    }
-    // Make the bound closed in if flatExpr is open. The inequality is always
-    // created in the upper bound form, so the adjustment is -1.
-    int64_t boundAdjustment = (isClosedBound || type == BoundType::EQ) ? 0 : -1;
-    // Constant term.
-    ineq[getNumCols() - 1] = (lower ? -flatExpr[flatExpr.size() - 1]
-                                    : flatExpr[flatExpr.size() - 1]) +
-                             boundAdjustment;
-    type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq);
-  }
-
-  return success();
-}
-
-LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
-                                                   AffineMap boundMap) {
-  return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB);
-}
-
 AffineMap
 FlatAffineValueConstraints::computeAlignedMap(AffineMap map,
                                               ValueRange operands) const {
   assert(map.getNumInputs() == operands.size() && "number of inputs mismatch");
 
   SmallVector<Value> dims, syms;
 #ifndef NDEBUG
   SmallVector<Value> newSyms;
for (unsigned i = getVarKindOffset(VarKind::SetDim),
     dims.push_back(values[i] ? *values[i] : Value());
   for (unsigned i = getVarKindOffset(VarKind::Symbol),
                 e = getVarKindEnd(VarKind::Symbol);
        i < e; ++i)
     syms.push_back(values[i] ? *values[i] : Value());
 
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
-  // All symbols are already part of this FlatAffineConstraints.
+  // All symbols are already part of this FlatAffineValueConstraints.
   assert(syms.size() == newSymsPtr->size() && "unexpected new/missing symbols");
   assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()) &&
          "unexpected new/missing symbols");
   return alignedMap;
 }
 
 LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
                                                    AffineMap boundMap,
static LogicalResult computeLocalVars(const FlatAffineValueConstraints &cst,
   bool changed;
   do {
     // Each time `changed` is true at the end of this iteration, one or more
     // local vars would have been detected as floordivs and set in memo; so the
     // number of null entries in memo[...] strictly reduces; so this converges.
     changed = false;
     for (unsigned i = 0, e = cst.getNumLocalVars(); i < e; ++i)
       if (!memo[numDims + numSyms + i] &&
-          detectAsFloorDiv(cst, /*pos=*/numDims + numSyms + i, context, memo))
+          cst.detectAsFloorDiv(/*pos=*/numDims + numSyms + i, context, memo))
         changed = true;
   } while (changed);
 
   ArrayRef<AffineExpr> localExprs =
       ArrayRef<AffineExpr>(memo).take_back(cst.getNumLocalVars());
   return success(
       llvm::all_of(localExprs, [](AffineExpr expr) { return expr; }));
 }