[MLIR] FlatAffineConstraints: Fixed bug where some divisions were not being detected

[Groverkss]

Summary

This patch fixes a bug in the existing implementation of detectAsFloorDiv,
where floordivs with numerator with non-zero constant term and floordivs with
numerator only consisting of a constant term were not being detected.

[Groverkss]

• Added test for bugfix

[arjunp]

Thanks for the patch! Please fix the clang-format errors. Other than that, this seems good to me. Let's wait for one of the others to take a look as well.

mlir/lib/Analysis/AffineStructures.cpp

1556

This now also allows numerators which are just zero, right? I think it makes sense to support that (and it will be needed to support subtraction on sets with divisions that happen to have numerator zero), but maybe it's worth mentioning in the commit message.

[Groverkss]

• Fixed clang-tidy errors and add punctuation in some comments

[Groverkss]

Addressed comments.

[vinayaka-polymage]

Thanks for this enhancement!

1. I tried an example, but didn't get the expected result. Can you please clarify ?
2. Format the commit msg to wrap around.
3. If my understanding of (1) is correct, can you also extend the doc comment of the function to use an example that has non-zero constant upper bound ?

mlir/lib/Analysis/AffineStructures.cpp

1561

This would be incorrect, I think. You would need to still initialize the expr with zero. I tried the example:
4q - 1 <= i + j <= 4q + 2.

Finally, the dividend should be i + j + 1, but I think it would detect i + j + 2 now.

On a side note, it means, testing is insufficient as it is only checking if all ids are detected, but not verifying the results.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1561

The detected dividend will be i + j + 1.

The upper-bound inequality is:

4q <= i + j + 1

rearranging it we get:

i + j - 4q + 1 >= 0

From which, the constant term extracted is 1. Is there something I'm missing?

[vinayaka-polymage]

Thanks for the clarification, will take another look in some time.

mlir/lib/Analysis/AffineStructures.cpp

1561

I just noticed that the loop below runs 0 to cst.getNumCols() - 1; this means the constant term is not added once again. So, you are correct.

I was under the impression that the loop is 0 to cst.getNumCols(), and the constant 1 in the above case would get added once again. So, this works too.

[Groverkss]

• Updated exaxmples in detectAsFloorDiv docs

[Groverkss]

Addressed comments.

mlir/lib/Analysis/AffineStructures.cpp

1561

Regarding the verification in testing: due to how the current implementation is exposed, it is hard to test if the result matches. There is another patch (https://reviews.llvm.org/D106662) that refactors this implementation and provides better a better way to test.

[vinayaka-polymage]

This comment was removed by vinayaka-polymage.

[vinayaka-polymage]

Looks good now, thanks for addressing the comments.

mlir/lib/Analysis/AffineStructures.cpp

static bool detectAsFloorDiv(const FlatAffineConstraints &cst, unsigned pos,
   SmallVector<unsigned, 4> lbIndices, ubIndices;
   cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
 
   // 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, if -32*k + 16*i + j >= 0
-  //                  32*k - 16*i - j + 31 >= 0   <=>
-  //             k = ( 16*i + j ) floordiv 32
-  unsigned seenDividends = 0;
+  // 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) {
-      // 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'.
+      // 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 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)
+      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++) {
+      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;
-      }
-    }
-  }
+      // 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);
+      auto dividendExpr = getAffineConstantExpr(divConstantTerm, context);
+      for (c = 0, f = cst.getNumCols() - 1; c < f; ++c) {
   return false;
 }
 
 // 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

 //===- AffineStructuresTest.cpp - Tests for AffineStructures ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/MLIRContext.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
 
 namespace mlir {
 
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.
+  FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {});
+
+  // 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);
+}
+
 } // namespace mlir