[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

159

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

[arjunp]

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

159

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

159

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 CountsSnapshot stores the count of each IdKind, and also of
+  /// each constraint type. getCounts() returns a CountsSnapshot 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 CountsSnapshot 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 CountsSnapshot {
+  public:
+    CountsSnapshot(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;
+  };
+  CountsSnapshot getCounts() const;
+  void truncate(const CountsSnapshot &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);
 
+  using PresburgerSpace::truncateIdKind;
+  /// Truncate the ids to the number in the space of the specified
+  /// CountsSnapshot.
+  void truncateIdKind(IdKind kind, const CountsSnapshot &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/include/mlir/Analysis/Presburger/PresburgerSpace.h

public:
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   virtual unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
 
   /// Removes identifiers of the specified kind in the column range [idStart,
   /// idLimit). The range is relative to the kind of identifier.
   virtual void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
+  /// Truncate the ids of the specified kind to the specified number by dropping
+  /// some ids at the end. `num` must be less than the current number.
+  void truncateIdKind(IdKind kind, unsigned num);
+
   /// Returns true if both the spaces are equal i.e. if both spaces have the
   /// same number of identifiers of each kind (excluding Local Identifiers).
   bool isEqual(const PresburgerSpace &other) 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.

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::CountsSnapshot IntegerRelation::getCounts() const {
+  return {PresburgerLocalSpace(*this), getNumInequalities(),
+          getNumEqualities()};
+}
+
+void IntegerRelation::truncateIdKind(IdKind kind,
+                                     const CountsSnapshot &counts) {
+  truncateIdKind(kind, counts.getSpace().getNumIdKind(kind));
+}
+
+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());
+  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::CountsSnapshot 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

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
   else if (kind == IdKind::Range)
     numRange -= numIdsEliminated;
   else if (kind == IdKind::Symbol)
     numSymbols -= numIdsEliminated;
   else
     llvm_unreachable("PresburgerSpace does not support local identifiers!");
 }
 
+void PresburgerSpace::truncateIdKind(IdKind kind, unsigned num) {
+  unsigned curNum = getNumIdKind(kind);
+  assert(num <= curNum && "Can't truncate to more ids!");
+  removeIdRange(kind, num, curNum);
+}
+
 unsigned PresburgerLocalSpace::insertId(IdKind kind, unsigned pos,
                                         unsigned num) {
   if (kind == IdKind::Local) {
     numLocals += num;
     return getIdKindOffset(IdKind::Local) + pos;
   }
   return PresburgerSpace::insertId(kind, pos, num);
 }