[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

1941–1943 ↗

(On Diff #393731)

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

1955–1959 ↗

(On Diff #393731)

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.

1997 ↗

(On Diff #393731)

Unnecessary newline.

2008–2009 ↗

(On Diff #393731)

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

2008–2009 ↗

(On Diff #393731)

I see, thanks.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1355–1356 ↗

(On Diff #393769)

1355–1357 ↗

(On Diff #393769)

1358 ↗

(On Diff #393769)

I think a better name would be normalizeDivisionByGCD.

1358 ↗

(On Diff #393769)

Could you rename expr to dividend for more clarity?

1366 ↗

(On Diff #393769)

1366–1367 ↗

(On Diff #393769)

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

1403–1404 ↗

(On Diff #393769)

mlir/unittests/Analysis/AffineStructuresTest.cpp

754

Just mention here that this tests normalization.

840

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 ↗

(On Diff #394317)

1367–1373 ↗

(On Diff #394317)

[bondhugula]

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

mlir/lib/Analysis/AffineStructures.cpp

1358 ↗

(On Diff #394317)

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

1360–1365 ↗

(On Diff #394317)

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 ↗

(On Diff #394317)

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

1409 ↗

(On Diff #394317)

"returned division" isn't meaningful. Rephrase?

[pashu123]

Addressing Uday's comments.

[pashu123]

Clang format.

[pashu123]

mlir/lib/Analysis/AffineStructures.cpp

1360–1365 ↗

(On Diff #394317)

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

1409 ↗

(On Diff #394317)

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

Capture by value is sufficient.

mlir/unittests/Analysis/AffineStructuresTest.cpp

598–601

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

754

Looks like this comment is still unaddressed.

773

1 division <-- One division ...

[pashu123]

Addressing bondhugula's comment.

[pashu123]

Comments.

[pashu123]

Comments.

[Groverkss]

mlir/unittests/Analysis/AffineStructuresTest.cpp

598–601

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

598–601

I have updated the test case.

[pashu123]

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

mlir/lib/Analysis/Presburger/Utils.cpp

 //===- Utils.cpp - General utilities for Presburger library ---------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Utility functions required by the Presburger Library.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Support/MathExtras.h"
 
 using namespace mlir;
 
+/// Normalize a division's `dividend` and the `divisor` by their GCD. For
+/// example: if the dividend and divisor are [2,0,4] and 4 respectively,
+/// they get normalized to [1,0,2] and 2.
+static void normalizeDivisionByGCD(SmallVectorImpl<int64_t> &dividend,
+                                   unsigned &divisor) {
+  if (divisor == 0 || dividend.empty())
+    return;
+  int64_t gcd = llvm::greatestCommonDivisor(dividend.front(), int64_t(divisor));
+
+  // The reason for ignoring the constant term is as follows.
+  // For a division:
+  //      floor((a + m.f(x))/(m.d))
+  // It can be replaced by:
+  //      floor((floor(a/m) + f(x))/d)
+  // Since `{a/m}/d` in the dividend satisfies 0 <= {a/m}/d < 1/d, it will not
+  // influence the result of the floor division and thus, can be ignored.
+  for (size_t i = 1, m = dividend.size() - 1; i < m; i++) {
+    gcd = llvm::greatestCommonDivisor(dividend[i], gcd);
+    if (gcd == 1)
+      return;
+  }
+
+  // Normalize the dividend and the denominator.
+  std::transform(dividend.begin(), dividend.end(), dividend.begin(),
+                 [gcd](int64_t &n) { return floor(n / gcd); });
+  divisor /= 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:

 /// 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.
+/// set to the denominator of the division. The final division expression is
+/// normalized by GCD.
 static LogicalResult getDivRepr(const IntegerPolyhedron &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");
static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos,
   for (i = 0, e = cst.getNumIds(); i < e; ++i)
     if (i != pos)
       expr[i] = cst.atIneq(ubIneq, i);
 
   // 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;
+  normalizeDivisionByGCD(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.

mlir/unittests/Analysis/AffineStructuresTest.cpp

 
 TEST(FlatAffineConstraintsTest, computeLocalReprConstantFloorDiv) {
   FlatAffineConstraints fac(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);
+  fac.addLocalFloorDiv({0, 0, 0, 0, 100}, 30);
+  fac.addLocalFloorDiv({0, 0, 0, 0, 0, 206}, 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, 3},
+                                                    {0, 0, 0, 0, 0, 0, 2}};
+  SmallVector<unsigned, 8> denoms = {1, 1};
 
   // Check if floordivs with constant numerator can be computed.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
 TEST(FlatAffineConstraintsTest, computeLocalReprRecursive) {
   FlatAffineConstraints fac(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);
   }
+
+  {
+    // Division Normalization test.
+    // (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());
+
+    // One 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);
+    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].
+    // This division would be normalized.
+    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