[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

1260

Add an assert for the length of exprs?

1266–1267

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'

1273–1274

In general, please avoid using auto unless it aids readability

1277–1286

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

1284

nit: Please put the full stop

1288

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

1299

nit: Please end the sentence with a full stop

1315

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

1331

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

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1281

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)

1294–1295

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

1314

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

1344

In general, please avoid using auto.

1344–1348

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

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

1627

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

1642–1659

I think you mean "the division variable"

1646

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

mlir/unittests/Analysis/AffineStructuresTest.cpp

596

nit: seems a bit odd to use a reference here

606–607

Should it be this?

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

621

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

624–656

Please add some comments saying what the tests are checking.

[Groverkss]

Thank you for the review!

mlir/lib/Analysis/AffineStructures.cpp

1315

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

mlir/unittests/Analysis/AffineStructuresTest.cpp

621

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

1315

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

mlir/unittests/Analysis/AffineStructuresTest.cpp

606–607

expr + divisor - 1, not minus?

621

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

1641–1642

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

1315

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

1344

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

1268

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.

1639–1641

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

1657

nit: gurantees -> guarantees

mlir/unittests/Analysis/AffineStructuresTest.cpp

612–614

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

622–623

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

168

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

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

702

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

1639–1641

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

163

indices may be confusing here. Consider rephrasing

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

?

166

Please rephrase the last two lines for clarity.

167

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

167

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

1255

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

1261

Assertion message please.

1341

SmallVector here to avoid a heap allocation.

1626

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

mlir/unittests/Analysis/AffineStructuresTest.cpp

595

This should be marked static?

[arjunp]

mlir/include/mlir/Analysis/AffineStructures.h

167

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

167

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;
 
+  /// Find pairs of indices of inequalities, using which an explicit
+  /// representation for each local variable can be computed.
+  /// Pairs are returned as indices of upperbound, lowerbound equalities.
+  /// If no such pair is found, pair is marked as None.
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>>
+  computeLocalRepr() 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); }
 
   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 std::vector<bool> &foundRepr, unsigned pos) {
+  assert(pos < cst.getNumIds() && "invalid position");
+  assert(foundRepr.size() == cst.getNumIds());
+
+  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 <= divisor * id - expr <= divisor - 1
+  // Rearranging, we have:
+  //                  divisor * id - expr >= 0    <-- Lower bound for 'id'
+  // -divisor * id + expr + (divisor - 1) >= 0    <-- Upper bound for 'id'
+  //
+  // For example, if -32*k + 16*i + j >= 0
+  //                  32*k - 16*i - j + 31 >= 0   <=>
+  //             k = ( 16*i + j ) floordiv 32
+  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 llvm::Optional<std::pair<unsigned, unsigned>>({ubPos, lbPos});
+    }
+  }
+
+  return llvm::None;
+}
+
+/// Find pairs of indices of inequalities, using which an explicit
+/// representation for each local variable can be computed.
+/// Pairs are returned as indices of upperbound, lowerbound equalities.
+/// If no such pair is found, pair is marked as None.
+std::vector<llvm::Optional<std::pair<unsigned, unsigned>>>
+FlatAffineConstraints::computeLocalRepr() const {
+
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
+      getNumLocalIds(), llvm::None);
+  std::vector<bool> 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);
+
+  return repr;
+}
+
 /// 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
+  std::vector<bool> 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, lbPos = ulPair->second;
 
-  // 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).
+  // Lower bound is of the form:
+  //      divisor * id >=  expr
   //
-  // For example, if -32*k + 16*i + j >= 0
-  //                  32*k - 16*i - j + 31 >= 0   <=>
-  //             k = ( 16*i + j ) floordiv 32
-  unsigned seenDividends = 0;
-  for (auto ubPos : ubIndices) {
-    for (auto lbPos : lbIndices) {
-      // Check if the lower bound's constant term 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'.
-      int64_t divisor = cst.atIneq(lbPos, pos);
-      int64_t lbConstTerm = cst.atIneq(lbPos, cst.getNumCols() - 1);
-      if (lbConstTerm != divisor - 1)
-        continue;
-      // Check if upper bound's constant term is 0.
-      if (cst.atIneq(ubPos, cst.getNumCols() - 1) != 0)
-        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;
-        if (c != pos && cst.atIneq(lbPos, c) != 0)
-          seenDividends++;
-      }
-      // Lb coeff's aren't negative of ub coeff's (for the non constant term
-      // part).
-      if (c < f)
-        continue;
-      if (seenDividends >= 1) {
-        // Construct the dividend expression.
-        auto dividendExpr = getAffineConstantExpr(0, context);
-        unsigned c, f;
-        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;
-}
+  // The divisor is the coefficent of the division variable
+  // the constant term for the division is:
+  //      constantTerm = constantTermOf(expr) - (divisor - 1)
+  int64_t constantTerm =
+      cst.atIneq(lbPos, cst.getNumCols() - 1) - (divisor - 1);
+
+  auto dividendExpr = getAffineConstantExpr(constantTerm, context);
+    // computeSingleVarRepr gurantees 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);
 }
 
+void checkDivisionRepresentation(FlatAffineConstraints &fac, bool all = true) {
+  // Check if representation can be computed for all local variables
+  auto res = fac.computeLocalRepr();
+
+  // If all is set, check if all divisions are computed
+  if (all) {
+    for (auto r : res)
+      EXPECT_NE(r, llvm::None);
+  }
+
+  unsigned divOffset = fac.getNumDimAndSymbolIds();
+  for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i) {
+    if (!res[i])
+      continue;
+
+    // Check if bounds are of the form:
+    // divisor * id >=  expr                    <-- Lower bound for 'id'
+    // divisor * id <=  expr - (divisor - 1)    <-- Upper bound for 'id'
+    unsigned ubPos = res[i]->first, lbPos = res[i]->second;
+    int64_t constantSum = fac.atIneq(lbPos, fac.getNumCols() - 1) +
+                          fac.atIneq(ubPos, fac.getNumCols() - 1);
+
+    // Check if (sum of constants > 0). Indirectly checking (divisor > 1).
+    EXPECT_TRUE(constantSum > 0);
+
+    // Check if coeff of variable is equal to divisor
+    EXPECT_EQ(constantSum + 1, fac.atIneq(lbPos, i + divOffset));
+
+    // Check if constraints are opposite of each other.
+    for (unsigned c = 0, f = fac.getNumIds(); c < f; ++c)
+      EXPECT_EQ(fac.atIneq(ubPos, c), -fac.atIneq(lbPos, c));
+  }
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprSimple) {
+  FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
+
+  fac.addLocalFloorDiv({1, 4}, 10);
+  fac.addLocalFloorDiv({1, 0, 100}, 10);
+
+  // Check if floordivs can be recovered when no other inequalities exist
+  // and floor divs do not depend on each other.
+  checkDivisionRepresentation(fac);
+}
+
+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);
+
+  // Check if floordivs with constant numerator can be recovered.
+  checkDivisionRepresentation(fac);
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprRecursive) {
+  FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {});
+  fac.addInequality({1, 0, 3, 1, 2});
+  fac.addInequality({1, 2, -8, 1, 10});
+  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});
+
+  // Check if floordivs which may depend on other floordivs can be recovered.
+  checkDivisionRepresentation(fac);
+}
+
 } // namespace mlir