[MLIR] NFC. Refactor IntegerRelation getSliceBounds

Refactor FlatLinearConstraints getSliceBounds. The method was too long
and nested. NFC.

Author: bondhugula

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

@@ -738,6 +738,12 @@ class IntegerRelation {
   /// Same as findSymbolicIntegerLexMin but produces lexmax instead of lexmin
   SymbolicLexOpt findSymbolicIntegerLexMax() const;
 
+  /// Searches for a constraint with a non-zero coefficient at `colIdx` in
+  /// equality (isEq=true) or inequality (isEq=false) constraints.
+  /// Returns true and sets row found in search in `rowIdx`, false otherwise.
+  bool findConstraintWithNonZeroAt(unsigned colIdx, bool isEq,
+                                   unsigned *rowIdx) const;
+
   /// Return the set difference of this set and the given set, i.e.,
   /// return `this \ set`.
   PresburgerRelation subtract(const PresburgerRelation &set) const;
@@ -820,12 +826,6 @@ class IntegerRelation {
   /// Normalized each constraints by the GCD of its coefficients.
   void normalizeConstraintsByGCD();
 
-  /// Searches for a constraint with a non-zero coefficient at `colIdx` in
-  /// equality (isEq=true) or inequality (isEq=false) constraints.
-  /// Returns true and sets row found in search in `rowIdx`, false otherwise.
-  bool findConstraintWithNonZeroAt(unsigned colIdx, bool isEq,
-                                   unsigned *rowIdx) const;
-
   /// Returns true if the pos^th column is all zero for both inequalities and
   /// equalities.
   bool isColZero(unsigned pos) const;

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -581,48 +581,35 @@ std::pair<AffineMap, AffineMap> FlatLinearConstraints::getLowerAndUpperBound(
   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 FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
-                                           MLIRContext *context,
-                                           SmallVectorImpl<AffineMap> *lbMaps,
-                                           SmallVectorImpl<AffineMap> *ubMaps,
-                                           bool closedUB) {
-  assert(offset + num <= getNumDimVars() && "invalid range");
-
-  // Basic simplification.
-  normalizeConstraintsByGCD();
-
-  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for variables at positions ["
-                          << offset << ", " << offset + num << ")\n");
-  LLVM_DEBUG(dumpPretty());
-
-  // Record computed/detected variables.
-  SmallVector<AffineExpr, 8> memo(getNumVars());
+/// Compute a representation of `num` identifiers starting at `offset` in `cst`
+/// as affine expressions involving other known identifiers. Each identifier's
+/// expression (in terms of known identifiers) is populated into `memo`.
+static void computeUnknownVars(const FlatLinearConstraints &cst,
+                               MLIRContext *context, unsigned offset,
+                               unsigned num,
+                               SmallVectorImpl<AffineExpr> &memo) {
   // Initialize dimensional and symbolic variables.
-  for (unsigned i = 0, e = getNumDimVars(); i < e; i++) {
+  for (unsigned i = 0, e = cst.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);
+  for (unsigned i = cst.getNumDimVars(), e = cst.getNumDimAndSymbolVars();
+       i < e; i++)
+    memo[i] = getAffineSymbolExpr(i - cst.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++) {
+    for (unsigned pos = 0, f = cst.getNumVars(); pos < f; pos++) {
       if (memo[pos])
         continue;
 
-      auto lbConst = getConstantBound64(BoundType::LB, pos);
-      auto ubConst = getConstantBound64(BoundType::UB, pos);
+      auto lbConst = cst.getConstantBound64(BoundType::LB, pos);
+      auto ubConst = cst.getConstantBound64(BoundType::UB, pos);
       if (lbConst.has_value() && ubConst.has_value()) {
         // Detect equality to a constant.
         if (*lbConst == *ubConst) {
@@ -633,7 +620,7 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
 
         // Detect a variable as modulo of another variable w.r.t a
         // constant.
-        if (detectAsMod(*this, pos, offset, num, *lbConst, *ubConst, context,
+        if (detectAsMod(cst, pos, offset, num, *lbConst, *ubConst, context,
                         memo)) {
           changed = true;
           continue;
@@ -642,24 +629,24 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
 
       // Detect a variable as a floordiv of an affine function of other
       // variables (divisor is a positive constant).
-      if (detectAsFloorDiv(*this, pos, context, memo)) {
+      if (detectAsFloorDiv(cst, pos, context, memo)) {
         changed = true;
         continue;
       }
 
       // Detect a variable as an expression of other variables.
       unsigned idx;
-      if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) {
+      if (!cst.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) {
+      for (j = 0, e = cst.getNumVars(); j < e; ++j) {
         if (j == pos)
           continue;
-        int64_t c = atEq64(idx, j);
+        int64_t c = cst.atEq64(idx, j);
         if (c == 0)
           continue;
         // If any of the involved IDs hasn't been found yet, we can't proceed.
@@ -673,8 +660,8 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
         continue;
 
       // Add constant term to AffineExpr.
-      expr = expr + atEq64(idx, getNumVars());
-      int64_t vPos = atEq64(idx, pos);
+      expr = expr + cst.atEq64(idx, cst.getNumVars());
+      int64_t vPos = cst.atEq64(idx, pos);
       assert(vPos != 0 && "expected non-zero here");
       if (vPos > 0)
         expr = (-expr).floorDiv(vPos);
@@ -689,6 +676,30 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
     // variable's explicit form is computed (in memo[pos]), it's not updated
     // again.
   } while (changed);
+}
+
+/// 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 FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
+                                           MLIRContext *context,
+                                           SmallVectorImpl<AffineMap> *lbMaps,
+                                           SmallVectorImpl<AffineMap> *ubMaps,
+                                           bool closedUB) {
+  assert(offset + num <= getNumDimVars() && "invalid range");
+
+  // Basic simplification.
+  normalizeConstraintsByGCD();
+
+  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for variables at positions ["
+                          << offset << ", " << offset + num << ")\n");
+  LLVM_DEBUG(dumpPretty());
+
+  // Record computed/detected variables.
+  SmallVector<AffineExpr, 8> memo(getNumVars());
+  computeUnknownVars(*this, context, offset, num, memo);
 
   int64_t ubAdjustment = closedUB ? 0 : 1;