[MLIR] FlatAffineConstraints: Allow extraction of explicit representation of local variables

[Groverkss]

Summary

This patch extends the existing functionality of computing an explicit
representation for local variables, to also get the explicit representation,
instead of only the inequality pairs.

This is required for a future patch to remove redundant local ids based on
their explicit representation.

[arjunp]

Seems good to me. Minor comments below.

mlir/include/mlir/Analysis/AffineStructures.h

165–166

I think the other way is easier to read: if no explicit representation could be found for the i^th local identifier, denominators[i] is set to 0.

mlir/lib/Analysis/AffineStructures.cpp

1430–1431

Returns the upper-lower bound inequality pair using which the floordiv can be computed.

Might be clearer to say that it "returns the upper and lower bound inequalities from which constrain the local variable to be a floordiv.

1432

If the representation could be computed, dividend and
/// denominator.

Incomplete sentence.

[Groverkss]

• Address Arjun's comments

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// 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.
-  void getLocalReprLbUbPairs(
+  /// 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.
+  ///
+  /// 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<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) const;
 
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
   /// Returns the value at the specified equality row and column.
   inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); }
   inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
 

mlir/lib/Analysis/AffineStructures.cpp

static LogicalResult getDivRepr(const FlatAffineConstraints &cst, unsigned pos,
 
   // Set divisor.
   divisor = denominator;
 
   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),
+/// 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.
 static Optional<std::pair<unsigned, unsigned>>
 computeSingleVarRepr(const FlatAffineConstraints &cst,
-                     const SmallVector<bool, 8> &foundRepr, unsigned pos) {
+                     const SmallVector<bool, 8> &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);
 
   for (unsigned ubPos : ubIndices) {
     for (unsigned lbPos : lbIndices) {
       // Attempt to get divison representation from ubPos, lbPos.
-      SmallVector<int64_t, 8> expr;
-      unsigned divisor;
-      if (failed(getDivRepr(cst, pos, ubPos, lbPos, expr, divisor)))
+      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] && expr[c] != 0)
+        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)
         continue;
 
       return std::make_pair(ubPos, lbPos);
     }
   }
 
   return llvm::None;
 }
 
-/// 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.
-void FlatAffineConstraints::getLocalReprLbUbPairs(
+void FlatAffineConstraints::getLocalReprs(
+    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+  std::vector<SmallVector<int64_t, 8>> dividends(getNumLocalIds());
+  SmallVector<unsigned, 4> denominators(getNumLocalIds());
+  getLocalReprs(dividends, denominators, repr);
+}
+
+void FlatAffineConstraints::getLocalReprs(
+    std::vector<SmallVector<int64_t, 8>> &dividends,
+    SmallVector<unsigned, 4> &denominators) const {
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
+      getNumLocalIds());
+  getLocalReprs(dividends, denominators, repr);
+}
+
+void FlatAffineConstraints::getLocalReprs(
+    std::vector<SmallVector<int64_t, 8>> &dividends,
+    SmallVector<unsigned, 4> &denominators,
     std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
-  assert(repr.size() == getNumLocalIds() &&
-         "Size of repr does not match number of local variables");
+
+  repr.resize(getNumLocalIds());
+  dividends.resize(getNumLocalIds());
+  denominators.resize(getNumLocalIds());
 
   SmallVector<bool, 8> foundRepr(getNumIds(), false);
   for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i)
     foundRepr[i] = true;
 
   unsigned divOffset = getNumDimAndSymbolIds();
   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]) {
-        if (auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i)) {
+        if (auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i,
+                                            dividends[i], denominators[i])) {
           foundRepr[i + divOffset] = true;
           repr[i] = res;
           changed = true;
         }
       }
     }
   } 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].hasValue())
+      denominators[i] = 0;
 }
 
 /// Tightens inequalities given that we are dealing with integer spaces. This is
 /// analogous to the GCD test but applied to inequalities. The constant term can
 /// be reduced to the preceding multiple of the GCD of the coefficients, i.e.,
 ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
 /// fast method - linear in the number of coefficients.
 // Example on how this affects practical cases: consider the scenario:
static bool detectAsFloorDiv(const FlatAffineConstraints &cst, unsigned pos,
   assert(pos < cst.getNumIds() && "invalid position");
 
   // Get upper-lower bound pair for this variable.
   SmallVector<bool, 8> foundRepr(cst.getNumIds(), false);
   for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i)
     if (exprs[i])
       foundRepr[i] = true;
 
-  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos);
+  SmallVector<int64_t, 8> dividend;
+  unsigned divisor;
+  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor);
 
   // No upper-lower bound pair found for this var.
   if (!ulPair)
     return false;
 
-  unsigned ubPos = ulPair->first;
-
-  // Upper bound is of the form:
-  //      -divisor * id + expr >= 0
-  // where `id` is equivalent to `expr floordiv divisor`.
-  //
-  // Since the division cannot be dependent on itself, the coefficient of
-  // of `id` in `expr` is zero. The coefficient of `id` in the upperbound
-  // is -divisor.
-  int64_t divisor = -cst.atIneq(ubPos, pos);
-  int64_t constantTerm = cst.atIneq(ubPos, cst.getNumCols() - 1);
-
   // Construct the dividend expression.
-  auto dividendExpr = getAffineConstantExpr(constantTerm, 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;
-    // computeSingleVarRepr guarantees that expr is known here.
-    dividendExpr = dividendExpr + ubVal * exprs[c];
-  }
+  auto dividendExpr = getAffineConstantExpr(dividend.back(), context);
+  for (unsigned c = 0, f = cst.getNumIds(); c < f; c++)
+    if (dividend[c] != 0)
+      dividendExpr = dividendExpr + dividend[c] * exprs[c];
 
   // Successfully detected the floordiv.
   exprs[pos] = dividendExpr.floorDiv(divisor);
   return true;
 }
 
 // Fills an inequality row with the value 'val'.
 static inline void fillInequality(FlatAffineConstraints *cst, unsigned r,

mlir/lib/Analysis/PresburgerSet.cpp

static void subtractRecursively(FlatAffineConstraints &b, Simplex &simplex,
   }
   FlatAffineConstraints sI = s.getFlatAffineConstraints(i);
   unsigned bInitNumLocals = b.getNumLocalIds();
 
   // 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());
-  sI.getLocalReprLbUbPairs(repr);
+  sI.getLocalReprs(repr);
 
   // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
   // sI's locals.
   b.mergeLocalIds(sI);
 
   // Mark which inequalities of sI are division inequalities and add all such
   // inequalities to b.
   llvm::SmallBitVector isDivInequality(sI.getNumInequalities());

mlir/unittests/Analysis/AffineStructuresTest.cpp

static void checkDivisionRepresentation(
   assert(divisions.size() == fac.getNumLocalIds() &&
          "Size of expected divisions does not match number of local variables");
   assert(
       denoms.size() == fac.getNumLocalIds() &&
       "Size of expected denominators does not match number of local variables");
 
   std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> res(
       fac.getNumLocalIds(), llvm::None);
-  fac.getLocalReprLbUbPairs(res);
+  fac.getLocalReprs(res);
 
   // Check if all expected divisions are computed.
   for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i)
     if (denoms[i] > 0)
       EXPECT_TRUE(res[i].hasValue());
     else
       EXPECT_FALSE(res[i].hasValue());