[MLIR] PresburgerSet: support divisions in subtract, intersect, and equality checks

[arjunp]

Summary

Depends on D106662.

[grosser]

mlir/include/mlir/Analysis/AffineStructures.h

138

Typo.

[arjunp]

Fixed typo.

[arjunp]

Fix build by making removeId overload with unspecified kind public.

[arjunp]

Personally, I believe that the removeId overload where the identifier kind is not specified should not be public. I think the interface should only be on a per-kind basis rather than exposing the implementation detail of the ordering among kinds.

The build failure occurs in LoopUtils.cpp. I would be happy to change it to instead call the removeId overload that takes the kind to be used as an argument. But I don't understand why these lines mark all but one variable as a symbol but comment that it marks all variables as symbols.

[bondhugula]

Please add a commit summary. Is this rebased on the current master? A lot of the relevant files have changed I think.

mlir/include/mlir/Analysis/AffineStructures.h

391

Missing doc comments.

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

165

Triple /// comments please - here and elsewhere.

[arjunp]

Splitting this into mutliple patches.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// checking for invalid constraints. Returns true if the GCD test fails for
   /// any equality, or if any invalid constraints are discovered on any row.
   /// Returns false otherwise.
   bool isEmpty() const;
 
   /// Runs the GCD test on all equality constraints. Returns 'true' if this test
   /// fails on any equality. Returns 'false' otherwise.
   /// This test can be used to disprove the existence of a solution. If it
-  /// returns true, no integer solution to the equality constraints can exist.
+  /// returns true, no integer solution to the equality constraints can exist.gg
   bool isEmptyByGCDTest() const;
 
   /// Returns true if the set of constraints is found to have no solution,
   /// false if a solution exists. Uses the same algorithm as findIntegerSample.
   bool isIntegerEmpty() const;
 
   // Returns a matrix where each row is a vector along which the polytope is
   // bounded. The span of the returned vectors is guaranteed to contain all
public:
   /// otherwise.
   bool containsId(Value id) const;
 
   /// Swap the posA^th identifier with the posB^th identifier.
   void swapId(unsigned posA, unsigned posB);
 
   // Add identifiers of the specified kind - specified positions are relative to
   // the kind of identifier. The coefficient column corresponding to the added
-  // identifier is initialized to zero. 'id' is the Value corresponding to the
-  // identifier that can optionally be provided.
+  // identifier is initialized to zero. 'ids' are the Values corresponding to
+  // the identifiers which can optionally be provided.
   void addDimId(unsigned pos, Value id = nullptr);
   void addSymbolId(unsigned pos, Value id = nullptr);
   void addLocalId(unsigned pos);
   void addId(IdKind kind, unsigned pos, Value id = nullptr);
 
+  void addDimIds(unsigned pos, unsigned count, ArrayRef<Value> newIds = {});
+  void addSymbolIds(unsigned pos, unsigned count, ArrayRef<Value> newIds = {});
+  void addLocalIds(unsigned pos, unsigned count);
+  void addIds(IdKind kind, unsigned pos, unsigned count,
+              ArrayRef<Value> newIds = {});
+
+  void removeId(IdKind kind, unsigned pos);
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
+
   /// Add the specified values as a dim or symbol id depending on its nature, if
   /// it already doesn't exist in the system. `id' has to be either a terminal
   /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
   /// Note: If `allowNonTerminal`, any symbol (incl. potentially non-terminal
   /// ones) is allowed.
   void addInductionVarOrTerminalSymbol(Value id, bool allowNonTerminal = false);
 
-  /// Composes the affine value map with this FlatAffineConstrains, adding the
+  /// Composes the affine value map with this FlatAffineConstraints, adding the
   /// results of the map as dimensions at the front [0, vMap->getNumResults())
   /// and with the dimensions set to the equalities specified by the value map.
   /// Returns failure if the composition fails (when vMap is a semi-affine map).
   /// The vMap's operand Value's are used to look up the right positions in
   /// the FlatAffineConstraints with which to associate. The dimensional and
   /// symbolic operands of vMap should match 1:1 (in the same order) with those
   /// of this constraint system, but the latter could have additional trailing
   /// operands.
public:
   // TODO: deal with integer exactness when necessary - can return a value to
   // mark exactness for example.
   void projectOut(unsigned pos, unsigned num);
   inline void projectOut(unsigned pos) { return projectOut(pos, 1); }
 
   /// Projects out the identifier that is associate with Value .
   void projectOut(Value id);
 
-  /// Removes the specified identifier from the system.
-  void removeId(unsigned pos);
-
   void removeEquality(unsigned pos);
   void removeInequality(unsigned pos);
 
+  void removeEqualityRange(unsigned start, unsigned end);
+  void removeInequalityRange(unsigned start, unsigned end);
+
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of trailing dimensional identifiers becomes
   /// symbols, or some of the leading symbols become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   /// Changes all symbol identifiers which are loop IVs to dim identifiers.
   void convertLoopIVSymbolsToDims();
 
public:
 
   // Removes all equalities and inequalities.
   void clearConstraints();
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
+  /// Return the index at which the specified kind of id starts.
+  unsigned getIdOffset(IdKind kind) const;
+
+  /// Assert that `value` is at most the number of ids of the specified kind.
+  void assertAtMostNumKind(unsigned value, IdKind kind) const;
+
   /// Returns false if the fields corresponding to various identifier counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
   bool hasConsistentState() const;
 
   /// 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;
private:
   /// i.e.,
   ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
   /// fast method (linear in the number of coefficients).
   void gcdTightenInequalities();
 
   /// Normalized each constraints by the GCD of its coefficients.
   void normalizeConstraintsByGCD();
 
+  /// Removes the specified identifier from the system.
+  void removeId(unsigned pos);
+
   /// 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);
 
   /// Total number of identifiers.
   unsigned numIds;
 

mlir/include/mlir/Analysis/Presburger/Matrix.h

public:
   void addToColumn(unsigned sourceColumn, unsigned targetColumn, int64_t scale);
 
   /// Negate the specified column.
   void negateColumn(unsigned column);
 
   /// Resize the matrix to the specified dimensions. If a dimension is smaller,
   /// the values are truncated; if it is bigger, the new values are default
   /// initialized.
+  void resize(unsigned newNRows, unsigned newNColumns);
   void resizeVertically(unsigned newNRows);
 
   /// Add an extra row at the bottom of the matrix and return its position.
   unsigned appendExtraRow();
 
   /// Print the matrix.
   void print(raw_ostream &os) const;
   void dump() const;

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

public:
   /// Returns the number of constraints in the tableau.
   unsigned numConstraints() const;
 
   /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding equality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
   void addEquality(ArrayRef<int64_t> coeffs);
 
+  // Add new variables to the end of the list of variables.
+  void appendVariable();
+  void appendVariables(unsigned count);
+
   /// Mark the tableau as being empty.
   void markEmpty();
 
   /// Get a snapshot of the current state. This is used for rolling back.
   unsigned getSnapshot() const;
 
   /// Rollback to a snapshot. This invalidates all later snapshots.
   void rollback(unsigned snapshot);
private:
 
   /// Add a new row to the tableau and the associated data structures.
   unsigned addRow(ArrayRef<int64_t> coeffs);
 
   /// Normalize the given row by removing common factors between the numerator
   /// and the denominator.
   void normalizeRow(unsigned row);
 
-  /// Swap the two rows in the tableau and associated data structures.
+  /// Swap the two rows/columns in the tableau and associated data structures.
   void swapRows(unsigned i, unsigned j);
+  void swapColumns(unsigned i, unsigned j);
 
   /// Restore the unknown to a non-negative sample value.
   ///
   /// Returns true if the unknown was successfully restored to a non-negative
   /// sample value, false otherwise.
   LogicalResult restoreRow(Unknown &u);
 
   /// Compute the maximum or minimum of the specified Unknown, depending on
   /// direction. The specified unknown may be pivoted. If the unknown is
   /// restricted, it will have a non-negative sample value on return.
   /// Should not be called if the Simplex is empty.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
   Optional<Fraction> computeOptimum(Direction direction, Unknown &u);
 
   /// Mark the specified unknown redundant. This operation is added to the undo
   /// log and will be undone by rollbacks. The specified unknown must be in row
   /// orientation.
   void markRowRedundant(Unknown &u);
 
   /// Enum to denote operations that need to be undone during rollback.
   enum class UndoLogEntry {
     RemoveLastConstraint,
+    RemoveLastVariable,
     UnmarkEmpty,
     UnmarkLastRedundant
   };
 
   /// Undo the operation represented by the log entry.
   void undo(UndoLogEntry entry);
 
   /// Find a row that can be used to pivot the column in the specified

mlir/include/mlir/Analysis/PresburgerSet.h

public:
   void print(raw_ostream &os) const;
   void dump() const;
 
   /// Return the complement of this set. Computing the complement of a set
   /// containing divisions is not yet supported.
   PresburgerSet complement() const;
 
   /// Return the set difference of this set and the given set, i.e.,
-  /// return `this \ set`. Subtracting when either set contains divisions is not
-  /// yet supported.
+  /// return `this \ set`. All local variables in `set` must correspond
+  /// to floor divisions, otherwise this function's behaviour is undefined.
   PresburgerSet subtract(const PresburgerSet &set) const;
 
   /// Return true if this set is equal to the given set, and false otherwise.
   /// Checking equality when either set contains divisions is not yet supported.
   bool isEqual(const PresburgerSet &set) const;
 
   /// Return a universe set of the specified type that contains all points.
   static PresburgerSet getUniverse(unsigned nDim = 0, unsigned nSym = 0);

mlir/lib/Analysis/AffineStructures.cpp

 void FlatAffineConstraints::addDimId(unsigned pos, Value id) {
   addId(IdKind::Dimension, pos, id);
 }
 
 void FlatAffineConstraints::addSymbolId(unsigned pos, Value id) {
   addId(IdKind::Symbol, pos, id);
 }
 
-/// Adds a dimensional identifier. The added column is initialized to
-/// zero.
 void FlatAffineConstraints::addId(IdKind kind, unsigned pos, Value id) {
+  addIds(IdKind::Local, pos, 1, id);
+}
+
+void FlatAffineConstraints::addLocalIds(unsigned pos, unsigned count) {
+  addIds(IdKind::Local, pos, count);
+}
+
+void FlatAffineConstraints::addDimIds(unsigned pos, unsigned count,
+                                      ArrayRef<Value> newIds) {
+  addIds(IdKind::Dimension, pos, count, newIds);
+}
+
+void FlatAffineConstraints::addSymbolIds(unsigned pos, unsigned count,
+                                         ArrayRef<Value> newIds) {
+  addIds(IdKind::Symbol, pos, count, newIds);
+}
+
+unsigned FlatAffineConstraints::getIdOffset(IdKind kind) const {
   if (kind == IdKind::Dimension)
-    assert(pos <= getNumDimIds());
+    return 0;
+  if (kind == IdKind::Symbol)
+    return getNumDimIds();
+  if (kind == IdKind::Local)
+    return getNumDimIds() + getNumSymbolIds();
+  llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
+}
+
+void FlatAffineConstraints::assertAtMostNumKind(unsigned val,
+                                                IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    assert(val <= getNumDimIds());
   else if (kind == IdKind::Symbol)
-    assert(pos <= getNumSymbolIds());
+    assert(val <= getNumSymbolIds());
+  else if (kind == IdKind::Local)
+    assert(val <= getNumLocalIds());
   else
-    assert(pos <= getNumLocalIds());
-
-  int absolutePos;
-  if (kind == IdKind::Dimension) {
-    absolutePos = pos;
-    numDims++;
-  } else if (kind == IdKind::Symbol) {
-    absolutePos = pos + getNumDimIds();
-    numSymbols++;
-  } else {
-    absolutePos = pos + getNumDimIds() + getNumSymbolIds();
-  }
-  numIds++;
+    llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
 
-  inequalities.insertColumn(absolutePos);
-  equalities.insertColumn(absolutePos);
+/// Adds dimensional identifiers. The added columns are initialized to
+/// zero.
+void FlatAffineConstraints::addIds(IdKind kind, unsigned pos, unsigned count,
+                                   ArrayRef<Value> newIds) {
+  assertAtMostNumKind(pos, kind);
+
+  unsigned absolutePos = pos + getIdOffset(kind);
+  if (kind == IdKind::Dimension)
+    numDims += count;
+  else if (kind == IdKind::Symbol)
+    numSymbols += count;
+  numIds += count;
 
-  // If an 'id' is provided, insert it; otherwise use None.
-  if (id)
-    ids.insert(ids.begin() + absolutePos, id);
+  inequalities.insertColumns(absolutePos, count);
+  equalities.insertColumns(absolutePos, count);
+
+  // If ids are provided, insert them; otherwise use None.
+  if (newIds.empty())
+    ids.insert(ids.begin() + absolutePos, count, None);
   else
-    ids.insert(ids.begin() + absolutePos, None);
+    ids.insert(ids.begin() + absolutePos, newIds.begin(), newIds.end());
   assert(ids.size() == getNumIds());
 }
 
+void FlatAffineConstraints::removeId(IdKind kind, unsigned pos) {
+  removeIdRange(kind, pos, pos + 1);
+}
+
+void FlatAffineConstraints::removeIdRange(IdKind kind, unsigned idStart,
+                                          unsigned idLimit) {
+  assert(idStart <= idLimit);
+  assertAtMostNumKind(idLimit, kind);
+  if (idStart == idLimit)
+    return;
+
+  removeIdRange(getIdOffset(kind) + idStart, getIdOffset(kind) + idLimit);
+}
+
 /// Checks if two constraint systems are in the same space, i.e., if they are
 /// associated with the same set of identifiers, appearing in the same order.
 static bool areIdsAligned(const FlatAffineConstraints &a,
                           const FlatAffineConstraints &b) {
   return a.getNumDimIds() == b.getNumDimIds() &&
          a.getNumSymbolIds() == b.getNumSymbolIds() &&
          a.getNumIds() == b.getNumIds() && a.getIds().equals(b.getIds());
 }

 void FlatAffineConstraints::removeEquality(unsigned pos) {
   equalities.removeRow(pos);
 }
 
 void FlatAffineConstraints::removeInequality(unsigned pos) {
   inequalities.removeRow(pos);
 }
 
+void FlatAffineConstraints::removeEqualityRange(unsigned begin, unsigned end) {
+  equalities.removeRows(begin, end - begin);
+}
+
+void FlatAffineConstraints::removeInequalityRange(unsigned begin,
+                                                  unsigned end) {
+  inequalities.removeRows(begin, end - begin);
+}
+
 /// Finds an equality that equates the specified identifier to a constant.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the
 /// symbols is also treated like a constant. Returns -1 if such an equality
 /// could not be found.
 static int findEqualityToConstant(const FlatAffineConstraints &cst,
                                   unsigned pos, bool symbolic = false) {
   assert(pos < cst.getNumIds() && "invalid position");

mlir/lib/Analysis/Presburger/Matrix.cpp

void Matrix::reserveRows(unsigned rows) {
   data.reserve(rows * nReservedColumns);
 }
 
 unsigned Matrix::appendExtraRow() {
   resizeVertically(nRows + 1);
   return nRows - 1;
 }
 
+void Matrix::resize(unsigned newNRows, unsigned newNColumns) {
+  if (newNColumns < nColumns)
+    removeColumns(newNColumns, nColumns - newNColumns);
+  if (newNColumns > nColumns)
+    insertColumns(nColumns, newNColumns - nColumns);
+  resizeVertically(newNRows);
+}
+
 void Matrix::resizeVertically(unsigned newNRows) {
   nRows = newNRows;
   data.resize(nRows * nReservedColumns);
 }
 
 void Matrix::swapRows(unsigned row, unsigned otherRow) {
   assert((row < getNumRows() && otherRow < getNumRows()) &&
          "Given row out of bounds");

mlir/lib/Analysis/Presburger/Simplex.cpp

void Simplex::swapRows(unsigned i, unsigned j) {
   if (i == j)
     return;
   tableau.swapRows(i, j);
   std::swap(rowUnknown[i], rowUnknown[j]);
   unknownFromRow(i).pos = i;
   unknownFromRow(j).pos = j;
 }
 
+void Simplex::swapColumns(unsigned i, unsigned j) {
+  if (i == j)
+    return;
+  tableau.swapColumns(i, j);
+  std::swap(colUnknown[i], colUnknown[j]);
+  unknownFromColumn(i).pos = i;
+  unknownFromColumn(j).pos = j;
+}
+
 /// Mark this tableau empty and push an entry to the undo stack.
 void Simplex::markEmpty() {
   undoLog.push_back(UndoLogEntry::UnmarkEmpty);
   empty = true;
 }
 
 /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
 /// is the current number of variables, then the corresponding inequality is
if (entry == UndoLogEntry::RemoveLastConstraint) {
     // tableau.
     swapRows(constraint.pos, nRow - 1);
     // It is not strictly necessary to shrink the tableau, but for now we
     // maintain the invariant that the tableau has exactly nRow rows.
     tableau.resizeVertically(nRow - 1);
     nRow--;
     rowUnknown.pop_back();
     con.pop_back();
+  } else if (entry == UndoLogEntry::RemoveLastVariable) {
+    assert(var.back().orientation == Orientation::Column &&
+           "Variable to be removed must be in column orientation!");
+    swapColumns(var.back().pos, nCol - 1);
+    tableau.resize(nRow, nCol - 1);
+    var.pop_back();
+    colUnknown.pop_back();
+    nCol--;
   } else if (entry == UndoLogEntry::UnmarkEmpty) {
     empty = false;
   } else if (entry == UndoLogEntry::UnmarkLastRedundant) {
     nRedundant--;
   }
 }
 
 /// Rollback to the specified snapshot.
 ///
 /// We undo all the log entries until the log size when the snapshot was taken
 /// is reached.
 void Simplex::rollback(unsigned snapshot) {
   while (undoLog.size() > snapshot) {
     undo(undoLog.back());
     undoLog.pop_back();
   }
 }
 
+void Simplex::appendVariable() {
+  undoLog.emplace_back(UndoLogEntry::RemoveLastVariable);
+  nCol++;
+  tableau.resize(nRow, nCol);
+  var.emplace_back(Orientation::Column, /*restricted=*/false, /*pos=*/nCol - 1);
+  colUnknown.push_back(var.size() - 1);
+}
+
+void Simplex::appendVariables(unsigned count) {
+  for (unsigned i = 0; i < count; ++i)
+    appendVariable();
+}
+
 /// Add all the constraints from the given FlatAffineConstraints.
 void Simplex::intersectFlatAffineConstraints(const FlatAffineConstraints &fac) {
   assert(fac.getNumIds() == numVariables() &&
          "FlatAffineConstraints must have same dimensionality as simplex");
   for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
     addInequality(fac.getInequality(i));
   for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
     addEquality(fac.getEquality(i));

mlir/lib/Analysis/PresburgerSet.cpp

PresburgerSet PresburgerSet::getUniverse(unsigned nDim, unsigned nSym) {
   result.unionFACInPlace(FlatAffineConstraints::getUniverse(nDim, nSym));
   return result;
 }
 
 PresburgerSet PresburgerSet::getEmptySet(unsigned nDim, unsigned nSym) {
   return PresburgerSet(nDim, nSym);
 }
 
+// Let a and b have n and m locals respectively.
+// Add m locals to a, after a's locals.
+// Also add n locals to b, before b's locals.
+void addLocalsToEachOther(FlatAffineConstraints &a, FlatAffineConstraints &b) {
+  unsigned initANumLocals = a.getNumLocalIds();
+  a.addLocalIds(a.getNumLocalIds(), b.getNumLocalIds());
+  // The number of locals of `a` has changed.
+  // We want to add as many locals to `b` as `a` initially had.
+  b.addLocalIds(0, initANumLocals);
+}
+
 // Return the intersection of this set with the given set.
 //
 // We directly compute (S_1 or S_2 ...) and (T_1 or T_2 ...)
 // as (S_1 and T_1) or (S_1 and T_2) or ...
+//
+// If S_i or T_j have local variables, then S_i and T_j contains the local
+// variables of both.
 PresburgerSet PresburgerSet::intersect(const PresburgerSet &set) const {
   assertDimensionsCompatible(set, *this);
 
   PresburgerSet result(nDim, nSym);
   for (const FlatAffineConstraints &csA : flatAffineConstraints) {
     for (const FlatAffineConstraints &csB : set.flatAffineConstraints) {
-      FlatAffineConstraints intersection(csA);
-      intersection.append(csB);
-      if (!intersection.isEmpty())
-        result.unionFACInPlace(std::move(intersection));
+      FlatAffineConstraints csACopy = csA, csBCopy = csB;
+      addLocalsToEachOther(csACopy, csBCopy);
+      csACopy.append(std::move(csBCopy));
+      if (!csACopy.isEmpty())
+        result.unionFACInPlace(std::move(csACopy));
     }
   }
   return result;
 }
 
 /// Return `coeffs` with all the elements negated.
 static SmallVector<int64_t, 8> getNegatedCoeffs(ArrayRef<int64_t> coeffs) {
   SmallVector<int64_t, 8> negatedCoeffs;
   negatedCoeffs.reserve(coeffs.size());
   for (int64_t coeff : coeffs)
     negatedCoeffs.emplace_back(-coeff);
   return negatedCoeffs;
 }
 
 /// Return the complement of the given inequality.
 ///
 /// The complement of a_1 x_1 + ... + a_n x_ + c >= 0 is
-/// a_1 x_1 + ... + a_n x_ + c < 0, i.e., -a_1 x_1 - ... - a_n x_ - c - 1 >= 0.
+/// a_1 x_1 + ... + a_n x_ + c < 0, i.e., -a_1 x_1 - ... - a_n x_ - c - 1 >= 0,
+/// since all the variables are constrained to be integers.
 static SmallVector<int64_t, 8> getComplementIneq(ArrayRef<int64_t> ineq) {
   SmallVector<int64_t, 8> coeffs;
   coeffs.reserve(ineq.size());
   for (int64_t coeff : ineq)
     coeffs.emplace_back(-coeff);
   --coeffs.back();
   return coeffs;
 }
 
 /// Return the set difference b \ s and accumulate the result into `result`.
 /// `simplex` must correspond to b.
 ///
-/// In the following, V denotes union, ^ denotes intersection, \ denotes set
+/// In the following, U denotes union, ^ denotes intersection, \ denotes set
 /// difference and ~ denotes complement.
-/// Let b be the FlatAffineConstraints and s = (V_i s_i) be the set. We want
-/// b \ (V_i s_i).
+/// Let b be the FlatAffineConstraints and s = (U_i s_i) be the set. We want
+/// b \ (U_i s_i).
 ///
 /// Let s_i = ^_j s_ij, where each s_ij is a single inequality. To compute
 /// b \ s_i = b ^ ~s_i, we partition s_i based on the first violated inequality:
-/// ~s_i = (~s_i1) V (s_i1 ^ ~s_i2) V (s_i1 ^ s_i2 ^ ~s_i3) V ...
-/// And the required result is (b ^ ~s_i1) V (b ^ s_i1 ^ ~s_i2) V ...
-/// We recurse by subtracting V_{j > i} S_j from each of these parts and
+/// ~s_i = (~s_i1) U (s_i1 ^ ~s_i2) U (s_i1 ^ s_i2 ^ ~s_i3) U ...
+/// And the required result is (b ^ ~s_i1) U (b ^ s_i1 ^ ~s_i2) U ...
+/// We recurse by subtracting U_{j > i} S_j from each of these parts and
 /// returning the union of the results. Each equality is handled as a
 /// conjunction of two inequalities.
 ///
+/// Note that the same approach works even if an inequality involves a floor
+/// division. For example, The complement of x <= 7*floor(x/7) is still
+/// x > 7*floor(x/7). Since b \ s_i contains the inequalities of both b and s_i
+/// (or the complements of those inequalities), b \ s_i may contain the
+/// divisions present in both b and s_i. Therefore, we need to add the local
+/// division variables of both b and s_i to each part in the result. This means
+/// adding the local variables of both b and s_i, as well as the corresponding
+/// division inequalities to each part. Since the division inequalities are
+/// added to each part, we can skip the parts where the complement of any
+/// division inequality is added, as these parts will become empty anyway.
+///
 /// As a heuristic, we try adding all the constraints and check if simplex
-/// says that the intersection is empty. Also, in the process we find out that
-/// some constraints are redundant. These redundant constraints are ignored.
+/// says that the intersection is empty. If it is, then subtracting this FAC is
+/// a no-op and we just skip it. Also, in the process we find out that some
+/// constraints are redundant. These redundant constraints are ignored.
+///
+/// b and simplex are callee saved, i.e., their values on return are
+/// semantically equivalent to their values when the function is called.
 static void subtractRecursively(FlatAffineConstraints &b, Simplex &simplex,
                                 const PresburgerSet &s, unsigned i,
                                 PresburgerSet &result) {
   if (i == s.getNumFACs()) {
     result.unionFACInPlace(b);
     return;
   }
-  const FlatAffineConstraints &sI = s.getFlatAffineConstraints(i);
-  assert(sI.getNumLocalIds() == 0 &&
-         "Subtracting sets with divisions is not yet supported!");
+  FlatAffineConstraints sI = s.getFlatAffineConstraints(i);
+  unsigned bInitNumLocals = b.getNumLocalIds();
+
+  // Find out which inequalities of sI correspond corresponding to division
+  // inequalities for the local variables of sI.
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
+      sI.getNumLocalIds());
+  sI.getLocalReprLbUbPairs(repr);
+
+  // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
+  // sI's locals.
+  addLocalsToEachOther(b, sI);
+
+  llvm::SmallBitVector isDivInequality(sI.getNumInequalities());
+  for (Optional<std::pair<unsigned, unsigned>> &maybePair : repr) {
+    if (maybePair) {
+      b.addInequality(sI.getInequality(maybePair->first));
+      b.addInequality(sI.getInequality(maybePair->second));
+
+      assert(maybePair->first != maybePair->second &&
+             "Upper and lower bounds must be different inequalities!");
+      isDivInequality[maybePair->first] = true;
+      isDivInequality[maybePair->second] = true;
+    } else {
+      llvm_unreachable(
+          "Subtraction is not supported when a representation of "
+          "the local variables of the subtrahend cannot be found!");
+    }
+  }
+
   unsigned initialSnapshot = simplex.getSnapshot();
   unsigned offset = simplex.numConstraints();
+  unsigned numLocalsAdded = b.getNumLocalIds() - bInitNumLocals;
+  simplex.appendVariables(numLocalsAdded);
+
+  unsigned snapshotBeforeIntersect = simplex.getSnapshot();
   simplex.intersectFlatAffineConstraints(sI);
 
   if (simplex.isEmpty()) {
     /// b ^ s_i is empty, so b \ s_i = b. We move directly to i + 1.
     simplex.rollback(initialSnapshot);
+    b.removeIdRange(FlatAffineConstraints::IdKind::Local, bInitNumLocals,
+                    b.getNumLocalIds());
     subtractRecursively(b, simplex, s, i + 1, result);
     return;
   }
 
   simplex.detectRedundant();
-  llvm::SmallBitVector isMarkedRedundant;
-  for (unsigned j = 0; j < 2 * sI.getNumEqualities() + sI.getNumInequalities();
-       j++)
-    isMarkedRedundant.push_back(simplex.isMarkedRedundant(offset + j));
 
-  simplex.rollback(initialSnapshot);
+  // Equalities are added to simplex as a pair of inequalities.
+  unsigned totalNewSimplexInequalities =
+      2 * sI.getNumEqualities() + sI.getNumInequalities();
+  llvm::SmallBitVector isMarkedRedundant(totalNewSimplexInequalities);
+  for (unsigned j = 0; j < totalNewSimplexInequalities; j++)
+    isMarkedRedundant[j] = simplex.isMarkedRedundant(offset + j);
+
+  simplex.rollback(snapshotBeforeIntersect);
 
   // Recurse with the part b ^ ~ineq. Note that b is modified throughout
   // subtractRecursively. At the time this function is called, the current b is
   // actually equal to b ^ s_i1 ^ s_i2 ^ ... ^ s_ij, and ineq is the next
   // inequality, s_{i,j+1}. This function recurses into the next level i + 1
   // with the part b ^ s_i1 ^ s_i2 ^ ... ^ s_ij ^ ~s_{i,j+1}.
   auto recurseWithInequality = [&, i](ArrayRef<int64_t> ineq) {
     size_t snapshot = simplex.getSnapshot();
auto processInequality = [&](ArrayRef<int64_t> ineq) {
     b.addInequality(ineq);
     simplex.addInequality(ineq);
   };
 
   // processInequality appends some additional constraints 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.
-  unsigned originalNumIneqs = b.getNumInequalities();
-  unsigned originalNumEqs = b.getNumEqualities();
+  unsigned bInitNumIneqs = b.getNumInequalities();
+  unsigned bInitNumEqs = b.getNumEqualities();
 
+  // Process all the inequalities, ignoring redundant inequalities and division
+  // inequalities. The result is correct whether or not we ignore these, but
+  // ignoring them makes the result simpler.
   for (unsigned j = 0, e = sI.getNumInequalities(); j < e; j++) {
     if (isMarkedRedundant[j])
       continue;
+    if (isDivInequality[j])
+      continue;
     processInequality(sI.getInequality(j));
   }
 
   offset = sI.getNumInequalities();
   for (unsigned j = 0, e = sI.getNumEqualities(); j < e; ++j) {
     const ArrayRef<int64_t> &coeffs = sI.getEquality(j);
-    // Same as the above loop for inequalities, done once each for the positive
-    // and negative inequalities that make up this equality.
+    // Similar to the above loop for inequalities, done once each for the
+    // positive and negative inequalities that make up this equality. Divisions
+    // are always represented in terms of inequalities and not equalities, so we
+    // do not check that here.
     if (!isMarkedRedundant[offset + 2 * j])
       processInequality(coeffs);
     if (!isMarkedRedundant[offset + 2 * j + 1])
       processInequality(getNegatedCoeffs(coeffs));
   }
 
   // Rollback b and simplex to their initial states.
-  for (unsigned i = b.getNumInequalities(); i > originalNumIneqs; --i)
-    b.removeInequality(i - 1);
-
-  for (unsigned i = b.getNumEqualities(); i > originalNumEqs; --i)
-    b.removeEquality(i - 1);
+  b.removeIdRange(FlatAffineConstraints::IdKind::Local, bInitNumLocals,
+                  b.getNumLocalIds());
+  b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
+  b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
 
   simplex.rollback(initialSnapshot);
 }
 
 /// Return the set difference fac \ set.
 ///
 /// The FAC here is modified in subtractRecursively, so it cannot be a const
 /// reference even though it is restored to its original state before returning
 /// from that function.
 PresburgerSet PresburgerSet::getSetDifference(FlatAffineConstraints fac,
                                               const PresburgerSet &set) {
   assertDimensionsCompatible(fac, set);
-  assert(fac.getNumLocalIds() == 0 &&
-         "Subtracting sets with divisions is not yet supported!");
   if (fac.isEmptyByGCDTest())
     return PresburgerSet::getEmptySet(fac.getNumDimIds(),
                                       fac.getNumSymbolIds());
 
   PresburgerSet result(fac.getNumDimIds(), fac.getNumSymbolIds());
   Simplex simplex(fac);
   subtractRecursively(fac, simplex, set, 0, result);
   return result;

mlir/unittests/Analysis/PresburgerSetTest.cpp

for (const SmallVector<int64_t, 4> &point : points) {
     else
       EXPECT_TRUE(inComplement);
   }
 }
 
 /// Construct a FlatAffineConstraints from a set of inequality and
 /// equality constraints.
 static FlatAffineConstraints
-makeFACFromConstraints(unsigned dims, ArrayRef<SmallVector<int64_t, 4>> ineqs,
-                       ArrayRef<SmallVector<int64_t, 4>> eqs) {
-  FlatAffineConstraints fac(ineqs.size(), eqs.size(), dims + 1, dims);
+makeFACFromConstraints(unsigned ids, ArrayRef<SmallVector<int64_t, 4>> ineqs,
+                       ArrayRef<SmallVector<int64_t, 4>> eqs,
+                       unsigned locals = 0) {
+  FlatAffineConstraints fac(ineqs.size(), eqs.size(), ids + 1, ids - locals, 0,
+                            locals);
   for (const SmallVector<int64_t, 4> &eq : eqs)
     fac.addEquality(eq);
   for (const SmallVector<int64_t, 4> &ineq : ineqs)
     fac.addInequality(ineq);
   return fac;
 }
 
 static FlatAffineConstraints
TEST(SetTest, isEqual) {
   PresburgerSet universeRect = square.unionSet(square.complement());
   PresburgerSet universeSquare = rect.unionSet(rect.complement());
   EXPECT_TRUE(universeRect.isEqual(universeSquare));
   EXPECT_FALSE(universeRect.isEqual(rect));
   EXPECT_FALSE(universeSquare.isEqual(square));
   EXPECT_FALSE(rect.complement().isEqual(square.complement()));
 }
 
+void expectEqual(PresburgerSet s, PresburgerSet t) {
+  EXPECT_TRUE(s.isEqual(t));
+}
+
+void expectEmpty(PresburgerSet s) { EXPECT_TRUE(s.isIntegerEmpty()); }
+
+TEST(SetTest, divisions) {
+  // Note: we currently need to add the equalities as inequalities to the FAC
+  // since detecting divisions based on equalities is not yet supported.
+  //
+  // evens = {x : exists q, x = 2q}.
+  PresburgerSet evens{
+      makeFACFromConstraints(2, {{1, -2, 0}, {-1, 2, 1}}, {{1, -2, 0}}, 1)};
+  // odds = {x : exists q, x = 2q + 1}.
+  PresburgerSet odds{
+      makeFACFromConstraints(2, {{1, -2, 0}, {-1, 2, 1}}, {{1, -2, -1}}, 1)};
+  // FlatAffineConstraints odd = odds.getFlatAffineConstraints(0);
+  PresburgerSet multiples3{
+      makeFACFromConstraints(2, {{1, -3, 0}, {-1, 3, 2}}, {{1, -3, 0}}, 1)};
+  // multiples6 = {x : exists q, x = 6q}.
+  PresburgerSet multiples6{
+      makeFACFromConstraints(2, {{1, -6, 0}, {-1, 6, 5}}, {{1, -6, 0}}, 1)};
+
+  // expectEmpty(PresburgerSet(evens).intersect(PresburgerSet(odds)));
+  // expectEqual(evens.unionSet(odds), PresburgerSet::getUniverse(1));
+  expectEqual(evens.complement(), odds);
+  expectEqual(odds.complement(), evens);
+  expectEqual(multiples3.intersect(evens), multiples6);
+}
+
 } // namespace mlir