[MLIR][FlatAffineConstraints] Add support for extracting divisions with tighter bounds

[Groverkss]

Summary

This patch adds support for extracting divisions when the set contains bounds
which are tighter than the division bounds. For example:

 3q - i + 2 >= 0                       <-- Lower bound for 'q'
-3q + i - 1 >= 0                       <-- Tighter upper bound for 'q'

Here, the actual upper bound for division for q would be -3q + i >= 0, but
since this actual upper bound is implied by a tighter upper bound, which awe can still
extract the divison.

[arjunp]

Thanks for the patch!

What we are really supporting is tighter bounds that imply the standard division bounds right? I would change the title to reflect that. Also, the first line of the summary is not clear to me.

While extracting divisions from local variables, due to additional constraints,
the original upper bound of the floor division may be removed.

Is it really removed while extracting divisions or before?

mlir/lib/Analysis/AffineStructures.cpp

1241–1255

I feel that this doesn't explain why it is valid to check this condition, i.e., why does this imply that we have a division? I would start by saying that we also support detecting divs from bounds that are strictly tighter than the division bounds described above still define a division representation, since the tighter bounds imply the division bounds.

Then finally I would give the example and mention that this can occur for example, because of the tighter bounds making the division bounds redundant, thus causing them to be removed.

1254–1256

"relaxed lower bound" seems misleading, because the lower bound here is either original one or actually tighter. (Even though the condition is relaxed in that we allow both the original bound and tighter bounds.)

1254–1256

This section was pretty confusing at first for me personally. I didn't immediately realize that you meant that Sum of constants >= 0 is equivalent to divisior - 1 - c >= 0, and so on. Writing iff on the left of the second and third lines for each constraint would help.

1255

Also, another thing I was wondering when I first saw this is, why can't we also support a tighter upper bound, i.e., c <= expr - divisior * id <= d.
Of course the answer is that WLOG you can set d = divisior - 1 by adding divisior - 1 - d to all three sides of the equation, increasing c and the constant term of expr. I would probably write something to this effect in the documentation:

WLOG we consider the upper bound to be `divisor - 1`, by defining that if the constant term of the upper bound inequality is `p`, then `divisor - 1 - p` is the constant term in `expr` and hence that inequality gives an upper bound of `divisor - 1` on `expr - divisor * id`.

Also, if you explain it this way, you can also directly obtain the constant term of expr at the start as divisor - 1 - p and then obtain c by subtracting the constant term of expression from the lower bound. Then you can check the bounds on c directly, which might be cleaner than explaining why we are checking bounds on the sum of constants below. Anyway I leave it to you to decide which is easier to understand.

1257

1257

Please set divisor before this and compare with divisor - 1. Then it will be immediately clear what we are doing here.

IMO if (!(constantSum <= divisor - 1)) is much more readable

1257

IMO if(!(constantSum >= 0)) is better

1291

Maybe not for this patch, but you can std::move here.

mlir/unittests/Analysis/AffineStructuresTest.cpp

677–729

Please add a test case like 0 <= x - 3q <= 3, where no division representation should be found.

682–683

Can you explain the test case a bit more? Where is the inequality that makes it redundant?

[Groverkss]

• Simplify implementation
• Update unit tests

mlir/lib/Analysis/AffineStructures.cpp

1255

I don't find this easier to understand :( I tried to simplify the doc as best as I could.

[arjunp]

LGTM, minor documentation comments below.

mlir/lib/Analysis/AffineStructures.cpp

1228–1229

We may as well use the same div as above (with tighter constraints) and say that these constraints also imply the same div. That makes it easier to read

1257–1258

1267–1283

1268

1269

1274

1279

1287–1291

mlir/unittests/Analysis/AffineStructuresTest.cpp

680

703

709

[Groverkss]

• Address Arjun's comments

mlir/lib/Analysis/AffineStructures.cpp

 /// Let `id` be the pos^th identifier, then `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
+///     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
+///     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
+//
+/// This function also supports detecting divisions from bounds that are
+/// strictly tighter than the division bounds described above, since tighter
+/// bounds imply the division bounds. For example:
+///     4q - i - j + 2 >= 0                       <-- Lower bound for 'q'
+///    -4q + i + j     >= 0                       <-- Tight upper bound for 'q'
+///
+/// To extract floor divisions with tighter bounds, we assume that that the
+/// constraints are of the form:
+///     c <= expr - divisior * id <= divisor - 1, where 0 <= c <= divisor - 1
+/// Rearranging, we have:
+///     divisor * id - expr + (divisor - 1) >= 0  <-- Lower bound for 'id'
+///    -divisor * id + expr - c             >= 0  <-- Upper bound for 'id'
 ///
 /// If successful, `expr` is set to dividend of the division and `divisor` is
 /// set to the denominator of the division.
 static LogicalResult getDivRepr(const FlatAffineConstraints &cst, unsigned pos,
                                 unsigned ubIneq, unsigned lbIneq,
                                 SmallVector<int64_t, 8> &expr,
                                 unsigned &divisor) {
 
   assert(pos <= cst.getNumIds() && "Invalid identifier position");
   assert(ubIneq <= cst.getNumInequalities() &&
          "Invalid upper bound inequality position");
   assert(lbIneq <= cst.getNumInequalities() &&
          "Invalid upper bound inequality position");
 
-  // Due to the form of the inequalities, sum of constants of the
-  // inequalities is (divisor - 1).
-  int64_t denominator = cst.atIneq(lbIneq, cst.getNumCols() - 1) +
-                        cst.atIneq(ubIneq, cst.getNumCols() - 1) + 1;
+  // Extract divisor from the lower bound.
+  divisor = cst.atIneq(lbIneq, pos);
 
-  // Divisor should be positive.
-  if (denominator <= 0)
-    return failure();
-
-  // Check if coeff of variable is equal to divisor.
-  if (denominator != cst.atIneq(lbIneq, pos))
-    return failure();
-
-  // Check if constraints are opposite of each other. Constant term
-  // is not required to be opposite and is not checked.
+  // First, check if the constraints are opposite of each other except the
+  // constant term.
   unsigned i = 0, e = 0;
   for (i = 0, e = cst.getNumIds(); i < e; ++i)
     if (cst.atIneq(ubIneq, i) != -cst.atIneq(lbIneq, i))
       break;
 
   if (i < e)
     return failure();
 
-  // Set expr with dividend of the division.
-  SmallVector<int64_t, 8> dividend(cst.getNumCols());
-  for (i = 0, e = cst.getNumCols(); i < e; ++i)
+  // Then, check if the constant term is of the proper form.
+  // Due to the form of the upper/lower bound inequalities, the sum of their
+  // constants is `divisor - 1 - c`. From this, we can extract c:
+  int64_t constantSum = cst.atIneq(lbIneq, cst.getNumCols() - 1) +
+                        cst.atIneq(ubIneq, cst.getNumCols() - 1);
+  int64_t c = divisor - 1 - constantSum;
+
+  // Check if `c` satisfies the condition `0 <= c <= divisor - 1`. This also
+  // implictly checks that `divisor` is positive.
+  if (!(c >= 0 && c <= divisor - 1))
+    return failure();
+
+  // The inequality pair can be used to extract the division.
+  // Set `expr` to the dividend of the division except the constant term, which
+  // is set below.
+  expr.resize(cst.getNumCols(), 0);
+  for (i = 0, e = cst.getNumIds(); i < e; ++i)
     if (i != pos)
-      dividend[i] = cst.atIneq(ubIneq, i);
-  expr = dividend;
+      expr[i] = cst.atIneq(ubIneq, i);
 
-  // Set divisor.
-  divisor = denominator;
+  // From the upper bound inequality's form, its constant term is equal to the
+  // constant term of `expr`, minus `c`. From this,
+  // constant term of `expr` = constant term of upper bound + `c`.
+  expr.back() = cst.atIneq(ubIneq, cst.getNumCols() - 1) + c;
 
   return success();
 }
 
 /// 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.
 /// Returns the upper and lower bound inequalities using which the floordiv can

mlir/unittests/Analysis/AffineStructuresTest.cpp

static void checkDivisionRepresentation(
     const std::vector<SmallVector<int64_t, 8>> &expectedDividends,
     const SmallVectorImpl<unsigned> &expectedDenominators) {
 
   std::vector<SmallVector<int64_t, 8>> dividends;
   SmallVector<unsigned, 4> denominators;
 
   fac.getLocalReprs(dividends, denominators);
 
-  // Check that the `dividends` and `expectedDividends` match.
-  EXPECT_TRUE(expectedDividends == dividends);
-
   // Check that the `denominators` and `expectedDenominators` match.
   EXPECT_TRUE(expectedDenominators == denominators);
+
+  // Check that the `dividends` and `expectedDividends` match. If the
+  // denominator for a division is zero, we ignore its dividend.
+  EXPECT_TRUE(dividends.size() == expectedDividends.size());
+  for (unsigned i = 0, e = dividends.size(); i < e; ++i)
+    if (denominators[i] != 0)
+      EXPECT_TRUE(expectedDividends[i] == dividends[i]);
 }
 
 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 = {
       {3, 0, 9, 2, 2, 0, 0, 10},
       {0, 1, -123, 2, 0, -4, 0, 10}};
 
   SmallVector<unsigned, 8> denoms = {3, 5, 3};
 
   // Check if floordivs which may depend on other floordivs can be computed.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
+TEST(FlatAffineConstraintsTest, computeLocalReprTightUpperBound) {
+  MLIRContext context;
+
+  {
+    FlatAffineConstraints fac = parseFAC("(i) : (i mod 3 - 1 >= 0)", &context);
+
+    // The set formed by the fac is:
+    //        3q - i + 2 >= 0             <-- Division lower bound
+    //       -3q + i - 1 >= 0
+    //       -3q + i     >= 0             <-- Division upper bound
+    // We remove redundant constraints to get the set:
+    //        3q - i + 2 >= 0             <-- Division lower bound
+    //       -3q + i - 1 >= 0             <-- Tighter division upper bound
+    // thus, making the upper bound tighter.
+    fac.removeRedundantConstraints();
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{1, 0, 0}};
+    SmallVector<unsigned, 8> denoms = {3};
+
+    // Check if the divisions can be computed even with a tighter upper bound.
+    checkDivisionRepresentation(fac, divisions, denoms);
+  }
+
+  {
+    FlatAffineConstraints fac = parseFAC(
+        "(i, j, q) : (4*q - i - j + 2 >= 0, -4*q + i + j >= 0)", &context);
+    // Convert `q` to a local variable.
+    fac.convertDimToLocal(2, 3);
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, 1}};
+    SmallVector<unsigned, 8> denoms = {4};
+
+    // Check if the divisions can be computed even with a tighter upper bound.
+    checkDivisionRepresentation(fac, divisions, denoms);
+  }
+}
+
+TEST(FlatAffineConstraintsTest, computeLocalReprNoRepr) {
+  MLIRContext context;
+  FlatAffineConstraints fac =
+      parseFAC("(x, q) : (x - 3 * q >= 0, -x + 3 * q + 3 >= 0)", &context);
+  // Convert q to a local variable.
+  fac.convertDimToLocal(1, 2);
+
+  std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0}};
+  SmallVector<unsigned, 8> denoms = {0};
+
+  // Check that no division is computed.
+  checkDivisionRepresentation(fac, divisions, denoms);
+}
+
 TEST(FlatAffineConstraintsTest, simplifyLocalsTest) {
   // (x) : (exists y: 2x + y = 1 and y = 2).
   FlatAffineConstraints fac(1, 0, 1);
   fac.addEquality({2, 1, -1});
   fac.addEquality({0, 1, -2});
 
   EXPECT_TRUE(fac.isEmpty());
 
   // (x) : (exists y, z, w: 3x + y = 1 and 2y = z and 3y = w and z = w).