[MLIR][Presburger] Support computing a representation of a set that only has locals that are divs

[arjunp]

Summary

This paves the way for integer-exact projection, and for supporting
non-division locals in subtraction, complement, and equality checks.

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

180–188

I think a better way to write this loop would be to write it in reverse. Swapping each non-div with the offset + numDivisions and incrementing numDivisions after each iteration.

180–188

Also, please write some comments explaining what this does. I find this somewhat non-trivial to understand.

1199–1201

Not for this patch, but I do not like that we compute the division representation each time we want it. This computation is very expensive. We should instead be storing these divisions separately with each set.

[arjunp]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

180–188

numDivisions does not get incremented on every iteration in that implementation. i is only incremented in some cases as well, same as here. Why are you saying that way is better?

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

180–188

Sorry for being unclear. I was just suggesting that I think it is more natural to write this loop in reverse. I find that loops working with divisions are easier to write when written in reverse. But since you are not using the division representation, this should be either way.

[arjunp]

Fix bug where the resulting set had a space incompatible with the space of the original set.

[arjunp]

Also addressed Kunwar's comments.

[arjunp]

Fix doc and add another test.

[Groverkss]

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

101–104

This is very confusing and I cannot understand what this does. Can you please write it more clearly?

758–760

I think writing about union of convex disjuncts is enough.

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

60

I would rather call them "extra" ids.

63

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

768

I think you should put a newline after this.

[arjunp]

Address kunwar's comments

[Groverkss]

LGTM.

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

 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 namespace presburger {
 
 class IntegerRelation;
 class IntegerPolyhedron;
+class PresburgerSet;
 
 /// 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:
   static bool classof(const IntegerRelation *cst) { return true; }
 
   // Clones this object.
   std::unique_ptr<IntegerRelation> clone() const;
 
   /// Returns a reference to the underlying space.
   const PresburgerSpace &getSpace() const { return space; }
 
+  /// Set the space to `oSpace`, which should have the same number of ids as
+  /// the current space.
+  void setSpace(const PresburgerSpace &oSpace);
+
+  /// Set the space to `oSpace`, which should not have any local ids.
+  /// `oSpace` can have fewer ids than the current space; in that case, the
+  /// the extra ids in `this` that are not accounted for by `oSpace` will be
+  /// considered as local ids. `oSpace` should not have more ids than the
+  /// current space; this will result in an assert failure.
+  void setSpaceExceptLocals(const PresburgerSpace &oSpace);
+
   /// Returns a copy of the space without locals.
   PresburgerSpace getSpaceWithoutLocals() const {
     return PresburgerSpace::getRelationSpace(space.getNumDomainIds(),
                                              space.getNumRangeIds(),
                                              space.getNumSymbolIds());
   }
 
   /// Appends constraints from `other` into `this`. This is equivalent to an
public:
   /// representation.
   ///
   /// The spaces of both relations should be compatible.
   ///
   /// Returns the number of non-merged local ids of `other`, i.e. the number of
   /// locals that have been added to `this`.
   unsigned mergeLocalIds(IntegerRelation &other);
 
+  /// Check whether all local ids have a division representation.
+  bool hasOnlyDivLocals() const;
+
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of dimensional identifiers immediately before
   /// the split become symbols, or some of the symbols immediately after the
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount) {
     space.setDimSymbolSeparation(newSymbolCount);
   }
 
public:
   std::unique_ptr<IntegerPolyhedron> clone() const;
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   unsigned insertId(IdKind 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:
   unsigned getNumLocalIds() const { return space.getNumLocalIds(); }
   unsigned getNumIds() const { return space.getNumIds(); }
 
   /// Return the number of disjuncts in the union.
   unsigned getNumDisjuncts() const;
 
   const PresburgerSpace &getSpace() const { return space; }
 
+  /// Set the space to `oSpace`. `oSpace` should not contain any local ids.
+  /// `oSpace` need not have the same number of ids as the current space;
+  /// it could have more or less. If it has less, the extra ids become
+  /// locals of the disjuncts. It can also have more, in which case the
+  /// disjuncts will have fewer locals. If its total number of ids
+  /// exceeds that of some disjunct, an assert failure will occur.
+  void setSpace(const PresburgerSpace &oSpace);
+
   /// Return a reference to the list of disjuncts.
   ArrayRef<IntegerRelation> getAllDisjuncts() const;
 
   /// Return the disjunct at the specified index.
   const IntegerRelation &getDisjunct(unsigned index) const;
 
   /// Mutate this set, turning it into the union of this set and the given
   /// disjunct.
public:
   Optional<uint64_t> computeVolume() const;
 
   /// Simplifies the representation of a PresburgerRelation.
   ///
   /// 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;
+
   /// 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

 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 id for every symbol in
   /// `constraints`, and no other ids.
   SymbolicLexSimplex(const IntegerPolyhedron &constraints,
                      const IntegerPolyhedron &symbolDomain)
-      : LexSimplexBase(constraints), domainPoly(symbolDomain),
-        domainSimplex(symbolDomain) {
-    assert(domainPoly.getNumIds() == constraints.getNumSymbolIds());
-    assert(domainPoly.getNumDimIds() == constraints.getNumSymbolIds());
+      : SymbolicLexSimplex(constraints,
+                           constraints.getIdKindOffset(IdKind::Symbol),
+                           symbolDomain) {
+    assert(constraints.getNumSymbolIds() == symbolDomain.getNumIds());
+  }
+
+  /// 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.getNumIds())
+  /// symbolDomain should only have dim ids.
+  SymbolicLexSimplex(const IntegerPolyhedron &constraints,
+                     unsigned symbolOffset,
+                     const IntegerPolyhedron &symbolDomain)
+      : LexSimplexBase(/*nVar=*/constraints.getNumIds(), symbolOffset,
+                       symbolDomain.getNumIds()),
+        domainPoly(symbolDomain), domainSimplex(symbolDomain) {
+    // TODO consider supporting this case. It amounts
+    // to just returning the input constraints.
+    assert(domainPoly.getNumIds() > 0 &&
+           "there must be some non-symbols to optimize!");
+    assert(domainPoly.getNumIds() == domainPoly.getNumDimIds());
+    intersectIntegerRelation(constraints);
   }
 
   /// The lexmin will be stored as a function `lexmin` from symbols to
   /// non-symbols in the result.
   ///
   /// For some values of the symbols, the lexmin may be unbounded.
   /// These parts of the symbol domain will be stored in `unboundedDomain`.
+  ///
+  /// The spaces of the sets in the result are compatible with the symbolDomain
+  /// passed in the SymbolicLexSimplex constructor.
   SymbolicLexMin computeSymbolicIntegerLexMin();
 
 private:
   /// Perform all pivots that do not require branching.
   ///
   /// Return failure if the tableau became empty, indicating that the polytope
   /// is always integer empty in the current symbol domain.
   /// Return success otherwise.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 std::unique_ptr<IntegerRelation> IntegerRelation::clone() const {
   return std::make_unique<IntegerRelation>(*this);
 }
 
 std::unique_ptr<IntegerPolyhedron> IntegerPolyhedron::clone() const {
   return std::make_unique<IntegerPolyhedron>(*this);
 }
 
+void IntegerRelation::setSpace(const PresburgerSpace &oSpace) {
+  assert(space.getNumIds() == oSpace.getNumIds() && "invalid space!");
+  space = oSpace;
+}
+
+void IntegerRelation::setSpaceExceptLocals(const PresburgerSpace &oSpace) {
+  assert(oSpace.getNumLocalIds() == 0 && "no locals should be present!");
+  assert(oSpace.getNumIds() <= getNumIds() && "invalid space!");
+  unsigned newNumLocals = getNumIds() - oSpace.getNumIds();
+  space = oSpace;
+  space.insertId(IdKind::Local, 0, newNumLocals);
+}
+
 void IntegerRelation::append(const IntegerRelation &other) {
   assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
 
   inequalities.reserveRows(inequalities.getNumRows() +
                            other.getNumInequalities());
   equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
 
   for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
void IntegerRelation::truncate(const CountsSnapshot &counts) {
   truncateIdKind(IdKind::Domain, counts);
   truncateIdKind(IdKind::Range, counts);
   truncateIdKind(IdKind::Symbol, counts);
   truncateIdKind(IdKind::Local, counts);
   removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
   removeEqualityRange(counts.getNumEqs(), getNumEqualities());
 }
 
+PresburgerSet IntegerPolyhedron::computeReprWithOnlyDivLocals() const {
+  // If there are no locals, we're done.
+  if (getNumLocalIds() == 0)
+    return PresburgerSet(*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;
+  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.getIdKindOffset(IdKind::Local);
+  for (unsigned i = 0, e = copy.getNumLocalIds(); 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.swapId(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);
+
+  // 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
+  // values of the "symbols" such that an assignment to the "non-symbols"
+  // exists, which is the union of the domain of the returned lexmin function
+  // and the returned set of assignments to the "symbols" that makes the lexmin
+  // unbounded.
+  SymbolicLexMin lexminResult =
+      SymbolicLexSimplex(copy, /*symbolOffset*/ 0,
+                         IntegerPolyhedron(PresburgerSpace::getSetSpace(
+                             /*numDims=*/copy.getNumIds() - numNonDivLocals)))
+          .computeSymbolicIntegerLexMin();
+  PresburgerSet result =
+      lexminResult.lexmin.getDomain().unionSet(lexminResult.unboundedDomain);
+
+  // The result set might lie in the wrong space -- all its ids are dims.
+  // Set it to the desired space and return.
+  PresburgerSpace space = getSpace();
+  space.removeIdRange(IdKind::Local, 0, getNumLocalIds());
+  result.setSpace(space);
+  return result;
+}
+
 SymbolicLexMin IntegerPolyhedron::findSymbolicIntegerLexMin() const {
   // Compute the symbolic lexmin of the dims and locals, with the symbols being
   // the actual symbols of this set.
   SymbolicLexMin result =
       SymbolicLexSimplex(*this, IntegerPolyhedron(PresburgerSpace::getSetSpace(
                                     /*numDims=*/getNumSymbolIds())))
           .computeSymbolicIntegerLexMin();
 
unsigned IntegerRelation::mergeLocalIds(IntegerRelation &other) {
 
   presburger::mergeLocalIds(*this, other, merge);
 
   // Since we do not remove duplicate divisions in relA, this is guranteed to be
   // non-negative.
   return relA.getNumLocalIds() - oldALocals;
 }
 
+bool IntegerRelation::hasOnlyDivLocals() const {
+  std::vector<MaybeLocalRepr> reprs;
+  getLocalReprs(reprs);
+  return llvm::all_of(reprs,
+                      [](const MaybeLocalRepr &repr) { return bool(repr); });
+}
+
 void IntegerRelation::removeDuplicateDivs() {
   std::vector<SmallVector<int64_t, 8>> divs;
   SmallVector<unsigned, 4> denoms;
 
   getLocalReprs(divs, denoms);
   auto merge = [this](unsigned i, unsigned j) -> bool {
     eliminateRedundantLocalId(i, j);
     return true;

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 using namespace mlir;
 using namespace presburger;
 
 PresburgerRelation::PresburgerRelation(const IntegerRelation &disjunct)
     : space(disjunct.getSpaceWithoutLocals()) {
   unionInPlace(disjunct);
 }
 
+void PresburgerRelation::setSpace(const PresburgerSpace &oSpace) {
+  assert(space.getNumLocalIds() == 0 && "no locals should be present");
+  space = oSpace;
+  for (IntegerRelation &disjunct : disjuncts)
+    disjunct.setSpaceExceptLocals(space);
+}
+
 unsigned PresburgerRelation::getNumDisjuncts() const {
   return disjuncts.size();
 }
 
 ArrayRef<IntegerRelation> PresburgerRelation::getAllDisjuncts() const {
   return disjuncts;
 }
 
LogicalResult SetCoalescer::coalescePair(unsigned i, unsigned j) {
 
   return failure();
 }
 
 PresburgerRelation PresburgerRelation::coalesce() const {
   return SetCoalescer(*this).coalesce();
 }
 
+bool PresburgerRelation::hasOnlyDivLocals() const {
+  return llvm::all_of(disjuncts, [](const IntegerRelation &rel) {
+    return rel.hasOnlyDivLocals();
+  });
+}
+
 void PresburgerRelation::print(raw_ostream &os) const {
   os << "Number of Disjuncts: " << getNumDisjuncts() << "\n";
   for (const IntegerRelation &disjunct : disjuncts) {
     disjunct.print(os);
     os << '\n';
   }
 }
 

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

expectComputedVolumeIsValidOverapprox(unbounded, /*trueVolume=*/{},
                                         /*resultBound=*/{});
 
   // Union of unbounded with bounded is unbounded.
   expectComputedVolumeIsValidOverapprox(unbounded.unionSet(diamond),
                                         /*trueVolume=*/{},
                                         /*resultBound=*/{});
 }
 
+// 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.convertIdKind(IdKind::SetDim, poly.getNumDimIds() - numToProject,
+                     poly.getNumDimIds(), IdKind::Local);
+  PresburgerSet 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.convertIdKind(IdKind::SetDim, poly.getNumDimIds() - numToProject,
+                     poly.getNumDimIds(), IdKind::Local);
+  PresburgerSet 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}},
+                          /*numToProject=*/0);
+  testComputeReprAtPoints(parsePoly("(x, e) : (x - 2*e == 0)"),
+                          {{1}, {2}, {3}, {4}, {5}}, /*numToProject=*/1);
+
+  // Tests to check that the space is preserved.
+  testComputeReprAtPoints(parsePoly("(x, y)[z, w] : ()"), {},
+                          /*numToProject=*/1);
+  testComputeReprAtPoints(parsePoly("(x, y)[z, w] : (z - (w floordiv 2) == 0)"),
+                          {},
+                          /*numToProject=*/1);
+
+  // Bezout's lemma: if a, b are constants,
+  // the set of values that ax + by can take is all multiples of gcd(a, b).
+  testComputeRepr(
+      parsePoly("(x, e, f) : (x - 15*e - 21*f == 0)"),
+      PresburgerSet(parsePoly({"(x) : (x - 3*(x floordiv 3) == 0)"})),
+      /*numToProject=*/2);
+}
+
 TEST(SetTest, subtractOutputSizeRegression) {
   PresburgerSet set1 =
       parsePresburgerSetFromPolyStrings(1, {"(i) : (i >= 0, 10 - i >= 0)"});
   PresburgerSet set2 =
       parsePresburgerSetFromPolyStrings(1, {"(i) : (i - 5 >= 0)"});
 
   PresburgerSet set3 =
       parsePresburgerSetFromPolyStrings(1, {"(i) : (i >= 0, 4 - i >= 0)"});