[MLIR] The return type in the computeSingleVarRepr function is modified to include equality expressions.

Earlier `computeSingleVarRepr` was returning a pair of upper bound and
lower bound indices of the inequality contraints that can be expressed
as a floordiv of an affine function. The equality expression can also be
expressed as a floordiv but contains only one index and hence the `LocalRepr`
class is introduced to facilitate this.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D117430

Author: pashu123

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

@@ -14,6 +14,7 @@
 #define MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
 
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
@@ -267,12 +268,10 @@ class IntegerPolyhedron {
   /// 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<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
-  void getLocalReprs(
-      std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
+  void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
+                     SmallVector<unsigned, 4> &denominators,
+                     std::vector<presburger_utils::MaybeLocalRepr> &repr) const;
+  void getLocalReprs(std::vector<presburger_utils::MaybeLocalRepr> &repr) const;
   void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators) const;
 

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

@@ -21,6 +21,23 @@ class IntegerPolyhedron;
 
 namespace presburger_utils {
 
+/// `ReprKind` enum is used to set the constraint type in `MaybeLocalRepr`.
+enum class ReprKind { Inequality, Equality, None };
+
+/// `MaybeLocalRepr` contains the indices of the contraints that can be
+/// expressed as a floordiv of an affine function. If it's an `equality`
+/// contraint `equalityIdx` is set, in case of `inequality` the `lowerBoundIdx`
+/// and `upperBoundIdx` is set. By default the kind attribute is set to None.
+struct MaybeLocalRepr {
+  ReprKind kind = ReprKind::None;
+  union {
+    unsigned equalityIdx;
+    struct {
+      unsigned lowerBoundIdx, upperBoundIdx;
+    } 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
@@ -29,10 +46,10 @@ namespace presburger_utils {
 /// 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.
-Optional<std::pair<unsigned, unsigned>>
-computeSingleVarRepr(const IntegerPolyhedron &cst, ArrayRef<bool> foundRepr,
-                     unsigned pos, SmallVector<int64_t, 8> &dividend,
-                     unsigned &divisor);
+MaybeLocalRepr computeSingleVarRepr(const IntegerPolyhedron &cst,
+                                    ArrayRef<bool> foundRepr, unsigned pos,
+                                    SmallVector<int64_t, 8> &dividend,
+                                    unsigned &divisor);
 
 } // namespace presburger_utils
 } // namespace mlir

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

@@ -21,6 +21,7 @@
 #define DEBUG_TYPE "presburger"
 
 using namespace mlir;
+using namespace presburger_utils;
 using llvm::SmallDenseMap;
 using llvm::SmallDenseSet;
 
@@ -799,8 +800,7 @@ bool IntegerPolyhedron::containsPoint(ArrayRef<int64_t> point) const {
   return true;
 }
 
-void IntegerPolyhedron::getLocalReprs(
-    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+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);
@@ -809,15 +809,14 @@ void IntegerPolyhedron::getLocalReprs(
 void IntegerPolyhedron::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators) const {
-  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
-      getNumLocalIds());
+  std::vector<MaybeLocalRepr> repr(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void IntegerPolyhedron::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators,
-    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+    std::vector<MaybeLocalRepr> &repr) const {
 
   repr.resize(getNumLocalIds());
   dividends.resize(getNumLocalIds());
@@ -835,11 +834,13 @@ void IntegerPolyhedron::getLocalReprs(
     changed = false;
     for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) {
       if (!foundRepr[i + divOffset]) {
-        if (auto res = presburger_utils::computeSingleVarRepr(
-                *this, foundRepr, divOffset + i, dividends[i],
-                denominators[i])) {
+        auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i,
+                                        dividends[i], denominators[i]);
+        if (res.kind == ReprKind::Inequality) {
           foundRepr[i + divOffset] = true;
-          repr[i] = res;
+          repr[i].kind = ReprKind::Inequality;
+          repr[i].repr.inEqualityPair = {res.repr.inEqualityPair.lowerBoundIdx,
+                                         res.repr.inEqualityPair.upperBoundIdx};
           changed = true;
         }
       }
@@ -849,7 +850,7 @@ void IntegerPolyhedron::getLocalReprs(
   // 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].hasValue())
+    if (repr[i].kind == ReprKind::None)
       denominators[i] = 0;
 }
 

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

@@ -8,10 +8,12 @@
 
 #include "mlir/Analysis/Presburger/PresburgerSet.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 PresburgerSet::PresburgerSet(const IntegerPolyhedron &poly)
     : nDim(poly.getNumDimIds()), nSym(poly.getNumSymbolIds()) {
@@ -209,8 +211,7 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
 
   // Find out which inequalities of sI correspond to division inequalities for
   // the local variables of sI.
-  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
-      sI.getNumLocalIds());
+  std::vector<MaybeLocalRepr> repr(sI.getNumLocalIds());
   sI.getLocalReprs(repr);
 
   // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
@@ -220,18 +221,20 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
   // Mark which inequalities of sI are division inequalities and add all such
   // inequalities to b.
   llvm::SmallBitVector isDivInequality(sI.getNumInequalities());
-  for (Optional<std::pair<unsigned, unsigned>> &maybePair : repr) {
-    assert(maybePair &&
+  for (MaybeLocalRepr &maybeInequality : repr) {
+    assert(maybeInequality.kind == ReprKind::Inequality &&
            "Subtraction is not supported when a representation of the local "
            "variables of the subtrahend cannot be found!");
+    auto lb = maybeInequality.repr.inEqualityPair.lowerBoundIdx;
+    auto ub = maybeInequality.repr.inEqualityPair.upperBoundIdx;
 
-    b.addInequality(sI.getInequality(maybePair->first));
-    b.addInequality(sI.getInequality(maybePair->second));
+    b.addInequality(sI.getInequality(lb));
+    b.addInequality(sI.getInequality(ub));
 
-    assert(maybePair->first != maybePair->second &&
+    assert(lb != ub &&
            "Upper and lower bounds must be different inequalities!");
-    isDivInequality[maybePair->first] = true;
-    isDivInequality[maybePair->second] = true;
+    isDivInequality[lb] = true;
+    isDivInequality[ub] = true;
   }
 
   unsigned offset = simplex.getNumConstraints();

mlir/lib/Analysis/Presburger/Utils.cpp

@@ -16,6 +16,7 @@
 #include "mlir/Support/MathExtras.h"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 /// 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,
@@ -144,7 +145,7 @@ static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos,
 /// be computed. If the representation could be computed, `dividend` and
 /// `denominator` are set. If the representation could not be computed,
 /// `llvm::None` is returned.
-Optional<std::pair<unsigned, unsigned>> presburger_utils::computeSingleVarRepr(
+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");
@@ -153,6 +154,7 @@ Optional<std::pair<unsigned, unsigned>> presburger_utils::computeSingleVarRepr(
 
   SmallVector<unsigned, 4> lbIndices, ubIndices;
   cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
+  MaybeLocalRepr repr;
 
   for (unsigned ubPos : ubIndices) {
     for (unsigned lbPos : lbIndices) {
@@ -178,9 +180,10 @@ Optional<std::pair<unsigned, unsigned>> presburger_utils::computeSingleVarRepr(
       if (c < f)
         continue;
 
-      return std::make_pair(ubPos, lbPos);
+      repr.kind = ReprKind::Inequality;
+      repr.repr.inEqualityPair = {ubPos, lbPos};
+      return repr;
     }
   }
-
-  return llvm::None;
+  return repr;
 }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

@@ -31,6 +31,7 @@
 #define DEBUG_TYPE "affine-structures"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 namespace {
 
@@ -851,11 +852,10 @@ static bool detectAsFloorDiv(const FlatAffineConstraints &cst, unsigned pos,
 
   SmallVector<int64_t, 8> dividend;
   unsigned divisor;
-  auto ulPair = presburger_utils::computeSingleVarRepr(cst, foundRepr, pos,
-                                                       dividend, divisor);
+  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor);
 
   // No upper-lower bound pair found for this var.
-  if (!ulPair)
+  if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality)
     return false;
 
   // Construct the dividend expression.