[MLIR] Add division normalization by GCD in `getDivRepr` fn.

[pashu123]

Summary

This commits adds division normalization in the getDivRepr function which extracts
the gcd from the dividend and divisor and normalizes them.

Signed-off-by: Prashant Kumar <[e-mail]>

[pashu123]

Clang format.

[pashu123]

Formatting.

[pashu123]

Updating for loops.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1969–1971

Please mention that by normalization we mean GCD normalization. Maybe even change the function name to reflect that.

1983–1987

You can ignore the constant term while taking gcd. The reasoning is as follows:

For a division:

floor((a + m f(x))/(m d))

you can replace it by

floor((floor(a/m) + f(x))/d)

Because, the difference {a/m}/d in the argument satisfies 0 <= {a/m}/d < 1/d and can therefore not influence the result of the floor.

1991

Unnecessary newline.

2001–2002

I think a better place for normalization is in getDivRepr after you obtain the division.

[pashu123]

Addressing groverkss comments and clang-format.

[pashu123]

Removing unneccesarry newline.

[pashu123]

Comments typo.

[pashu123]

mlir/lib/Analysis/AffineStructures.cpp

2001–2002

I see, thanks.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1355–1356

1355–1357

1358

I think a better name would be normalizeDivisionByGCD.

1358

Could you rename expr to dividend for more clarity?

1366

1366–1367

Please write the reasoning of why it can be replaced here.

1404–1406

mlir/unittests/Analysis/AffineStructuresTest.cpp

817

Just mention here that this tests normalization.

902

Please mention that this division should be normalized. Something like z = [3x + 3y] / 9 -> z = [x+y]/3.

[pashu123]

Addressing groverkss comments.

[pashu123]

Updated comments.

[pashu123]

nit

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1358

1367–1373

[bondhugula]

The commit summary and title is missing context on where the normalization is being added. Please update.

mlir/lib/Analysis/AffineStructures.cpp

1358

It'd be good to change denominator -> divisor to be consistent with dividend.

1360–1365

This whole thing doesn't look clean to me. Calling something gcd when the dividend is empty is out of line.

We probably need a logic cleanup or better variable names.

1363

Please use !empty() instead of size().

1404–1405

"returned division" isn't meaningful. Rephrase?

[pashu123]

Addressing Uday's comments.

[pashu123]

Clang format.

[pashu123]

mlir/lib/Analysis/AffineStructures.cpp

1360–1365

I have placed the dividend.empty() check in the if condition above.

1404–1405

How about final division expression?

[pashu123]

Clang format warnings.

[Groverkss]

LGTM. But, I would wait for @bondhugula to accept this in case he has any remaining concerns.

[pashu123]

@bondhugula I have addressed all your comments. Please take a look.

[pashu123]

In D115595#3204258, @pashu123 wrote:

@bondhugula I have addressed all your comments. Please take a look.

Reminder.

[pashu123]

Rebasing latest changes.

[bondhugula]

mlir/lib/Analysis/Presburger/Utils.cpp

44 ↗

(On Diff #397478)

Capture by value is sufficient.

mlir/unittests/Analysis/AffineStructuresTest.cpp

694–697

Why is this test case updated to a trivial one? Please try to preserve the original test's strength here. @Groverkss?

817

Looks like this comment is still unaddressed.

836

1 division <-- One division ...

[pashu123]

Addressing bondhugula's comment.

[pashu123]

Comments.

[pashu123]

Comments.

[Groverkss]

mlir/unittests/Analysis/AffineStructuresTest.cpp

694–697

After the changes this patch introduces, constant divisions are simplified to an integer. Here it becomes ([0/1], [0/1]) instead of ([10/30], [99/101]). This test was supposed to check if we can extract constant divisions.

This test could be changed to adding divisions to ([100/30], [206/101]) and the result should be ([3/1], [2/1]). Just to check if constants are still extracted properly (extracting 0 does not seem like the best test to me).

[pashu123]

updating tests.

[pashu123]

mlir/unittests/Analysis/AffineStructuresTest.cpp

694–697

I have updated the test case.

[pashu123]

@Groverkss @bondhugula I have addressed all your comments PTAL.

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;
 }
 
+/// Normalize a division's `dividend` and the `denominator` by their GCD. For
+/// example: if the dividend and denominator are [2,0,4] and 4 respectively,
+/// they get normalized to [1,0,2] and 2.
+static void normalizeDivisionsByGCD(SmallVectorImpl<int64_t> &dividend,
+                                    unsigned &denominator) {
+  if (denominator == 0)
+    return;
+  int64_t gcd =
+      dividend.size() > 0
+          ? llvm::greatestCommonDivisor(dividend[0], int64_t(denominator))
+          : 1;
+
+  // The constant term is ignored while taking gcd since
+  // floor((a + m.f(x))/(m.d)) can be replaced by floor((floor(a/m) + f(x))/d).
+  for (size_t i = 1, m = dividend.size() - 1; i < m; i++) {
+    gcd = llvm::greatestCommonDivisor(dividend[i], gcd);
+    if (gcd == 1)
+      return;
+  }
+
+  // Normalize the expr and the denominator.
+  std::transform(dividend.begin(), dividend.end(), dividend.begin(),
+                 [&gcd](int64_t &n) { return floor(n / gcd); });
+  denominator /= gcd;
+}
+
 /// Check if the pos^th identifier can be represented as a division using upper
 /// bound inequality at position `ubIneq` and lower bound inequality at position
 /// `lbIneq`.
 ///
 /// 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
 ///
 ///       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
 ///
 /// If successful, `expr` is set to dividend of the division and `divisor` is
-/// set to the denominator of the division.
+/// set to the denominator of the division. The returned division is normalized
+/// by GCD.
 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() &&
static LogicalResult getDivRepr(const FlatAffineConstraints &cst, unsigned pos,
   SmallVector<int64_t, 8> dividend(cst.getNumCols());
   for (i = 0, e = cst.getNumCols(); i < e; ++i)
     if (i != pos)
       dividend[i] = cst.atIneq(ubIneq, i);
   expr = dividend;
 
   // Set divisor.
   divisor = denominator;
+  normalizeDivisionsByGCD(expr, divisor);
 
   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.
static void eliminateRedundantLocalId(FlatAffineConstraints &fac, unsigned pos1,
   pos2 += localOffset;
   for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
     fac.atIneq(i, pos1) += fac.atIneq(i, pos2);
   for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
     fac.atEq(i, pos1) += fac.atEq(i, pos2);
   fac.removeId(pos2);
 }
 
 /// Adds additional local ids to the sets such that they both have the union
 /// of the local ids in each set, without changing the set of points that
 /// lie in `this` and `other`.
 ///
 /// To detect local ids that always take the same in both sets, each local id is
 /// represented as a floordiv with constant denominator in terms of other ids.
 /// After extracting these divisions, local ids with the same division
 /// representation are considered duplicate and are merged. It is possible that
 /// division representation for some local id cannot be obtained, and thus these
 /// local ids are not considered for detecting duplicates.
 void FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
   assert(getNumDimIds() == other.getNumDimIds() &&
          "Number of dimension ids should match");
   assert(getNumSymbolIds() == other.getNumSymbolIds() &&
          "Number of symbol ids should match");
 
   FlatAffineConstraints &facA = *this;
   FlatAffineConstraints &facB = other;
 
   // Merge local ids of facA and facB without using division information,
   // i.e. append local ids of `facB` to `facA` and insert local ids of `facA`
   // to `facB` at start of its local ids.
   unsigned initLocals = facA.getNumLocalIds();
   insertLocalId(facA.getNumLocalIds(), facB.getNumLocalIds());
   facB.insertLocalId(0, initLocals);
 
   // Get division representations from each FAC.
   std::vector<SmallVector<int64_t, 8>> divsA, divsB;
   SmallVector<unsigned, 4> denomsA, denomsB;
   facA.getLocalReprs(divsA, denomsA);
   facB.getLocalReprs(divsB, denomsB);
 
   // Copy division information for facB into `divsA` and `denomsA`, so that
   // these have the combined division information of both FACs. Since newly
   // added local variables in facA and facB have no constraints, they will not
   // have any division representation.
   std::copy(divsB.begin() + initLocals, divsB.end(),
             divsA.begin() + initLocals);
   std::copy(denomsB.begin() + initLocals, denomsB.end(),
             denomsA.begin() + initLocals);
 
   // Find and merge duplicate divisions.

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, computeLocalReprConstantFloorDiv) {
 
   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<unsigned, 8> denoms = {30, 101};
+  std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0, 0, 0, 0, 0},
+                                                    {0, 0, 0, 0, 0, 0, 0}};
+  SmallVector<unsigned, 8> denoms = {1, 1};
 
   // 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});
TEST(FlatAffineConstraintsTest, mergeDivisionsSimple) {
 
     // Local space should be same.
     EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(fac1.getNumLocalIds(), 2u);
     EXPECT_EQ(fac2.getNumLocalIds(), 2u);
   }
+
+  {
+    // (x) : (exists z, y  = [x / 2] : x = 3y and x + z + 1 >= 0).
+    FlatAffineConstraints fac1(1, 0, 1);
+    // This division would be normalized.
+    fac1.addLocalFloorDiv({3, 0, 0}, 6); // y = [3x / 6] -> [x/2].
+    fac1.addEquality({1, 0, -3, 0});     // x = 3z.
+    fac1.addInequality({1, 1, 0, 1});    // x + y + 1 >= 0.
+
+    // (x) : (exists y = [x / 2], z : x = 5y).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
+    fac2.addEquality({1, -5, 0});     // x = 5y.
+    fac2.appendLocalId();             // Add local id z.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 1 division should be matched + 2 unmatched local ids.
+    EXPECT_EQ(fac1.getNumLocalIds(), 3u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 3u);
+  }
 }
 
 TEST(FlatAffineConstraintsTest, mergeDivisionsNestedDivsions) {
   {
     // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z >= x).
     FlatAffineConstraints fac1(1);
     fac1.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
     fac1.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
TEST(FlatAffineConstraintsTest, mergeDivisionsNestedDivsions) {
 
     // Local space should be same.
     EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
 
     // 3 divisions should be matched.
     EXPECT_EQ(fac1.getNumLocalIds(), 3u);
     EXPECT_EQ(fac2.getNumLocalIds(), 3u);
   }
+  {
+    // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z >= x).
+    FlatAffineConstraints fac1(1);
+    // This division would be normalized.
+    fac1.addLocalFloorDiv({2, 0}, 4);    // y = [2x / 4] -> [x / 2].
+    fac1.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
+    fac1.addInequality({-1, 1, 1, 0});   // y + z >= x.
+
+    // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z <= x).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
+    fac2.addLocalFloorDiv({3, 3, 0}, 9); // z = [3x + 3y / 9] -> [x + y / 3].
+    fac2.addInequality({1, -1, -1, 0});  // y + z <= x.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 2 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 2u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 2u);
+  }
 }
 
 TEST(FlatAffineConstraintsTest, mergeDivisionsConstants) {
   {
     // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z >= x).
     FlatAffineConstraints fac1(1);
     fac1.addLocalFloorDiv({1, 1}, 2);    // y = [x + 1 / 2].
     fac1.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
TEST(FlatAffineConstraintsTest, mergeDivisionsConstants) {
 
     // Local space should be same.
     EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(fac1.getNumLocalIds(), 2u);
     EXPECT_EQ(fac2.getNumLocalIds(), 2u);
   }
+  {
+    // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z >= x).
+    FlatAffineConstraints fac1(1);
+    fac1.addLocalFloorDiv({1, 1}, 2); // y = [x + 1 / 2].
+    // Normalization test.
+    fac1.addLocalFloorDiv({3, 0, 6}, 9); // z = [3x + 6 / 9] -> [x + 2 / 3].
+    fac1.addInequality({-1, 1, 1, 0});   // y + z >= x.
+
+    // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
+    FlatAffineConstraints fac2(1);
+    // Normalization test.
+    fac2.addLocalFloorDiv({2, 2}, 4);    // y = [2x + 2 / 4] -> [x + 1 / 2].
+    fac2.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
+    fac2.addInequality({1, -1, -1, 0});  // y + z <= x.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 2 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 2u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 2u);
+  }
 }
 
 } // namespace mlir