[MLIR][Presburger] IntegerRelation: implement partial rollback support

[arjunp]

Summary

It is often necessary to "rollback" IntegerRelations to an earlier state. Although providing full rollback support is non-trivial, we really only need to support the case where the only changes made are to append ids or append constraints, and then rollback these additions. This patch adds support to rollback in such situations by recording the number of ids and constraints of each kind and providing support to truncate the IntegerRelation to those counts by removing appended ids and constraints. This already simplifies subtraction a little bit and will also be useful in the implementation of symbolic integer lexmin.

[Groverkss]

Can you write a patch description with the reasoning for this patch?

[Groverkss]

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

155

Can we have a better name than Counts? Maybe IntegerRelationState or something more descriptive?

[arjunp]

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

155

I named it counts to make it clear it's not really a proper thing storing the full state. Do you have a better suggestion that also conveys that it's not a proper snapshot?

[Groverkss]

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

155

Maybe CountState? I just find Counts not very descriptive.

[arjunp]

Move truncateIdKind to PresburgerSpace and change Counts -> CountsSnapshot

[Groverkss]

LGTM.

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

public:
   inline ArrayRef<int64_t> getEquality(unsigned idx) const {
     return equalities.getRow(idx);
   }
 
   inline ArrayRef<int64_t> getInequality(unsigned idx) const {
     return inequalities.getRow(idx);
   }
 
+  /// The struct Counts stores the count of each IdKind, and also of each
+  /// constraint type. getCounts() returns a Counts object describing the
+  /// current state of the IntegerRelation. truncate() truncates all ids of each
+  /// IdKind and all constraints of both kinds beyond the counts in the
+  /// specified Counts object. This can be used to achieve rudimentary
+  /// rollback support. As long as none of the existing constraints or ids are
+  /// disturbed, and only additional ids or constraints are added, this addition
+  /// can be rolled back using truncate.
+  struct Counts {
+  public:
+    Counts(const PresburgerLocalSpace &space, unsigned numIneqs,
+           unsigned numEqs)
+        : space(space), numIneqs(numIneqs), numEqs(numEqs) {}
+    const PresburgerLocalSpace &getSpace() const { return space; };
+    unsigned getNumIneqs() const { return numIneqs; }
+    unsigned getNumEqs() const { return numEqs; }
+
+  private:
+    PresburgerLocalSpace space;
+    unsigned numIneqs, numEqs;
+  };
+  Counts getCounts() const;
+  void truncate(const Counts &counts);
+
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. The coefficient columns
   /// corresponding to the added identifiers are initialized to zero. 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;
 
   /// Append `num` identifiers of the specified kind after the last identifier.
protected:
   /// IntegerRelation.
   virtual void printSpace(raw_ostream &os) const;
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.
   void removeIdRange(unsigned idStart, unsigned idLimit);
 
+  /// Truncate the ids of the specified kind to the specified count by dropping
+  /// some ids at the end.
+  void truncateIdKind(IdKind kind, unsigned num);
+  /// Truncate the ids to the number in the space of the specified Counts.
+  void truncateIdKind(IdKind kind, const Counts &counts);
+
   /// A parameter that controls detection of an unrealistic number of
   /// constraints. If the number of constraints is this many times the number of
   /// variables, we consider such a system out of line with the intended use
   /// case of IntegerRelation.
   // The rationale for 32 is that in the typical simplest of cases, an
   // identifier is expected to have one lower bound and one upper bound
   // constraint. With a level of tiling or a connection to another identifier
   // through a div or mod, an extra pair of bounds gets added. As a limit, we

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void removeConstraintsInvolvingIdRange(IntegerRelation &poly, unsigned begin,
   // next constraint to be tested is always i - 2.
   for (unsigned i = poly.getNumEqualities(); i > 0; i--)
     if (!rangeIsZero(poly.getEquality(i - 1).slice(begin, count)))
       poly.removeEquality(i - 1);
   for (unsigned i = poly.getNumInequalities(); i > 0; i--)
     if (!rangeIsZero(poly.getInequality(i - 1).slice(begin, count)))
       poly.removeInequality(i - 1);
 }
+
+IntegerRelation::Counts IntegerRelation::getCounts() const {
+  return {PresburgerLocalSpace(*this), getNumInequalities(),
+          getNumEqualities()};
+}
+
+void IntegerRelation::truncateIdKind(IdKind kind, unsigned num) {
+  unsigned curNum = getNumIdKind(kind);
+  assert(num <= curNum && "Can't truncate to more ids!");
+  removeIdRange(kind, num, curNum);
+}
+
+void IntegerRelation::truncateIdKind(IdKind kind, const Counts &counts) {
+  truncateIdKind(kind, counts.getSpace().getNumIdKind(kind));
+}
+
+void IntegerRelation::truncate(const Counts &counts) {
+  truncateIdKind(IdKind::Domain, counts);
+  truncateIdKind(IdKind::Range, counts);
+  truncateIdKind(IdKind::Symbol, counts);
+  truncateIdKind(IdKind::Local, counts);
+  removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
+  removeInequalityRange(counts.getNumEqs(), getNumEqualities());
+}
+
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned insertPos = PresburgerLocalSpace::insertId(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);
   return insertPos;
 }

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

if (i == s.getNumDisjuncts()) {
     return;
   }
   IntegerRelation sI = s.getDisjunct(i);
 
   // Below, we append some additional constraints and ids to b. We want to
   // rollback b to its initial state before returning, which we will do by
   // removing all constraints beyond the original number of inequalities
   // and equalities, so we store these counts first.
-  const unsigned bInitNumIneqs = b.getNumInequalities();
-  const unsigned bInitNumEqs = b.getNumEqualities();
-  const unsigned bInitNumLocals = b.getNumLocalIds();
+  const IntegerRelation::Counts bCounts = b.getCounts();
   // Similarly, we also want to rollback simplex to its original state.
   const unsigned initialSnapshot = simplex.getSnapshot();
 
   auto restoreState = [&]() {
-    b.removeIdRange(IdKind::Local, bInitNumLocals, b.getNumLocalIds());
-    b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
-    b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
+    b.truncate(bCounts);
     simplex.rollback(initialSnapshot);
   };
 
   // Automatically restore the original state when we return.
   auto stateRestorer = llvm::make_scope_exit(restoreState);
 
   // Find out which inequalities of sI correspond to division inequalities for
   // the local variables of sI.
for (MaybeLocalRepr &maybeInequality : repr) {
 
     assert(lb != ub &&
            "Upper and lower bounds must be different inequalities!");
     isDivInequality[lb] = true;
     isDivInequality[ub] = true;
   }
 
   unsigned offset = simplex.getNumConstraints();
-  unsigned numLocalsAdded = b.getNumLocalIds() - bInitNumLocals;
+  unsigned numLocalsAdded =
+      b.getNumLocalIds() - bCounts.getSpace().getNumLocalIds();
   simplex.appendVariable(numLocalsAdded);
 
   unsigned snapshotBeforeIntersect = simplex.getSnapshot();
   simplex.intersectIntegerRelation(sI);
 
   if (simplex.isEmpty()) {
     // b ^ s_i is empty, so b \ s_i = b. We move directly to i + 1.
     // We are ignoring level i completely, so we restore the state