[MLIR] Extract division representation from equality expressions.

[pashu123]

Summary

Extract the division representation from equality constraints.
For example:

32*k == 16*i + j - 31                 <-- k is the localVariable
expr = 16*i + j - 31, divisor = 32
k = (16*i + j - 32) floordiv 32

The dividend of the division is set to [16, 1, -32] and the divisor is set
to 32.

[pashu123]

Minor comments.

[pashu123]

Updating comments.

[pashu123]

Clang format.

[Groverkss]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

223

Why is this default value removed? If you do decide to remove it, please specify how to make equalities optional (i.e. by passing a nullptr).

mlir/lib/Analysis/Presburger/Utils.cpp

180–182

186

I don't think you need to pass this by reference.

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

737–772

Could you add some tests where there are both inequalities and equalities and divisions are extracted from both?

[arjunp]

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

803

Maybe not for this patch, but the std::vector can be changed to SmallVector now.

838–843

mlir/lib/Analysis/Presburger/Utils.cpp

163

I would suggest either writing a comment that the equality must involve the pos-th variable and assert temp_div != 0, or return failure if temp_div == 0.

[pashu123]

Addressing arjun's and Kunwar's comments.

[pashu123]

minor comments.

[pashu123]

Addressing comments.

[pashu123]

Removing unneccesary debugs.

[pashu123]

Updating comments.

[pashu123]

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

803

Indeed, we can do that in the next patch.

838–843

I missed it. This is very elegant. Thanks.

[Groverkss]

This comment has been deleted.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

803

Please update the docs for getLocalReprs in header file to reflect the new return type MaybeLocalRepr and that it supports not only inequalities but also equalities.

mlir/lib/Analysis/Presburger/Utils.cpp

169–170

210–215

Please update these docs to reflect the return type. Same in the header file.

[pashu123]

Updating comments.

[pashu123]

Clang format warnings.

[arjunp]

mlir/lib/Transforms/Utils/LoopUtils.cpp

3293

Why is this change required?

[pashu123]

removing unneccesary changes.

[pashu123]

mlir/lib/Transforms/Utils/LoopUtils.cpp

3293

Reverting, this is not required.

[Groverkss]

mlir/lib/Analysis/Presburger/Utils.cpp

188

I don't see any reason for this to be const.

[pashu123]

mlir/lib/Analysis/Presburger/Utils.cpp

188

We don't want pos to be changed by the function and hence making it const makes sense here. There are many benefits associated: const tells the compiler that a variable or method is immutable. This helps the compiler optimize the code and helps the developer know if a function has a side effect. Also, using const & prevents the compiler from copying data unnecessarily.

[pashu123]

mlir/lib/Analysis/Presburger/Utils.cpp

188

I think pass by reference and const makes more sense.

[pashu123]

nit

[pashu123]

Removing const.

[Groverkss]

LGTM. Just a minor comment.

mlir/lib/Analysis/Presburger/Utils.cpp

169–170

I think you missed this comment.

[pashu123]

Minor comments.

[arjunp]

mlir/lib/Analysis/Presburger/Utils.cpp

163

Oh, I think it should be called tempDiv.

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

public:
 
   /// Gather positions of all lower and upper bounds of the identifier at `pos`,
   /// and optionally any equalities on it. In addition, the bounds are to be
   /// independent of identifiers in position range [`offset`, `offset` + `num`).
   void
   getLowerAndUpperBoundIndices(unsigned pos,
                                SmallVectorImpl<unsigned> *lbIndices,
                                SmallVectorImpl<unsigned> *ubIndices,
                                SmallVectorImpl<unsigned> *eqIndices = nullptr,
                                unsigned offset = 0, unsigned num = 0) const;
 
   /// Checks for emptiness by performing variable elimination on all
   /// identifiers, running the GCD test on each equality constraint, and
   /// checking for invalid constraints. Returns true if the GCD test fails for
   /// any equality, or if any invalid constraints are discovered on any row.
   /// Returns false otherwise.
   bool isEmpty() const;
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 inequalities identified by their position indices, using
-  /// which an explicit representation for each local variable can be computed.
-  /// The pairs are stored as indices of upperbound, lowerbound inequalities. If
-  /// no such pair can be found, it is stored as llvm::None.
+  /// Find equality and pairs of inequality contraints identified by their
+  /// position indices, using which an explicit representation for each local
+  /// variable can be computed. The indices of the constraints are stored in
+  /// `MaybeLocalRepr` struct. If no such pair can be found, the kind attribute
+  /// in `MaybeLocalRepr` is set to None.
   ///
   /// The dividends of the explicit representations are stored in `dividends`
   /// and the denominators in `denominators`. If no explicit representation
   /// could be found for the `i^th` local identifier, `denominators[i]` is set
   /// to 0.
   void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators,
                      std::vector<presburger_utils::MaybeLocalRepr> &repr) const;

mlir/include/mlir/Analysis/Presburger/Utils.h

union {
     } inEqualityPair;
   } repr;
 };
 
 /// 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 be computed. If the representation could be computed, `dividend` and
-/// `denominator` are set. If the representation could not be computed,
-/// `llvm::None` is returned.
+/// Returns the `MaybeLocalRepr` struct which contains the indices of the
+/// constraints that can be expressed as a floordiv of an affine function. If
+/// the representation could be computed, `dividend` and `denominator` are set.
+/// If the representation could not be computed, the kind attribute in
+/// `MaybeLocalRepr` is set to None.
 MaybeLocalRepr computeSingleVarRepr(const IntegerPolyhedron &cst,
                                     ArrayRef<bool> foundRepr, unsigned pos,
                                     SmallVector<int64_t, 8> &dividend,
                                     unsigned &divisor);
 
 } // namespace presburger_utils
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_UTILS_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

bool IntegerPolyhedron::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;
 }
 
 void IntegerPolyhedron::getLocalReprs(std::vector<MaybeLocalRepr> &repr) const {
   std::vector<SmallVector<int64_t, 8>> dividends(getNumLocalIds());
   SmallVector<unsigned, 4> denominators(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void IntegerPolyhedron::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators) const {
void IntegerPolyhedron::getLocalReprs(
   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]) {
         auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i,
                                         dividends[i], denominators[i]);
-        if (res.kind == ReprKind::Inequality) {
-          foundRepr[i + divOffset] = true;
-          repr[i].kind = ReprKind::Inequality;
-          repr[i].repr.inEqualityPair = {res.repr.inEqualityPair.lowerBoundIdx,
-                                         res.repr.inEqualityPair.upperBoundIdx};
-          changed = true;
-        }
-      }
-    }
+        if (res.kind == ReprKind::None)
+          continue;
+        repr[i] = res;
   } while (changed);
 
   // Set 0 denominator for identifiers for which no division representation
   // could be found.
   for (unsigned i = 0, e = repr.size(); i < e; ++i)
     if (repr[i].kind == ReprKind::None)
       denominators[i] = 0;
 }

mlir/lib/Analysis/Presburger/Utils.cpp

static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos,
   // 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 represented as a division using
+/// equality at position `eqInd`.
+///
+/// For example:
+///     32*k == 16*i + j - 31                 <-- `eqInd` for 'k'
+///     expr = 16*i + j - 31, divisor = 32
+///     k = (16*i + j - 31) floordiv 32
+///
+/// If successful, `expr` is set to dividend of the division and `divisor` is
+/// set to the denominator of the division. The final division expression is
+/// normalized by GCD.
+static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos,
+                                unsigned eqInd, SmallVector<int64_t, 8> &expr,
+                                unsigned &divisor) {
+
+  assert(pos <= cst.getNumIds() && "Invalid identifier position");
+  assert(eqInd <= cst.getNumEqualities() && "Invalid equality position");
+
+  // Extract divisor, the divisor can be negative and hence its sign information
+  // is stored in `signDiv` to reverse the sign of dividend's coefficients.
+  // Equality must involve the pos-th variable and hence `temp_div` != 0.
+  int64_t temp_div = cst.atEq(eqInd, pos);
+  if(temp_div == 0)
+      return failure();
+  int64_t signDiv = temp_div < 0 ? -1 : 1;
+
+  // Divisor is always positive integer.
+  divisor = temp_div * signDiv;
+
+  expr.resize(cst.getNumCols(), 0);
+  for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i)
+    if (i != pos)
+      expr[i] = signDiv * cst.atEq(eqInd, i);
+
+  expr.back() = signDiv * cst.atEq(eqInd, cst.getNumCols() - 1);
+  normalizeDivisionByGCD(expr, divisor);
+
+  return success();
+}
+
+// Returns `false` if the constraints depends on a variable for which an
+// explicit representation has not been found yet, otherwise returns `true`.
+static bool checkExplicitRepresentation(const IntegerPolyhedron &cst,
+                                        ArrayRef<bool> foundRepr,
+                                        SmallVectorImpl<int64_t> &dividend,
+                                        const unsigned pos) {
+  // Exit to avoid circular dependencies between divisions.
+  unsigned c, f;
+  for (c = 0, f = cst.getNumIds(); c < f; ++c) {
+    if (c == pos)
+      continue;
+    if (!foundRepr[c] && dividend[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)
+    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.
-/// Returns the upper and lower bound inequalities using which the floordiv can
-/// be computed. If the representation could be computed, `dividend` and
-/// `denominator` are set. If the representation could not be computed,
-/// `llvm::None` is returned.
+/// Returns the `MaybeLocalRepr` struct which contains the indices of the
+/// constraints that can be expressed as a floordiv of an affine function. If
+/// the representation could be computed, `dividend` and `denominator` are set.
+/// If the representation could not be computed, the kind attribute in
+/// `MaybeLocalRepr` is set to None.
 MaybeLocalRepr presburger_utils::computeSingleVarRepr(
     const IntegerPolyhedron &cst, ArrayRef<bool> foundRepr, unsigned pos,
     SmallVector<int64_t, 8> &dividend, unsigned &divisor) {
   assert(pos < cst.getNumIds() && "invalid position");
   assert(foundRepr.size() == cst.getNumIds() &&
          "Size of foundRepr does not match total number of variables");
 
-  SmallVector<unsigned, 4> lbIndices, ubIndices;
-  cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
-  MaybeLocalRepr repr;
+  SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
+  cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices, &eqIndices);
+  MaybeLocalRepr repr{};
 
   for (unsigned ubPos : ubIndices) {
     for (unsigned lbPos : lbIndices) {
       // Attempt to get divison representation from ubPos, lbPos.
       if (failed(getDivRepr(cst, pos, ubPos, lbPos, dividend, divisor)))
         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.
-      unsigned c, f;
-      for (c = 0, f = cst.getNumIds(); c < f; ++c) {
-        if (c == pos)
-          continue;
-        if (!foundRepr[c] && dividend[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)
+      if (!checkExplicitRepresentation(cst, foundRepr, dividend, pos))
         continue;
 
       repr.kind = ReprKind::Inequality;
       repr.repr.inEqualityPair = {ubPos, lbPos};
       return repr;
     }
   }
+  for (unsigned eqPos : eqIndices) {
+    // Attempt to get divison representation from eqPos.
+    if (failed(getDivRepr(cst, pos, eqPos, dividend, divisor)))
+      continue;
+
+    if (!checkExplicitRepresentation(cst, foundRepr, dividend, pos))
+      continue;
+
+    repr.kind = ReprKind::Equality;
+    repr.repr.equalityIdx = eqPos;
+    return repr;
+  }
   return repr;
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

if (!cst.getConstantBoundOnDimSize(0, /*lb=*/nullptr,
                                        /*boundFloorDivisor=*/nullptr,
                                        /*ub=*/nullptr, &fullTileLbPos,
                                        &fullTileUbPos)) {
       LLVM_DEBUG(llvm::dbgs() << "Can't get constant diff pair for a loop\n");
       return nullptr;
     }
 
     SmallVector<unsigned, 4> lbIndices, ubIndices;
-    cst.getLowerAndUpperBoundIndices(/*pos=*/0, &lbIndices, &ubIndices);
+    cst.getLowerAndUpperBoundIndices(/*pos=*/0, &lbIndices, &ubIndices,
+                                     /*eqIndices=*/nullptr);
 
     auto fLb = cst.getInequality(fullTileLbPos);
     auto fUb = cst.getInequality(fullTileUbPos);
     fullTileLb.assign(fLb.begin(), fLb.end());
     fullTileUb.assign(fUb.begin(), fUb.end());
 
     // Full tile lower bound should be >= than any other lower bound.
     for (auto lbIndex : lbIndices)

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

MLIRContext context;
     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(poly, divisions, denoms);
   }
 }
 
+TEST(IntegerPolyhedronTest, computeLocalReprFromEquality) {
+  MLIRContext context;
+  {
+    IntegerPolyhedron poly =
+        parsePoly("(i, j, q) : (-4*q + i + j == 0)", &context);
+    // Convert `q` to a local variable.
+    poly.convertDimToLocal(2, 3);
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
+    SmallVector<unsigned, 8> denoms = {4};
+
+    checkDivisionRepresentation(poly, divisions, denoms);
+  }
+  {
+    IntegerPolyhedron poly =
+        parsePoly("(i, j, q) : (4*q - i - j == 0)", &context);
+    // Convert `q` to a local variable.
+    poly.convertDimToLocal(2, 3);
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
+    SmallVector<unsigned, 8> denoms = {4};
+
+    checkDivisionRepresentation(poly, divisions, denoms);
+  }
+  {
+    IntegerPolyhedron poly =
+        parsePoly("(i, j, q) : (3*q + i + j - 2 == 0)", &context);
+    // Convert `q` to a local variable.
+    poly.convertDimToLocal(2, 3);
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, -2}};
+    SmallVector<unsigned, 8> denoms = {3};
+
+    checkDivisionRepresentation(poly, divisions, denoms);
+  }
+}
+
+TEST(IntegerPolyhedronTest, computeLocalReprFromEqualityAndInequality) {
+  MLIRContext context;
+  {
+    IntegerPolyhedron poly = parsePoly("(i, j, q, k) : (-3*k + i + j == 0, 4*q - "
+                                       "i - j + 2 >= 0, -4*q + i + j >= 0)",
+                                       &context);
+    // Convert `q` and `k` to local variables.
+    poly.convertDimToLocal(2, 4);
+
+    std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, 0, 1},
+                                                      {-1, -1, 0, 0, 0}};
+    SmallVector<unsigned, 8> denoms = {4, 3};
+
+    checkDivisionRepresentation(poly, divisions, denoms);
+  }
+}
+
 TEST(IntegerPolyhedronTest, computeLocalReprNoRepr) {
   MLIRContext context;
   IntegerPolyhedron poly =
       parsePoly("(x, q) : (x - 3 * q >= 0, -x + 3 * q + 3 >= 0)", &context);
   // Convert q to a local variable.
   poly.convertDimToLocal(1, 2);
 
   std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0}};