[MLIR][Presburger] Support divisions in union of two PWMAFunction

[Groverkss]

Summary

This patch adds support for divisions in the union of two PWMAFunction. This is
now possible because of previous patches, which made divisions explicitly
stored in MultiAffineFunction (MAF). This patch also refactors the previous
implementation, moving the implementation for obtaining a set of points where a
MAF is lexicographically "better" than the other to MAF.

[Groverkss]

• Fix spell error

[arjunp]

Why is the Compare change in this patch?

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

26

maybe use Ordering to follow std::strong_ordering

99–102

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

286

356

maybe mention in a comment that this is templated since it's passed as a a lambda. wish we had better partial application in c++!

[bondhugula]

Please add a commit summary.

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

26

Compare -> ComparisonKind

Please use more descriptive naming.

[Groverkss]

Address comments

[Groverkss]

I used OrderingKind enum in this patch because I wanted to move the lexSet implementation to MultiAffineFunction.

[bondhugula]

Removing change request.

[Groverkss]

Rebase

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include <optional>
 
 namespace mlir {
 namespace presburger {
 
+/// Enum representing a binary comparison operator: equal, not equal, less than,
+/// less than or equal, greater than, greater than or equal.
+enum class OrderingKind { EQ, NE, LT, LE, GT, GE };
+
 /// This class represents a multi-affine function with the domain as Z^d, where
 /// `d` is the number of domain variables of the function. For example:
 ///
 /// (x, y) -> (x + 2, 2*x - 3y + 5, 2*x + y).
 ///
 /// The output expressions are represented as a matrix with one row for every
 /// output, one column for each var including division variables, and an extra
 /// column at the end for the constant term.
public:
   PresburgerSpace getDomainSpace() const { return space.getDomainSpace(); }
   PresburgerSpace getOutputSpace() const { return space.getRangeSpace(); }
 
   /// Get a matrix with each row representing row^th output expression.
   const Matrix &getOutputMatrix() const { return output; }
   /// Get the `i^th` output expression.
   ArrayRef<MPInt> getOutputExpr(unsigned i) const { return output.getRow(i); }
 
-  // Remove the specified range of outputs.
+  /// Get the divisions used in this function.
+  const DivisionRepr &getDivs() const { return divs; }
+
+  /// Remove the specified range of outputs.
   void removeOutputs(unsigned start, unsigned end);
 
   /// Given a MAF `other`, merges division variables such that both functions
   /// have the union of the division vars that exist in the functions.
   void mergeDivs(MultiAffineFunction &other);
 
   //// Return the output of the function at the given point.
   SmallVector<MPInt, 8> valueAt(ArrayRef<MPInt> point) const;
   SmallVector<MPInt, 8> valueAt(ArrayRef<int64_t> point) const {
     return valueAt(getMPIntVec(point));
   }
 
   /// Return whether the `this` and `other` are equal when the domain is
   /// restricted to `domain`. This is the case if they lie in the same space,
   /// and their outputs are equal for every point in `domain`.
   bool isEqual(const MultiAffineFunction &other) const;
   bool isEqual(const MultiAffineFunction &other,
                const IntegerPolyhedron &domain) const;
   bool isEqual(const MultiAffineFunction &other,
                const PresburgerSet &domain) const;
 
   void subtract(const MultiAffineFunction &other);
 
+  /// Return the set of domain points where the output of `this` and `other`
+  /// are ordered lexicographically according to the given ordering.
+  /// For example, if the given comparison is `LT`, then the returned set
+  /// contains all points where the first output of `this` is lexicographically
+  /// less than `other`.
+  PresburgerSet getLexSet(OrderingKind comp,
+                          const MultiAffineFunction &other) const;
+
   /// Get this function as a relation.
   IntegerRelation getAsRelation() const;
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
   /// Assert that the MAF is consistent.
public:
   PresburgerSet getDomain() const;
 
   /// Return the output of the function at the given point.
   std::optional<SmallVector<MPInt, 8>> valueAt(ArrayRef<MPInt> point) const;
   std::optional<SmallVector<MPInt, 8>> valueAt(ArrayRef<int64_t> point) const {
     return valueAt(getMPIntVec(point));
   }
 
+  /// Return all the pieces of this piece-wise function.
+  ArrayRef<Piece> getAllPieces() const { return pieces; }
+
   /// Return whether `this` and `other` are equal as PWMAFunctions, i.e. whether
   /// they have the same dimensions, the same domain and they take the same
   /// value at every point in the domain.
   bool isEqual(const PWMAFunction &other) const;
 
   /// Return a function defined on the union of the domains of this and func,
   /// such that when only one of the functions is defined, it outputs the same
   /// as that function, and if both are defined, it outputs the lexmax/lexmin of

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

public:
 
   static PresburgerSpace getSetSpace(unsigned numDims = 0,
                                      unsigned numSymbols = 0,
                                      unsigned numLocals = 0) {
     return PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols,
                            numLocals);
   }
 
-  // Get the domain/range space of this space. The returned space is a set
-  // space.
+  /// Get the domain/range space of this space. The returned space is a set
+  /// space.
   PresburgerSpace getDomainSpace() const;
   PresburgerSpace getRangeSpace() const;
 
+  /// Get the space without local variables.
+  PresburgerSpace getSpaceWithoutLocals() const;
+
   unsigned getNumDomainVars() const { return numDomain; }
   unsigned getNumRangeVars() const { return numRange; }
   unsigned getNumSetDimVars() const { return numRange; }
   unsigned getNumSymbolVars() const { return numSymbols; }
   unsigned getNumLocalVars() const { return numLocals; }
 
   unsigned getNumDimVars() const { return numDomain + numRange; }
   unsigned getNumDimAndSymbolVars() const {

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

void MultiAffineFunction::mergeDivs(MultiAffineFunction &other) {
   output.insertColumns(divOffset + nDivs, newDivs);
   divs = other.divs;
 
   // Check consistency.
   assertIsConsistent();
   other.assertIsConsistent();
 }
 
+PresburgerSet
+MultiAffineFunction::getLexSet(OrderingKind comp,
+                               const MultiAffineFunction &other) const {
+  assert(getSpace().isCompatible(other.getSpace()) &&
+         "Output space of funcs should be compatible");
+
+  // Create copies of functions and merge their local space.
+  MultiAffineFunction funcA = *this;
+  MultiAffineFunction funcB = other;
+  funcA.mergeDivs(funcB);
+
+  // We first create the set `result`, corresponding to the set where output
+  // of funcA is lexicographically larger/smaller than funcB. This is done by
+  // creating a PresburgerSet with the following constraints:
+  //
+  //    (outA[0] > outB[0]) U
+  //    (outA[0] = outB[0], outA[1] > outA[1]) U
+  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] > outA[2]) U
+  //    ...
+  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] > outB[n-1])
+  //
+  // where `n` is the number of outputs.
+  // If `lexMin` is set, the complement inequality is used:
+  //
+  //    (outA[0] < outB[0]) U
+  //    (outA[0] = outB[0], outA[1] < outA[1]) U
+  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] < outA[2]) U
+  //    ...
+  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] < outB[n-1])
+  PresburgerSpace resultSpace = funcA.getDomainSpace();
+  PresburgerSet result =
+      PresburgerSet::getEmpty(resultSpace.getSpaceWithoutLocals());
+  IntegerPolyhedron levelSet(
+      /*numReservedInequalities=*/1 + 2 * resultSpace.getNumLocalVars(),
+      /*numReservedEqualities=*/funcA.getNumOutputs(),
+      /*numReservedCols=*/resultSpace.getNumVars() + 1, resultSpace);
+
+  // Add division inequalities to `levelSet`.
+  for (unsigned i = 0, e = funcA.getNumDivs(); i < e; ++i) {
+    levelSet.addInequality(getDivUpperBound(funcA.divs.getDividend(i),
+                                            funcA.divs.getDenom(i),
+                                            funcA.divs.getDivOffset() + i));
+    levelSet.addInequality(getDivLowerBound(funcA.divs.getDividend(i),
+                                            funcA.divs.getDenom(i),
+                                            funcA.divs.getDivOffset() + i));
+  }
+
+  for (unsigned level = 0; level < funcA.getNumOutputs(); ++level) {
+    // Create the expression `outA - outB` for this level.
+    SmallVector<MPInt, 8> subExpr =
+        subtractExprs(funcA.getOutputExpr(level), funcB.getOutputExpr(level));
+
+    // TODO: Implement all comparison cases.
+    switch (comp) {
+    case OrderingKind::LT:
+      // For less than, we add an upper bound of -1:
+      //        outA - outB <= -1
+      //        outA <= outB - 1
+      //        outA < outB
+      levelSet.addBound(IntegerPolyhedron::BoundType::UB, subExpr, MPInt(-1));
+      break;
+    case OrderingKind::GT:
+      // For greater than, we add a lower bound of 1:
+      //        outA - outB >= 1
+      //        outA > outB + 1
+      //        outA > outB
+      levelSet.addBound(IntegerPolyhedron::BoundType::LB, subExpr, MPInt(1));
+      break;
+    case OrderingKind::GE:
+    case OrderingKind::LE:
+    case OrderingKind::EQ:
+    case OrderingKind::NE:
+      assert(false && "Not implemented case");
+    }
+
+    // Union the set with the result.
+    result.unionInPlace(levelSet);
+    // The last inequality in `levelSet` is the bound we inserted. We remove
+    // that for next iteration.
+    levelSet.removeInequality(levelSet.getNumInequalities() - 1);
+    // Add equality `outA - outB == 0` for this level for next iteration.
+    levelSet.addEquality(subExpr);
+  }
+
+  return result;
+}
+
 /// Two PWMAFunctions are equal if they have the same dimensionalities,
 /// the same domain, and take the same value at every point in the domain.
 bool PWMAFunction::isEqual(const PWMAFunction &other) const {
   if (!space.isCompatible(other.space))
     return false;
 
   if (!this->getDomain().isEqual(other.getDomain()))
     return false;
 
   // Check if, whenever the domains of a piece of `this` and a piece of `other`
   // overlap, they take the same output value. If `this` and `other` have the
   // same domain (checked above), then this check passes iff the two functions
   // have the same output at every point in the domain.
   return llvm::all_of(this->pieces, [&other](const Piece &pieceA) {
     return llvm::all_of(other.pieces, [&pieceA](const Piece &pieceB) {
       PresburgerSet commonDomain = pieceA.domain.intersect(pieceB.domain);
       return pieceA.output.isEqual(pieceB.output, commonDomain);
     });
   });
 }
 
 void PWMAFunction::addPiece(const Piece &piece) {
   assert(piece.isConsistent() && "Piece should be consistent");
+  assert(piece.domain.intersect(getDomain()).isIntegerEmpty() &&
+         "Piece should be disjoint from the function");
   pieces.push_back(piece);
 }
 
 void PWMAFunction::print(raw_ostream &os) const {
   space.print(os);
   os << getNumPieces() << " pieces:\n";
   for (const Piece &piece : pieces) {
     os << "Domain of piece:\n";
PWMAFunction PWMAFunction::unionFunction(
   PresburgerSet dom = getDomain();
   for (const Piece &pieceB : func.pieces)
     result.addPiece({pieceB.domain.subtract(dom), pieceB.output});
 
   return result;
 }
 
 /// A tiebreak function which breaks ties by comparing the outputs
-/// lexicographically. If `lexMin` is true, then the ties are broken by
-/// taking the lexicographically smaller output and otherwise, by taking the
-/// lexicographically larger output.
-template <bool lexMin>
+/// lexicographically based on the given comparison operator.
+/// This is templated since it is passed as a lambda.
+template <OrderingKind comp>
 static PresburgerSet tiebreakLex(const PWMAFunction::Piece &pieceA,
                                  const PWMAFunction::Piece &pieceB) {
-  // TODO: Support local variables here.
-  assert(pieceA.output.getSpace().isCompatible(pieceB.output.getSpace()) &&
-         "Pieces should be compatible");
-  assert(pieceA.domain.getSpace().getNumLocalVars() == 0 &&
-         "Local variables are not supported yet.");
-
-  PresburgerSpace compatibleSpace = pieceA.domain.getSpace();
-  const PresburgerSpace &space = pieceA.domain.getSpace();
-
-  // We first create the set `result`, corresponding to the set where output
-  // of pieceA is lexicographically larger/smaller than pieceB. This is done by
-  // creating a PresburgerSet with the following constraints:
-  //
-  //    (outA[0] > outB[0]) U
-  //    (outA[0] = outB[0], outA[1] > outA[1]) U
-  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] > outA[2]) U
-  //    ...
-  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] > outB[n-1])
-  //
-  // where `n` is the number of outputs.
-  // If `lexMin` is set, the complement inequality is used:
-  //
-  //    (outA[0] < outB[0]) U
-  //    (outA[0] = outB[0], outA[1] < outA[1]) U
-  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] < outA[2]) U
-  //    ...
-  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] < outB[n-1])
-  PresburgerSet result = PresburgerSet::getEmpty(compatibleSpace);
-  IntegerPolyhedron levelSet(
-      /*numReservedInequalities=*/1,
-      /*numReservedEqualities=*/pieceA.output.getNumOutputs(),
-      /*numReservedCols=*/space.getNumVars() + 1, space);
-  for (unsigned level = 0; level < pieceA.output.getNumOutputs(); ++level) {
-
-    // Create the expression `outA - outB` for this level.
-    SmallVector<MPInt, 8> subExpr = subtractExprs(
-        pieceA.output.getOutputExpr(level), pieceB.output.getOutputExpr(level));
-
-    if (lexMin) {
-      // For lexMin, we add an upper bound of -1:
-      //        outA - outB <= -1
-      //        outA <= outB - 1
-      //        outA < outB
-      levelSet.addBound(IntegerPolyhedron::BoundType::UB, subExpr, MPInt(-1));
-    } else {
-      // For lexMax, we add a lower bound of 1:
-      //        outA - outB >= 1
-      //        outA > outB + 1
-      //        outA > outB
-      levelSet.addBound(IntegerPolyhedron::BoundType::LB, subExpr, MPInt(1));
-    }
-
-    // Union the set with the result.
-    result.unionInPlace(levelSet);
-    // There is only 1 inequality in `levelSet`, so the index is always 0.
-    levelSet.removeInequality(0);
-    // Add equality `outA - outB == 0` for this level for next iteration.
-    levelSet.addEquality(subExpr);
-  }
-
-  // We then intersect `result` with the domain of pieceA and pieceB, to only
-  // tiebreak on the domain where both are defined.
+  PresburgerSet result = pieceA.output.getLexSet(comp, pieceB.output);
   result = result.intersect(pieceA.domain).intersect(pieceB.domain);
 
   return result;
 }
 
 PWMAFunction PWMAFunction::unionLexMin(const PWMAFunction &func) {
-  return unionFunction(func, tiebreakLex</*lexMin=*/true>);
+  return unionFunction(func, tiebreakLex</*comp=*/OrderingKind::LT>);
 }
 
 PWMAFunction PWMAFunction::unionLexMax(const PWMAFunction &func) {
-  return unionFunction(func, tiebreakLex</*lexMin=*/false>);
+  return unionFunction(func, tiebreakLex</*comp=*/OrderingKind::GT>);
 }
 
 void MultiAffineFunction::subtract(const MultiAffineFunction &other) {
   assert(space.isCompatible(other.space) &&
          "Spaces should be compatible for subtraction.");
 
   MultiAffineFunction copyOther = other;
   mergeDivs(copyOther);

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

PresburgerSpace PresburgerSpace::getDomainSpace() const {
   // TODO: Preserve identifiers here.
   return PresburgerSpace::getSetSpace(numDomain, numSymbols, numLocals);
 }
 
 PresburgerSpace PresburgerSpace::getRangeSpace() const {
   return PresburgerSpace::getSetSpace(numRange, numSymbols, numLocals);
 }
 
+PresburgerSpace PresburgerSpace::getSpaceWithoutLocals() const {
+  PresburgerSpace space = *this;
+  space.removeVarRange(VarKind::Local, 0, numLocals);
+  return space;
+}
+
 unsigned PresburgerSpace::getNumVarKind(VarKind kind) const {
   if (kind == VarKind::Domain)
     return getNumDomainVars();
   if (kind == VarKind::Range)
     return getNumRangeVars();
   if (kind == VarKind::Symbol)
     return getNumSymbolVars();
   if (kind == VarKind::Local)

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

TEST(PWMAFunction, unionLexMinComplex) {
   PWMAFunction result = parsePWMAF({
       {"(x, y) : (x == 1, y >= 0, y <= 1)", "(x, y) -> (-x, y)"},
       {"(x, y) : (x == 0, y >= 0, y <= 1)", "(x, y) -> (0, y - 1)"},
   });
 
   EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
   EXPECT_TRUE(func2.unionLexMin(func1).isEqual(result));
 }
+
+TEST(PWMAFunction, unionLexMinWithDivs) {
+  {
+    PWMAFunction func1 = parsePWMAF({
+        {"(x, y) : (x mod 5 == 0)", "(x, y) -> (x, 1)"},
+    });
+
+    PWMAFunction func2 = parsePWMAF({
+        {"(x, y) : (x mod 7 == 0)", "(x, y) -> (x + y, 2)"},
+    });
+
+    PWMAFunction result = parsePWMAF({
+        {"(x, y) : (x mod 5 == 0, x mod 7 >= 1)", "(x, y) -> (x, 1)"},
+        {"(x, y) : (x mod 7 == 0, x mod 5 >= 1)", "(x, y) -> (x + y, 2)"},
+        {"(x, y) : (x mod 5 == 0, x mod 7 == 0, y >= 0)", "(x, y) -> (x, 1)"},
+        {"(x, y) : (x mod 7 == 0, x mod 5 == 0, y <= -1)",
+         "(x, y) -> (x + y, 2)"},
+    });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+  }
+
+  {
+    PWMAFunction func1 = parsePWMAF({
+        {"(x) : (x >= 0, x <= 1000)", "(x) -> (x floordiv 16)"},
+    });
+
+    PWMAFunction func2 = parsePWMAF({
+        {"(x) : (x >= 0, x <= 1000)", "(x) -> ((x + 10) floordiv 17)"},
+    });
+
+    PWMAFunction result = parsePWMAF({
+        {"(x) : (x >= 0, x <= 1000, x floordiv 16 <= (x + 10) floordiv 17)",
+         "(x) -> (x floordiv 16)"},
+        {"(x) : (x >= 0, x <= 1000, x floordiv 16 >= (x + 10) floordiv 17 + 1)",
+         "(x) -> ((x + 10) floordiv 17)"},
+    });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+  }
+}