[MLIR] FlatAffineConstraints: Refactored computation of explicit representation for identifiers

[Groverkss]

Summary

This patch refactors the existing implementation of computing an explicit
representation of an identifier as a floordiv in terms of other identifiers and
exposes this computation as a public function.

The computation of this representation is required to support local identifiers
in PresburgerSet subtract, complement and isEqual.

[arjunp]

Thanks for the patch!

I didn't mention it everywhere but put please put fullstops at the end of sentences in comments.

mlir/lib/Analysis/AffineStructures.cpp

1247

Add an assert for the length of exprs?

1253–1254

Might be more helpful as a reference to understand the implementation if you include the most intuitive form first and then the normal form as it is stored.

// `id` is equivalent to `expr floordiv divisor` if there are constraints of the form
  // 0 <= divisor * id - expr <= divisor - 1
  // Rearranging, we have 
  // divisor * id - expr >=  0    <-- Lower bound for 'id'
  // -divisor * id + expr + (divisor - 1) >= 0    <-- bound for 'id'

1260–1261

In general, please avoid using auto unless it aids readability

1264–1273

Maybe just setting divisor = cst.atIneq(lbPos, cst.getNumCols() - 1) + cst.atIneq(ubPos, cst.getNumCols() - 1) + 1; would be cleaner?

1271

nit: Please put the full stop

1275

Mention in the comment that the constant term is not and should not be checked.

1286

nit: Please end the sentence with a full stop

1302

Can't you just use {ubPos, lbPos}?

1318

nit: you can use getNumIds() instead of getNumCols() - 1

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1268

Technically, even if the divisor is one it would work right? (I guess it would just be an equality then. But it still yields a representation for that variable)

1281–1282

nit: if you factor out checking an upper-lower bound pair into a separate function, then all this would look cleaner I feel.

1301

I don't like setting this inside. I would take foundRepr as a const reference and set it in the caller.

1331

In general, please avoid using auto.

1331–1335

The if (auto opt = ... idiom might be more readable:

if (auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i)) {
  repr[i] = res;
  changed = false;
}

1620

I believe it is preferred to drop these braces in LLVM style

1630

The lower bound is just divisor * id >= expr right?

1636–1652

I think you mean "the division variable"

mlir/unittests/Analysis/AffineStructuresTest.cpp

598

nit: seems a bit odd to use a reference here

608–609

Should it be this?

// divisor * id >=  expr     <-- Lower bound for 'id'
// divisor * id <=  expr + divisor - 1                  <-- Upper bound for 'id'

623

Can you pass the numerators and denominators for the floordivs and check that the inequalities returned are actually correct for each floordiv?

626–658

Please add some comments saying what the tests are checking.

[Groverkss]

Thank you for the review!

mlir/lib/Analysis/AffineStructures.cpp

1302

I'm not sure why but it does not work. Gives a no-matching constructor error for me.

mlir/unittests/Analysis/AffineStructuresTest.cpp

623

I'm not sure if that is required. If two inequalities are found which satisfy the checks given before this, the variable can be represented as a floordiv. It may not give the same floordiv that was added.

[Groverkss]

Addressed review comments

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1302

Does llvm::Optional({ubPos, lbPos}) also not work?

mlir/unittests/Analysis/AffineStructuresTest.cpp

608–609

expr + divisor - 1, not minus?

623

It seems a bit funny to copy paste the algorithm's logic again to check if the algorithm is correct.

The tests should trivially verify the answer IMO. I would just use tests where there is only one possible answer and check directly.

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1623–1624

Shouldn't the constant term just be the constant term of expr?

[Groverkss]

Fixed accidently removed commits

[Groverkss]

• Improved testing for computing divisions and fixed some bugs

[Groverkss]

Addressed more review comments.

mlir/lib/Analysis/AffineStructures.cpp

1302

This does not work since llvm::Optional requires template arguments. std::make_pair does work so I have used that instead.

1331

Due to the return type being too big, I have used auto here to improve readability. Let me know if it would still be better to replace it.

[Groverkss]

• Added punctuation in comments

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1255

Can you use the 0 <= expr - divisor*id <= divisor - 1 form of the inequalities everywhere? I feel this is much more intuitive since you can immediately see that expr - divisor * id is the remainder and this should be between 0 and divisor - 1.

1623–1624

If you change the representation to divisor - 1 >= expr - divisor*id >= 0 then this will become the lower bound.

1651

nit: gurantees -> guarantees

mlir/unittests/Analysis/AffineStructuresTest.cpp

614–616

nit: I think EXPECT_FALSE(res[i].hasValue()) is more intuitive

624–625

If you're changing the representation you'll have to change it here as well

[bondhugula]

Can you please separate out the bugfix from the refactoring?

mlir/include/mlir/Analysis/AffineStructures.h

164

These can lead to expensive copies. Could you take an output argument instead?

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

704

Use SmallVector to avoid heap allocation.

[Groverkss]

Moved the bugfix to a separate patch: https://reviews.llvm.org/D107214

[Groverkss]

Rebased to latest main branch

[Groverkss]

• Addressed comments

[Groverkss]

Addressed Comments.

mlir/lib/Analysis/AffineStructures.cpp

1623–1624

If I understand correctly, the upper bound is defined based on the value of the division variable and not on the value of the remainder of the division. So it will still be the upper bound. I don't think this needs to be changed.

[arjunp]

Looks good to me, please resolve the clang-format issues though.

Regarding the use of auto, I'm not sure where the line is drawn generally. It seems okay to me. Let's see what the others think.

[Groverkss]

• Fixed clang-format issues

[grosser]

@Groverkss, the commit message still references the bugfix that was factored out.

[Groverkss]

Hi @bondhugula, @vinayaka-polymage, I think I've addressed all the comments. Can you please check if there are any more changes to be made?

[Groverkss]

Hi @bondhugula , @vinayaka-polymage , gentle reminder 🙂 Are there any other changes to be made?

[arjunp]

Seems good to me. Let's also wait to see if @bondhugula or @vinayaka-polymage have any additional comments.

[vinayaka-polymage]

Hi @arjunp @Groverkss , sorry - I am currently unavailable until the end of the week, and will not be able to review this. I will check this during the weekend if it is still open.

[bondhugula]

This overall is looking good to me. Some questions and comments for minor improvements I feel.

mlir/include/mlir/Analysis/AffineStructures.h

159

indices may be confusing here. Consider rephrasing

... pairs of inequalities identified by their position indices, ...

?

162

Please rephrase the last two lines for clarity.

163

It's perhaps better to be explicit here: computeLocalRepresentation?

163

Is this name appropriate? It's not computing the local representation but only finding the pairs to use to construct the local representation. Better name?

getLocalReprLbUbPairs?

On this note, do we need to expose this publicly? Shouldn't detectAsFloorDiv be the one to expose and the former is just a local helper?

mlir/lib/Analysis/AffineStructures.cpp

1242

Put foundRepr in backticks. Similarly, at any places in doc comments.

1248

Assertion message please.

1328

SmallVector here to avoid a heap allocation.

1619

Can you use a SmallVector here to prevent a heap allocation?

mlir/unittests/Analysis/AffineStructuresTest.cpp

598

This should be marked static?

[arjunp]

mlir/include/mlir/Analysis/AffineStructures.h

163

PresburgerSet will require access to this function to support divisions in subtract, intersect, and equality checks. We could make this private if we are find with making PresburgerSet a friend class.

[Groverkss]

• Addressed @bondhugula 's comments

[Groverkss]

Rebase

[bondhugula]

Looks good - thanks!

mlir/include/mlir/Analysis/AffineStructures.h

163

It's fine as is then.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// Returns an integer sample point if one exists, or an empty Optional
   /// otherwise.
   Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
 
   /// Returns true if the given point satisfies the constraints, or false
   /// otherwise.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
-  /// Clones this object.
+  /// Find pairs of inequalities identified by their position indices, using
+  /// which an explicit representation for each local variable can be computed
+  /// The pairs are stored as indices of upperbound, lowerbound
+  /// inequalities. If no such pair can be found, it is stored as llvm::None.
+  void getLocalReprLbUbPairs(
+      std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
+
+  // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
   /// Returns the value at the specified equality row and column.
   inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); }
   inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
 
   /// Returns the value at the specified inequality row and column.
   inline int64_t atIneq(unsigned i, unsigned j) const {

mlir/lib/Analysis/AffineStructures.cpp

bool FlatAffineConstraints::containsPoint(ArrayRef<int64_t> point) const {
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     if (valueAt(getInequality(i), point) < 0)
       return false;
   }
   return true;
 }
 
+/// Check if the pos^th identifier can be expressed as a floordiv of an affine
+/// function of other identifiers (where the divisor is a positive constant),
+/// `foundRepr` contains a boolean for each identifier indicating if the
+/// explicit representation for that identifier has already been computed.
+static Optional<std::pair<unsigned, unsigned>>
+computeSingleVarRepr(const FlatAffineConstraints &cst,
+                     const SmallVector<bool, 8> &foundRepr, unsigned pos) {
+  assert(pos < cst.getNumIds() && "invalid position");
+  assert(foundRepr.size() == cst.getNumIds() &&
+         "Size of foundRepr does not match total number of variables");
+
+  SmallVector<unsigned, 4> lbIndices, ubIndices;
+  cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
+
+  // `id` is equivalent to `expr floordiv divisor` if there
+  // are constraints of the form:
+  //      0 <= expr - divisor * id <= divisor - 1
+  // Rearranging, we have:
+  //       divisor * id - expr + (divisor - 1) >= 0  <-- Lower bound for 'id'
+  //      -divisor * id + expr                 >= 0  <-- Upper bound for 'id'
+  //
+  // For example:
+  //       32*k >= 16*i + j - 31                 <-- Lower bound for 'k'
+  //       32*k  <= 16*i + j                     <-- Upper bound for 'k'
+  //       expr = 16*i + j, divisor = 32
+  //       k = ( 16*i + j ) floordiv 32
+  //
+  //       4q >= i + j - 2                       <-- Lower bound for 'q'
+  //       4q <= i + j + 1                       <-- Upper bound for 'q'
+  //       expr = i + j + 1, divisor = 4
+  //       q = (i + j + 1) floordiv 4
+  for (unsigned ubPos : ubIndices) {
+    for (unsigned lbPos : lbIndices) {
+      // Due to the form of the inequalities, sum of constants of the
+      // inequalities is (divisor - 1).
+      int64_t divisor = cst.atIneq(lbPos, cst.getNumCols() - 1) +
+                        cst.atIneq(ubPos, cst.getNumCols() - 1) + 1;
+
+      // Divisor should be positive.
+      if (divisor <= 0)
+        continue;
+
+      // Check if coeff of variable is equal to divisor.
+      if (divisor != cst.atIneq(lbPos, pos))
+        continue;
+
+      // Check if constraints are opposite of each other. Constant term
+      // is not required to be opposite and is not checked.
+      unsigned c = 0, f = 0;
+      for (c = 0, f = cst.getNumIds(); c < f; ++c)
+        if (cst.atIneq(ubPos, c) != -cst.atIneq(lbPos, c))
+          break;
+
+      if (c < f)
+        continue;
+
+      // Check if the inequalities depend on a variable for which
+      // an explicit representation has not been found yet.
+      // Exit to avoid circular dependencies between divisions.
+      for (c = 0, f = cst.getNumIds(); c < f; ++c) {
+        if (c == pos)
+          continue;
+        if (!foundRepr[c] && cst.atIneq(lbPos, c) != 0)
+          break;
+      }
+
+      // Expression can't be constructed as it depends on a yet unknown
+      // identifier.
+      // TODO: Visit/compute the identifiers in an order so that this doesn't
+      // happen. More complex but much more efficient.
+      if (c < f)
+        continue;
+
+      return std::make_pair(ubPos, lbPos);
+    }
+  }
+
+  return llvm::None;
+}
+
+/// Find pairs of inequalities identified by their position indices, using
+/// which an explicit representation for each local variable can be computed
+/// The pairs are stored as indices of upperbound, lowerbound
+/// inequalities. If no such pair can be found, it is stored as llvm::None.
+void FlatAffineConstraints::getLocalReprLbUbPairs(
+    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+  assert(repr.size() == getNumLocalIds() &&
+         "Size of repr does not match number of local variables");
+
+  SmallVector<bool, 8> foundRepr(getNumIds(), false);
+  for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i)
+    foundRepr[i] = true;
+
+  unsigned divOffset = getNumDimAndSymbolIds();
+  bool changed;
+  do {
+    // Each time changed is true, at end of this iteration, one or more local
+    // vars have been detected as floor divs.
+    changed = false;
+    for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) {
+      if (!foundRepr[i + divOffset]) {
+        if (auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i)) {
+          foundRepr[i + divOffset] = true;
+          repr[i] = res;
+          changed = true;
+        }
+      }
+    }
+  } while (changed);
+}
+
 /// Tightens inequalities given that we are dealing with integer spaces. This is
 /// analogous to the GCD test but applied to inequalities. The constant term can
 /// be reduced to the preceding multiple of the GCD of the coefficients, i.e.,
 ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
 /// fast method - linear in the number of coefficients.
 // Example on how this affects practical cases: consider the scenario:
 // 64*i >= 100, j = 64*i; without a tightening, elimination of i would yield
 // j >= 100 instead of the tighter (exact) j >= 128.

 /// 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 FlatAffineConstraints &cst, unsigned pos,
                              MLIRContext *context,
                              SmallVectorImpl<AffineExpr> &exprs) {
   assert(pos < cst.getNumIds() && "invalid position");
 
-  SmallVector<unsigned, 4> lbIndices, ubIndices;
-  cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
+  // Get upper-lower bound pair for this variable.
+  SmallVector<bool, 8> foundRepr(cst.getNumIds(), false);
+  for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i)
+    if (exprs[i])
+      foundRepr[i] = true;
+
+  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos);
+
+  // No upper-lower bound pair found for this var.
+  if (!ulPair)
+    return false;
+
+  unsigned ubPos = ulPair->first;
+
+  // Upper bound is of the form:
+  //      -divisor * id + expr >= 0
+  // where `id` is equivalent to `expr floordiv divisor`.
+  //
+  // Since the division cannot be dependent on itself, the coefficient of
+  // of `id` in `expr` is zero. The coefficient of `id` in the upperbound
+  // is -divisor.
+  int64_t divisor = -cst.atIneq(ubPos, pos);
+  int64_t constantTerm = cst.atIneq(ubPos, cst.getNumCols() - 1);
 
-  // Check if any lower bound, upper bound pair is of the form:
-  // divisor * id >=  expr - (divisor - 1)    <-- Lower bound for 'id'
-  // divisor * id <=  expr                    <-- Upper bound for 'id'
-  // Then, 'id' is equivalent to 'expr floordiv divisor'.  (where divisor > 1).
-  //
-  // For example:
-  //    32*k >= 16*i + j - 31                 <-- Lower bound for 'k'
-  //    32*k  <= 16*i + j                     <-- Upper bound for 'k'
-  //    expr = 16*i + j, divisor = 32
-  //    k = ( 16*i + j ) floordiv 32
-  //
-  //    4q >= i + j - 2                       <-- Lower bound for 'q'
-  //    4q <= i + j + 1                       <-- Upper bound for 'q'
-  //    expr = i + j + 1, divisor = 4
-  //    q = (i + j + 1) floordiv 4
-  for (auto ubPos : ubIndices) {
-    for (auto lbPos : lbIndices) {
-      // Due to the form of the inequalities, the sum of constants of upper
-      // bound and lower bound is divisor - 1. The 'divisor' here is
-      // cst.atIneq(lbPos, pos) and we already know that it's positive (since
-      // cst.Ineq(lbPos, ...) is a lower bound expr for 'pos'.
-      // Check if this sum of constants is divisor - 1.
-      int64_t divisor = cst.atIneq(lbPos, pos);
-      int64_t constantSum = cst.atIneq(lbPos, cst.getNumCols() - 1) +
-                            cst.atIneq(ubPos, cst.getNumCols() - 1);
-      if (constantSum != divisor - 1)
-        continue;
-      // For the remaining part, check if the lower bound expr's coeff's are
-      // negations of corresponding upper bound ones'.
-      unsigned c, f;
-      for (c = 0, f = cst.getNumCols() - 1; c < f; ++c)
-        if (cst.atIneq(lbPos, c) != -cst.atIneq(ubPos, c))
-          break;
-      // Lb coeff's aren't negative of ub coeff's (for the non constant term
-      // part).
-      if (c < f)
-        continue;
-      // Due to the form of the upper bound inequality, the constant term of
-      // `expr` is the constant term of upper bound inequality.
-      int64_t divConstantTerm = cst.atIneq(ubPos, cst.getNumCols() - 1);
-      // Construct the dividend expression.
-      auto dividendExpr = getAffineConstantExpr(divConstantTerm, context);
-      for (c = 0, f = cst.getNumCols() - 1; c < f; ++c) {
-        if (c == pos)
-          continue;
-        int64_t ubVal = cst.atIneq(ubPos, c);
-        if (ubVal == 0)
-          continue;
-        if (!exprs[c])
-          break;
-        dividendExpr = dividendExpr + ubVal * exprs[c];
-      }
-      // Expression can't be constructed as it depends on a yet unknown
-      // identifier.
-      // TODO: Visit/compute the identifiers in an order so that this doesn't
-      // happen. More complex but much more efficient.
-      if (c < f)
-        continue;
-      // Successfully detected the floordiv.
-      exprs[pos] = dividendExpr.floorDiv(divisor);
-      return true;
-    }
-  }
-  return false;
-}
+  auto dividendExpr = getAffineConstantExpr(constantTerm, context);
+  unsigned c, f;
+  for (c = 0, f = cst.getNumCols() - 1; c < f; c++) {
+    // computeSingleVarRepr guarantees that expr is known here.
+
 
 // Fills an inequality row with the value 'val'.
 static inline void fillInequality(FlatAffineConstraints *cst, unsigned r,
                                   int64_t val) {
   for (unsigned c = 0, f = cst->getNumCols(); c < f; c++) {
     cst->atIneq(r, c) = val;
   }
 }

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, clearConstraints) {
 
   fac.clearConstraints();
 
   fac.addInequality({1, 0});
   EXPECT_EQ(fac.atIneq(0, 0), 1);
   EXPECT_EQ(fac.atIneq(0, 1), 0);
 }
 
-TEST(FlatAffineConstraintsTest, constantDivs) {
-  // This test checks if floordivs with numerator containing non zero constant
-  // term can be computed from a FlatAffineConstraints instance.
+/// Check if the expected division representation of local variables matches the
+/// computed representation. The expected division representation is given as
+/// a vector of expressions set in `divisions` and the corressponding
+/// denominator in `denoms`. If expected denominator for a variable is
+/// non-positive, the local variable is expected to not have a computed
+/// representation.
+static void checkDivisionRepresentation(
+    FlatAffineConstraints &fac,
+    const std::vector<SmallVector<int64_t, 8>> &divisions,
+    const SmallVector<int64_t, 8> &denoms) {
+
+  assert(divisions.size() == fac.getNumLocalIds() &&
+         "Size of expected divisions does not match number of local variables");
+  assert(
+      denoms.size() == fac.getNumLocalIds() &&
+      "Size of expected denominators does not match number of local variables");
+
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> res(
+      fac.getNumLocalIds(), llvm::None);
+  fac.getLocalReprLbUbPairs(res);
+
+  // Check if all expected divisions are computed.
+  for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i)
+    if (denoms[i] > 0)
+      EXPECT_TRUE(res[i].hasValue());
+    else
+      EXPECT_FALSE(res[i].hasValue());
+
+  unsigned divOffset = fac.getNumDimAndSymbolIds();
+  for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i) {
+    if (!res[i])
+      continue;
+
+    // Check if the bounds are of the form:
+    //      0 <= expr - divisor * id <= divisor - 1
+    // Rearranging, we have:
+    //       divisor * id - expr + (divisor - 1) >= 0  <-- Lower bound for 'id'
+    //      -divisor * id + expr                 >= 0  <-- Upper bound for 'id'
+    // where `id = expr floordiv divisor`.
+    unsigned ubPos = res[i]->first, lbPos = res[i]->second;
+    const SmallVector<int64_t, 8> &expr = divisions[i];
+
+    // Check if lower bound is of the correct form.
+    int64_t computedDivisorLb = fac.atIneq(lbPos, i + divOffset);
+    EXPECT_EQ(computedDivisorLb, denoms[i]);
+    for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) {
+      if (c == i + divOffset)
+        continue;
+      EXPECT_EQ(fac.atIneq(lbPos, c), -expr[c]);
+    }
+    // Check if constant term of lower bound matches expected constant term.
+    EXPECT_EQ(fac.atIneq(lbPos, fac.getNumCols() - 1),
+              -expr.back() + (denoms[i] - 1));
+
+    // Check if upper bound is of the correct form.
+    int64_t computedDivisorUb = fac.atIneq(ubPos, i + divOffset);
+    EXPECT_EQ(computedDivisorUb, -denoms[i]);
+    for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) {
+      if (c == i + divOffset)
+        continue;
+      EXPECT_EQ(fac.atIneq(ubPos, c), expr[c]);
+    }
+    // Check if constant term of upper bound matches expected constant term.
+    EXPECT_EQ(fac.atIneq(ubPos, fac.getNumCols() - 1), expr.back());
+  }
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprSimple) {
+  FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
+
+  fac.addLocalFloorDiv({1, 4}, 10);
+  fac.addLocalFloorDiv({1, 0, 100}, 10);
+
+  std::vector<SmallVector<int64_t, 8>> divisions = {{1, 0, 0, 4},
+                                                    {1, 0, 0, 100}};
+  SmallVector<int64_t, 8> denoms = {10, 10};
+
+  // Check if floordivs can be computed when no other inequalities exist
+  // and floor divs do not depend on each other.
+  checkDivisionRepresentation(fac, divisions, denoms);
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprConstantFloorDiv) {
+  FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {});
+
+  fac.addInequality({1, 0, 3, 1, 2});
+  fac.addInequality({1, 2, -8, 1, 10});
+  fac.addEquality({1, 2, -4, 1, 10});
+
+  fac.addLocalFloorDiv({0, 0, 0, 0, 10}, 30);
+  fac.addLocalFloorDiv({0, 0, 0, 0, 0, 99}, 101);
+
+  std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0, 0, 0, 0, 10},
+                                                    {0, 0, 0, 0, 0, 0, 99}};
+  SmallVector<int64_t, 8> denoms = {30, 101};
+
+  // Check if floordivs with constant numerator can be computed.
+  checkDivisionRepresentation(fac, divisions, denoms);
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprRecursive) {
   FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {});
+  fac.addInequality({1, 0, 3, 1, 2});
+  fac.addInequality({1, 2, -8, 1, 10});
+  fac.addEquality({1, 2, -4, 1, 10});
+
+  fac.addLocalFloorDiv({0, -2, 7, 2, 10}, 3);
+  fac.addLocalFloorDiv({3, 0, 9, 2, 2, 10}, 5);
+  fac.addLocalFloorDiv({0, 1, -123, 2, 0, -4, 10}, 3);
+
+  fac.addInequality({1, 2, -2, 1, -5, 0, 6, 100});
+  fac.addInequality({1, 2, -8, 1, 3, 7, 0, -9});
+
+  std::vector<SmallVector<int64_t, 8>> divisions = {{0, -2, 7, 2, 0, 0, 0, 10},
+                                                    {3, 0, 9, 2, 2, 0, 0, 10},
+                                                    {0, 1, -123, 2, 0, -4, 10}};
+  SmallVector<int64_t, 8> denoms = {3, 5, 3};
 
-  // Build a FlatAffineConstraints instance with floordivs containing numerator
-  // with non zero constant term.
-  fac.addLocalFloorDiv({0, 1, 0, 0, 10}, 30);
-  fac.addLocalFloorDiv({1, 0, 0, 0, 0, 99}, 101);
-
-  // Add inequalities using the local variables created above.
-  fac.addInequality({1, 0, 0, 0, 1, 0, 2});
-  fac.addInequality({1, 0, 0, 0, 0, 1, 5});
-
-  // FlatAffineConstraints::getAsIntegerSet returns a null integer set if an
-  // explicit representation for each local variable could not be found.
-  MLIRContext ctx;
-  IntegerSet iSet = fac.getAsIntegerSet(&ctx);
-  EXPECT_TRUE((bool)iSet);
+  // Check if floordivs which may depend on other floordivs can be computed.
+  checkDivisionRepresentation(fac, divisions, denoms);
 }
 
 } // namespace mlir