[MLIR][Presburger] subtract: support non-div locals

[arjunp]

Summary

Also added test cases. Also extend support for computeReprWithOnlyDivLocals from IntegerPolyhedron to IntegerRelation and PresburgerRelation.

Depends on D128736.

[Groverkss]

LGTM.

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

131

[arjunp]

Address kunwar's comment

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 namespace presburger {
 
 class IntegerRelation;
 class IntegerPolyhedron;
 class PresburgerSet;
+class PresburgerRelation;
 
 /// An IntegerRelation represents the set of points from a PresburgerSpace that
 /// satisfy a list of affine constraints. Affine constraints can be inequalities
 /// or equalities in the form:
 ///
 /// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n >= 0
 /// Equality  : c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n == 0
 ///
public:
   void applyDomain(const IntegerRelation &rel);
 
   /// Given a relation `rel`, apply the relation to the range of this relation.
   ///
   /// Formally, R1.applyRange(R2) is the same as R1.compose(R2) but we provide
   /// this for uniformity with `applyDomain`.
   void applyRange(const IntegerRelation &rel);
 
+  /// Compute an equivalent representation of the same set, such that all local
+  /// vars in all disjuncts have division representations. This representation
+  /// may involve local vars that correspond to divisions, and may also be a
+  /// union of convex disjuncts.
+  PresburgerRelation computeReprWithOnlyDivLocals() const;
+
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Checks all rows of equality/inequality constraints for trivial
   /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced
   /// after elimination. Returns true if an invalid constraint is found;
   /// false otherwise.
public:
   std::unique_ptr<IntegerPolyhedron> clone() const;
 
   /// Insert `num` variables of the specified kind at position `pos`.
   /// Positions are relative to the kind of variable. Return the absolute
   /// column position (i.e., not relative to the kind of variable) of the
   /// first added variable.
   unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1) override;
 
-  /// Compute an equivalent representation of the same set, such that all local
-  /// ids have division representations. This representation may involve
-  /// local ids that correspond to divisions, and may also be a union of convex
-  /// disjuncts.
-  PresburgerSet computeReprWithOnlyDivLocals() const;
-
   /// Compute the symbolic integer lexmin of the polyhedron.
   /// This finds, for every assignment to the symbols, the lexicographically
   /// minimum value attained by the dimensions. For example, the symbolic lexmin
   /// of the set
   ///
   /// (x, y)[a, b, c] : (a <= x, b <= x, x <= c)
   ///
   /// can be written as

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

public:
   ///
   /// In particular, removes all disjuncts which are subsets of other
   /// disjuncts in the union.
   PresburgerRelation coalesce() const;
 
   /// Check whether all local ids in all disjuncts have a div representation.
   bool hasOnlyDivLocals() const;
 
+  /// Compute an equivalent representation of the same relation, such that all
+  /// local ids in all disjuncts have division representations. This
+  /// representation may involve local ids that correspond to divisions, and may
+  /// also be a union of convex disjuncts.
+  PresburgerRelation computeReprWithOnlyDivLocals() const;
+
   /// Print the set's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Construct an empty PresburgerRelation with the specified number of
   /// dimension and symbols.
   explicit PresburgerRelation(const PresburgerSpace &space) : space(space) {

mlir/include/mlir/Analysis/Presburger/Simplex.h

 /// case where the row isn't violated, and then come back and run the case
 /// where it is.
 class SymbolicLexSimplex : public LexSimplexBase {
 public:
   /// `constraints` is the set for which the symbolic lexmin will be computed.
   /// `symbolDomain` is the set of values of the symbols for which the lexmin
   /// will be computed. `symbolDomain` should have a dim var for every symbol in
   /// `constraints`, and no other vars.
-  SymbolicLexSimplex(const IntegerPolyhedron &constraints,
+  SymbolicLexSimplex(const IntegerRelation &constraints,
                      const IntegerPolyhedron &symbolDomain)
       : SymbolicLexSimplex(constraints,
                            constraints.getVarKindOffset(VarKind::Symbol),
                            symbolDomain) {
     assert(constraints.getNumSymbolVars() == symbolDomain.getNumVars());
   }
 
   /// An overload to select some other subrange of ids as symbols for lexmin.
   /// The symbol ids are the range of ids with absolute index
   /// [symbolOffset, symbolOffset + symbolDomain.getNumVars())
   /// symbolDomain should only have dim ids.
-  SymbolicLexSimplex(const IntegerPolyhedron &constraints,
-                     unsigned symbolOffset,
+  SymbolicLexSimplex(const IntegerRelation &constraints, unsigned symbolOffset,
                      const IntegerPolyhedron &symbolDomain)
       : LexSimplexBase(/*nVar=*/constraints.getNumVars(), symbolOffset,
                        symbolDomain.getNumVars()),
         domainPoly(symbolDomain), domainSimplex(symbolDomain) {
     // TODO consider supporting this case. It amounts
     // to just returning the input constraints.
     assert(domainPoly.getNumVars() > 0 &&
            "there must be some non-symbols to optimize!");

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void IntegerRelation::truncate(const CountsSnapshot &counts) {
   truncateVarKind(VarKind::Domain, counts);
   truncateVarKind(VarKind::Range, counts);
   truncateVarKind(VarKind::Symbol, counts);
   truncateVarKind(VarKind::Local, counts);
   removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
   removeEqualityRange(counts.getNumEqs(), getNumEqualities());
 }
 
-PresburgerSet IntegerPolyhedron::computeReprWithOnlyDivLocals() const {
+PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const {
   // If there are no locals, we're done.
   if (getNumLocalVars() == 0)
-    return PresburgerSet(*this);
+    return PresburgerRelation(*this);
 
   // Move all the non-div locals to the end, as the current API to
   // SymbolicLexMin requires these to form a contiguous range.
   //
   // Take a copy so we can perform mutations.
-  IntegerPolyhedron copy = *this;
+  IntegerRelation copy = *this;
   std::vector<MaybeLocalRepr> reprs;
   copy.getLocalReprs(reprs);
 
   // Iterate through all the locals. The last `numNonDivLocals` are the locals
   // that have been scanned already and do not have division representations.
   unsigned numNonDivLocals = 0;
   unsigned offset = copy.getVarKindOffset(VarKind::Local);
   for (unsigned i = 0, e = copy.getNumLocalVars(); i < e - numNonDivLocals;) {
     if (!reprs[i]) {
       // Whenever we come across a local that does not have a division
       // representation, we swap it to the `numNonDivLocals`-th last position
       // and increment `numNonDivLocal`s. `reprs` also needs to be swapped.
       copy.swapVar(offset + i, offset + e - numNonDivLocals - 1);
       std::swap(reprs[i], reprs[e - numNonDivLocals - 1]);
       ++numNonDivLocals;
       continue;
     }
     ++i;
   }
 
   // If there are no non-div locals, we're done.
   if (numNonDivLocals == 0)
-    return PresburgerSet(*this);
+    return PresburgerRelation(*this);
 
   // We computeSymbolicIntegerLexMin by considering the non-div locals as
   // "non-symbols" and considering everything else as "symbols". This will
   // compute a function mapping assignments to "symbols" to the
   // lexicographically minimal valid assignment of "non-symbols", when a
   // satisfying assignment exists. It separately returns the set of assignments
   // to the "symbols" such that a satisfying assignment to the "non-symbols"
   // exists but the lexmin is unbounded. We basically want to find the set of

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

static SmallVector<int64_t, 8> getIneqCoeffsFromIdx(const IntegerRelation &rel,
   idx -= rel.getNumInequalities();
   ArrayRef<int64_t> eqCoeffs = rel.getEquality(idx / 2);
 
   if (idx % 2 == 0)
     return llvm::to_vector<8>(eqCoeffs);
   return getNegatedCoeffs(eqCoeffs);
 }
 
+PresburgerRelation PresburgerRelation::computeReprWithOnlyDivLocals() const {
+  if (hasOnlyDivLocals())
+    return *this;
+
+  // The result is just the union of the reprs of the disjuncts.
+  PresburgerRelation result(getSpace());
+  for (const IntegerRelation &disjunct : disjuncts)
+    result.unionInPlace(disjunct.computeReprWithOnlyDivLocals());
+  return result;
+}
+
 /// Return the set difference b \ s.
 ///
 /// In the following, U denotes union, /\ denotes intersection, \ denotes set
 /// difference and ~ denotes complement.
 ///
 /// Let s = (U_i s_i). We want  b \ (U_i s_i).
 ///
 /// Let s_i = /\_j s_ij, where each s_ij is a single inequality. To compute

 /// ignored.
 ///
 static PresburgerRelation getSetDifference(IntegerRelation b,
                                            const PresburgerRelation &s) {
   assert(b.getSpace().isCompatible(s.getSpace()) && "Spaces should match");
   if (b.isEmptyByGCDTest())
     return PresburgerRelation::getEmpty(b.getSpaceWithoutLocals());
 
+  if (!s.hasOnlyDivLocals())
+    return getSetDifference(b, s.computeReprWithOnlyDivLocals());
+
   // Remove duplicate divs up front here to avoid existing
   // divs disappearing in the call to mergeLocalVars below.
   b.removeDuplicateDivs();
 
   PresburgerRelation result =
       PresburgerRelation::getEmpty(b.getSpaceWithoutLocals());
   Simplex simplex(b);
 

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

TEST(SetTest, isEqual) {
   EXPECT_FALSE(universeSquare.isEqual(square));
   EXPECT_FALSE(rect.complement().isEqual(square.complement()));
 }
 
 void expectEqual(const PresburgerSet &s, const PresburgerSet &t) {
   EXPECT_TRUE(s.isEqual(t));
 }
 
+void expectEqual(const IntegerPolyhedron &s, const IntegerPolyhedron &t) {
+  EXPECT_TRUE(s.isEqual(t));
+}
+
 void expectEmpty(const PresburgerSet &s) { EXPECT_TRUE(s.isIntegerEmpty()); }
 
 TEST(SetTest, divisions) {
   // evens = {x : exists q, x = 2q}.
   PresburgerSet evens{parsePoly("(x) : (x - 2 * (x floordiv 2) == 0)")};
 
   //  odds = {x : exists q, x = 2q + 1}.
   PresburgerSet odds{parsePoly("(x) : (x - 2 * (x floordiv 2) - 1 == 0)")};
TEST(SetTest, divisionsDefByEq) {
   // even multiples of 3 = multiples of 6.
   expectEqual(multiples3.intersect(evens), multiples6);
 
   PresburgerSet evensDefByIneq{
       parsePoly("(x) : (x - 2 * (x floordiv 2) == 0)")};
   expectEqual(evens, PresburgerSet(evensDefByIneq));
 }
 
+TEST(SetTest, divisionNonDivLocals) {
+  // This is a tetrahedron with vertices at
+  // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 1000), and (1000, 1000, 1000).
+  //
+  // The only integer point in this is at (1000, 1000, 1000).
+  // We project this to the xy plane.
+  IntegerPolyhedron tetrahedron =
+      parsePolyAndMakeLocals("(x, y, z) : (y >= 0, z - y >= 0, 3000*x - 2998*y "
+                             "- 1000 - z >= 0, -1500*x + 1499*y + 1000 >= 0)",
+                             /*numLocals=*/1);
+
+  // This is a triangle with vertices at (1/3, 0), (2/3, 0) and (1000, 1000).
+  // The only integer point in this is at (1000, 1000).
+  //
+  // It also happens to be the projection of the above onto the xy plane.
+  IntegerPolyhedron triangle = parsePoly("(x,y) : (y >= 0, "
+                                         "3000 * x - 2999 * y - 1000 >= 0, "
+                                         "-3000 * x + 2998 * y + 2000 >= 0)");
+  EXPECT_TRUE(triangle.containsPoint({1000, 1000}));
+  EXPECT_FALSE(triangle.containsPoint({1001, 1001}));
+  expectEqual(triangle, tetrahedron);
+
+  convertSuffixDimsToLocals(triangle, 1);
+  IntegerPolyhedron line = parsePoly("(x) : (x - 1000 == 0)");
+  expectEqual(line, triangle);
+
+  // Triangle with vertices (0, 0), (5, 0), (15, 5).
+  // Projected on x, it becomes [0, 13] U {15} as it becomes too narrow towards
+  // the apex and so does not have have any integer point at x = 14.
+  // At x = 15, the apex is an integer point.
+  PresburgerSet triangle2{parsePolyAndMakeLocals("(x,y) : (y >= 0, "
+                                                 "x - 3*y >= 0, "
+                                                 "2*y - x + 5 >= 0)",
+                                                 /*numLocals=*/1)};
+  PresburgerSet zeroToThirteen{parsePoly("(x) : (13 - x >= 0, x >= 0)")};
+  PresburgerSet fifteen{parsePoly("(x) : (x - 15 == 0)")};
+  expectEqual(triangle2.subtract(zeroToThirteen), fifteen);
+}
+
 TEST(SetTest, subtractDuplicateDivsRegression) {
   // Previously, subtracting sets with duplicate divs might result in crashes
   // due to existing divs being removed when merging local ids, due to being
   // identified as being duplicates for the first time.
   IntegerPolyhedron setA(PresburgerSpace::getSetSpace(1));
   setA.addLocalFloorDiv({1, 0}, 2);
   setA.addLocalFloorDiv({1, 0, 0}, 2);
   EXPECT_TRUE(setA.isEqual(setA));

 
 // The last `numToProject` dims will be projected out, i.e., converted to
 // locals.
 void testComputeReprAtPoints(IntegerPolyhedron poly,
                              ArrayRef<SmallVector<int64_t, 4>> points,
                              unsigned numToProject) {
   poly.convertVarKind(VarKind::SetDim, poly.getNumDimVars() - numToProject,
                       poly.getNumDimVars(), VarKind::Local);
-  PresburgerSet repr = poly.computeReprWithOnlyDivLocals();
+  PresburgerRelation repr = poly.computeReprWithOnlyDivLocals();
   EXPECT_TRUE(repr.hasOnlyDivLocals());
   EXPECT_TRUE(repr.getSpace().isCompatible(poly.getSpace()));
   for (const SmallVector<int64_t, 4> &point : points) {
     EXPECT_EQ(poly.containsPointNoLocal(point).hasValue(),
               repr.containsPoint(point));
   }
 }
 
 void testComputeRepr(IntegerPolyhedron poly, const PresburgerSet &expected,
                      unsigned numToProject) {
   poly.convertVarKind(VarKind::SetDim, poly.getNumDimVars() - numToProject,
                       poly.getNumDimVars(), VarKind::Local);
-  PresburgerSet repr = poly.computeReprWithOnlyDivLocals();
+  PresburgerRelation repr = poly.computeReprWithOnlyDivLocals();
   EXPECT_TRUE(repr.hasOnlyDivLocals());
   EXPECT_TRUE(repr.getSpace().isCompatible(poly.getSpace()));
   EXPECT_TRUE(repr.isEqual(expected));
 }
 
 TEST(SetTest, computeReprWithOnlyDivLocals) {
   testComputeReprAtPoints(parsePoly("(x, y) : (x - 2*y == 0)"),
                           {{1, 0}, {2, 1}, {3, 0}, {4, 2}, {5, 3}},