[MLIR] FlatAffineConstraints : Store explicit division representation for local identifiers

[Groverkss]

Summary

This patch improves support for local identifiers by explicitly storing representations of their values as divisions of an affine expression by a constant denominator, when such a representation is known.

Exact re-computation of division representations from inequalities is expensive. Many operations on PresburgerSets when local identifiers are involved require these representations. Storing these representations allows for faster computation in PresburgerSet operations where local identifiers are involved.

[arjunp]

Thanks for the patch!

I did not point it out in each instance, but in general, please use full sentences and punctuation for all comments.

mlir/include/mlir/Analysis/AffineStructures.h

82

You can initialize this directly in the member initializer list above right?

214

Can we add a getDivNumerator?

305

Division representation? Also, please put the full stop, here and everywhere.

319

nit: unnecessary extra line?

320

nit: the floordiv

324

Maybe say that the FAC is considered to be in an inconsistent state?

676

nit: corresponding to

mlir/lib/Analysis/AffineStructures.cpp

158

Initialize divDenom?

227

Please use pre-increment.

233

Please use pre-increment.

234

If you add getDivNumerator, you can replace this loop with getDivNumerator(r) == other.getDivNumerator(r).

2199

This is not needed right? This is not set anywhere and should already be zero? (You can probably still add a comment saying that it's zero)

2202

Maybe mention, with a zero inserted for the new local

2250

nit: swap division representationS OF both?

2251

nit: unsigned?

2256

If one is a local and the other isn't, you should clear the representation for the local right?

2630

I think it might be more consistent to skip the + and just put a space

mlir/unittests/Analysis/AffineStructuresTest.cpp

601

Here and below, expect that the optional actually has a value as well?

604

Also check (0, 2)?

604

You can also check the inequalities that were added.

606

Can you add at position 0 instead? Will the other one get moved appropriately?

612

Shouldn't it be -1?

617

Can you add a division with non-zero constant term and check it?

639

Try removing the first division and check if the other one moves?

640

You can also use 1u

640

Also check for consistency?

[arjunp]

mlir/unittests/Analysis/AffineStructuresTest.cpp

612

It should be fac2 I think

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

158

Initializing divDenom does not make sense according to me. IntegerSet has no concept of local variables so divDenom is always of size 0. divisions need to be initialized due to column size.

234

I'm not sure about this. getDivNumerator would return an ArrayRef<int64_t>. Does using == work on ArrayRefs?

mlir/unittests/Analysis/AffineStructuresTest.cpp

640

The consistency check is a private method IIUC. Not sure how to use it here.

[Groverkss]

• Addressed arjunp's comments
• Fixed punctuation in comments

[Groverkss]

Restored accidently deleted previous commits

[Groverkss]

• Fixed failing tests

[Groverkss]

• Fixed formatting issues

[bondhugula]

Exact re-computation of division representations from inequalities is expensive. Many operations on PresburgerSets when local identifiers are involved require these representations.
Storing these representations allows for faster computation in PresburgerSet operations where local identifiers are involved.

This approach is storing redundant information in the system for faster detection/computation later. One can't be sure of updating any inequality without invalidating such information. There are some non-trivial tradeoffs to discuss here. This is a major change and non-trivial enough to discuss thoroughly on discourse first? Could you please post a description of this - sort of an expanded commit summary with additional detail on tradeoffs?

mlir/include/mlir/Analysis/AffineStructures.h

178

This doc comment doesn't parse well - rephrase?

185

Missing assertion on valid pos range.

187

Likewise.

215

Doc comment here along the same lines as as getDivDenonm....

mlir/include/mlir/Analysis/AffineStructures.h

public:
   FlatAffineConstraints(unsigned numReservedInequalities,
                         unsigned numReservedEqualities,
                         unsigned numReservedCols, unsigned numDims = 0,
                         unsigned numSymbols = 0, unsigned numLocals = 0,
                         ArrayRef<Optional<Value>> idArgs = {})
       : numIds(numDims + numSymbols + numLocals), numDims(numDims),
         numSymbols(numSymbols),
         equalities(0, numIds + 1, numReservedEqualities, numReservedCols),
-        inequalities(0, numIds + 1, numReservedInequalities, numReservedCols) {
+        inequalities(0, numIds + 1, numReservedInequalities, numReservedCols),
+        divisions(numLocals, numIds + 1, numLocals, numReservedCols) {
     assert(numReservedCols >= numIds + 1);
     assert(idArgs.empty() || idArgs.size() == numIds);
     ids.reserve(numReservedCols);
     if (idArgs.empty())
       ids.resize(numIds, None);
     else
       ids.append(idArgs.begin(), idArgs.end());
+    divDenom.resize(numLocals, llvm::None);
   }
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                         unsigned numLocals = 0,
                         ArrayRef<Optional<Value>> idArgs = {})
       : FlatAffineConstraints(/*numReservedInequalities=*/0,
public:
   inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
 
   inline int64_t atIneq(unsigned i, unsigned j) const {
     return inequalities(i, j);
   }
 
   inline int64_t &atIneq(unsigned i, unsigned j) { return inequalities(i, j); }
 
+  /// Returns the value at specified division dividend row and column
+  inline int64_t atDiv(unsigned i, unsigned j) const { return divisions(i, j); }
+  inline int64_t &atDiv(unsigned i, unsigned j) { return divisions(i, j); }
+
+  /// Returns the denominator for specified division
+  /// Returns None if the local identifier does not have an explicit
+  /// representation
+  inline Optional<int64_t> atDivDenom(unsigned pos) const {
+    return divDenom[pos];
+  }
+  inline Optional<int64_t> &atDivDenom(unsigned pos) { return divDenom[pos]; }
+
   /// Returns the number of columns in the constraint system.
   inline unsigned getNumCols() const { return numIds + 1; }
 
   inline unsigned getNumEqualities() const { return equalities.getNumRows(); }
 
   inline unsigned getNumInequalities() const {
     return inequalities.getNumRows();
   }
 
   inline unsigned getNumReservedEqualities() const {
     return equalities.getNumReservedRows();
   }
 
   inline unsigned getNumReservedInequalities() const {
     return inequalities.getNumReservedRows();
   }
 
   inline ArrayRef<int64_t> getEquality(unsigned idx) const {
     return equalities.getRow(idx);
   }
 
   inline ArrayRef<int64_t> getInequality(unsigned idx) const {
     return inequalities.getRow(idx);
   }
 
   /// Adds constraints (lower and upper bounds) for the specified 'affine.for'
   /// operation's Value using IR information stored in its bound maps. The
   /// right identifier is first looked up using forOp's Value. Asserts if the
   /// Value corresponding to the 'affine.for' operation isn't found in the
   /// constraint system. Returns failure for the yet unimplemented/unsupported
   /// cases.  Any new identifiers that are found in the bound operands of the
   /// 'affine.for' operation are added as trailing identifiers (either
   /// dimensional or symbolic depending on whether the operand is a valid
   /// symbol).
public:
   // Adds an equality from the coefficients specified in eq.
   void addEquality(ArrayRef<int64_t> eq);
 
   /// Adds a constant lower bound constraint for the specified identifier.
   void addConstantLowerBound(unsigned pos, int64_t lb);
   /// Adds a constant upper bound constraint for the specified identifier.
   void addConstantUpperBound(unsigned pos, int64_t ub);
 
+  /// Sets the division for pos^th local identifier
+  void setDivision(unsigned pos, ArrayRef<int64_t> dividend, int64_t divisor);
+
+  /// Clears division representation for the pos^th local identifier
+  void clearDivision(unsigned pos);
+
+  /// Clears division representation for all local identifiers
+  void clearAllDivisions();
+
   /// Adds a new local identifier as the floordiv of an affine function of other
   /// identifiers, the coefficients of which are provided in 'dividend' and with
   /// respect to a positive constant 'divisor'. Two constraints are added to the
   /// system to capture equivalence with the floordiv:
   /// q = dividend floordiv c    <=>   c*q <= dividend <= c*q + c - 1.
+  ///
+  /// The added division representation for floordiv is considered to be part of
+  /// the constraint system. If the added inequalities are removed or modified
+  /// such that the division representation of that floordiv becomes invalid,
+  /// the division representation of that local identifier must be updated or
+  /// cleared. Otherwise, the state of divisions becomes inconsistent.
   void addLocalFloorDiv(ArrayRef<int64_t> dividend, int64_t divisor);
 
   /// Adds a constant lower bound constraint for the specified expression.
   void addConstantLowerBound(ArrayRef<int64_t> expr, int64_t lb);
   /// Adds a constant upper bound constraint for the specified expression.
   void addConstantUpperBound(ArrayRef<int64_t> expr, int64_t ub);
 
   /// Sets the identifier at the specified position to a constant.
private:
   unsigned numSymbols;
 
   /// Coefficients of affine equalities (in == 0 form).
   Matrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
   Matrix inequalities;
 
+  /// Coefficients of affine expressions corresponding the dividend of division
+  /// representing local identifiers.
+  Matrix divisions;
+
+  /// Denominators of divisions representing local identifiers.
+  /// Identifiers whose explicit representation is unknown are marked None
+  SmallVector<Optional<int64_t>, 8> divDenom;
+
   /// Values corresponding to the (column) identifiers of this constraint
   /// system appearing in the order the identifiers correspond to columns.
   /// Temporary ones or those that aren't associated to any Value are set to
   /// None.
   SmallVector<Optional<Value>, 8> ids;
 
   /// A parameter that controls detection of an unrealistic number of
   /// constraints. If the number of constraints is this many times the number of

mlir/lib/Analysis/AffineStructures.cpp

std::unique_ptr<FlatAffineConstraints> FlatAffineConstraints::clone() const {
   return std::make_unique<FlatAffineConstraints>(*this);
 }
 
 // Construct from an IntegerSet.
 FlatAffineConstraints::FlatAffineConstraints(IntegerSet set)
     : numIds(set.getNumDims() + set.getNumSymbols()), numDims(set.getNumDims()),
       numSymbols(set.getNumSymbols()),
       equalities(0, numIds + 1, set.getNumEqualities(), numIds + 1),
-      inequalities(0, numIds + 1, set.getNumInequalities(), numIds + 1) {
+      inequalities(0, numIds + 1, set.getNumInequalities(), numIds + 1),
+      divisions(0, numIds + 1) {
   ids.resize(numIds, None);
 
   // Flatten expressions and add them to the constraint system.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatAffineConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
     assert(false && "flattening unimplemented for semi-affine integer sets");
     return;
void FlatAffineConstraints::append(const FlatAffineConstraints &other) {
   equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
 
   for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
     addInequality(other.getInequality(r));
   }
   for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
     addEquality(other.getEquality(r));
   }
+
+  // Division representations that do not match are marked as unknown
+  for (unsigned r = 0, e = getNumLocalIds(); r < e; r++) {
+    if (atDivDenom(r) != other.atDivDenom(r)) {
+      atDivDenom(r) = llvm::None;
+      continue;
+    }
+
+    for (unsigned k = 0; k < getNumIds(); k++) {
+      if (atDiv(r, k) != other.atDiv(r, k)) {
+        atDivDenom(r) = llvm::None;
+        break;
+      }
+    }
+  }
 }
 
 void FlatAffineConstraints::addLocalId(unsigned pos) {
   addId(IdKind::Local, pos);
 }
 
 void FlatAffineConstraints::addDimId(unsigned pos, Value id) {
   addId(IdKind::Dimension, pos, id);
void FlatAffineConstraints::addId(IdKind kind, unsigned pos, Value id) {
   if (kind == IdKind::Dimension) {
     absolutePos = pos;
     numDims++;
   } else if (kind == IdKind::Symbol) {
     absolutePos = pos + getNumDimIds();
     numSymbols++;
   } else {
     absolutePos = pos + getNumDimIds() + getNumSymbolIds();
+
+    // Insert unknown division for local identifier
+    divisions.insertRow(pos);
+    divDenom.insert(divDenom.begin() + pos, llvm::None);
   }
   numIds++;
 
   inequalities.insertColumn(absolutePos);
   equalities.insertColumn(absolutePos);
+  divisions.insertColumn(absolutePos);
 
   // If an 'id' is provided, insert it; otherwise use None.
   if (id)
     ids.insert(ids.begin() + absolutePos, id);
   else
     ids.insert(ids.begin() + absolutePos, None);
   assert(ids.size() == getNumIds());
 }
void FlatAffineConstraints::normalizeConstraintsByGCD() {
   }
 }
 
 bool FlatAffineConstraints::hasConsistentState() const {
   if (!inequalities.hasConsistentState())
     return false;
   if (!equalities.hasConsistentState())
     return false;
+  if (!divisions.hasConsistentState())
+    return false;
   if (ids.size() != getNumIds())
     return false;
+  if (divDenom.size() != getNumLocalIds())
+    return false;
+  if (divDenom.size() != divisions.getNumRows())
+    return false;
 
   // Catches errors where numDims, numSymbols, numIds aren't consistent.
   if (numDims > numIds || numSymbols > numIds || numDims + numSymbols > numIds)
     return false;
 
   return true;
 }
 
void FlatAffineConstraints::removeIdRange(unsigned idStart, unsigned idLimit) {
 
   // We are going to be removing one or more identifiers from the range.
   assert(idStart < numIds && "invalid idStart position");
 
   // TODO: Make 'removeIdRange' a lambda called from here.
   // Remove eliminated identifiers from the constraints..
   equalities.removeColumns(idStart, idLimit - idStart);
   inequalities.removeColumns(idStart, idLimit - idStart);
+  divisions.removeColumns(idStart, idLimit - idStart);
 
   // Update members numDims, numSymbols and numIds.
   unsigned numDimsEliminated = 0;
   unsigned numLocalsEliminated = 0;
   unsigned numColsEliminated = idLimit - idStart;
   if (idStart < numDims) {
     numDimsEliminated = std::min(numDims, idLimit) - idStart;
   }
   // Check how many local id's were removed. Note that our identifier order is
   // [dims, symbols, locals]. Local id start at position numDims + numSymbols.
   if (idLimit > numDims + numSymbols) {
     numLocalsEliminated = std::min(
         idLimit - std::max(idStart, numDims + numSymbols), getNumLocalIds());
+
+    // Remove local identifier representations from [0, numLocalsEliminated]
+    divisions.removeRows(0, numLocalsEliminated);
+    divDenom.erase(divDenom.begin(), divDenom.begin() + numLocalsEliminated);
   }
+
   unsigned numSymbolsEliminated =
       numColsEliminated - numDimsEliminated - numLocalsEliminated;
 
   numDims -= numDimsEliminated;
   numSymbols -= numSymbolsEliminated;
   numIds = numIds - numColsEliminated;
 
   ids.erase(ids.begin() + idStart, ids.begin() + idLimit);
void FlatAffineConstraints::addConstantUpperBound(ArrayRef<int64_t> expr,
                                                   int64_t ub) {
   assert(expr.size() == getNumCols());
   unsigned row = inequalities.appendExtraRow();
   for (unsigned i = 0, e = expr.size(); i < e; ++i)
     inequalities(row, i) = -expr[i];
   inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) += ub;
 }
 
+void FlatAffineConstraints::setDivision(unsigned pos,
+                                        ArrayRef<int64_t> dividend,
+                                        int64_t divisor) {
+  assert(pos < getNumLocalIds() && "Local identifier position out of range");
+  assert(dividend.size() == getNumCols() && "incorrect dividend size");
+  assert(divisor > 0 && "positive divisor expected");
+
+  for (unsigned i = 0, e = getNumIds(); i < e; ++i)
+    atDiv(pos, i) = dividend[i];
+  atDivDenom(pos) = divisor;
+}
+
+void FlatAffineConstraints::clearDivision(unsigned pos) {
+  atDivDenom(pos) = llvm::None;
+}
+
+void FlatAffineConstraints::clearAllDivisions() {
+  for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) 
+    clearDivision(i);
+}
+
 /// Adds a new local identifier as the floordiv of an affine function of other
 /// identifiers, the coefficients of which are provided in 'dividend' and with
 /// respect to a positive constant 'divisor'. Two constraints are added to the
 /// system to capture equivalence with the floordiv.
 ///      q = expr floordiv c    <=>   c*q <= expr <= c*q + c - 1.
 void FlatAffineConstraints::addLocalFloorDiv(ArrayRef<int64_t> dividend,
                                              int64_t divisor) {
   assert(dividend.size() == getNumCols() && "incorrect dividend size");
   assert(divisor > 0 && "positive divisor expected");
 
   addLocalId(getNumLocalIds());
 
   // Add two constraints for this new identifier 'q'.
   SmallVector<int64_t, 8> bound(dividend.size() + 1);
 
-  // dividend - q * divisor >= 0
+  // dividend >= 0
   std::copy(dividend.begin(), dividend.begin() + dividend.size() - 1,
             bound.begin());
   bound.back() = dividend.back();
+  bound[getNumIds() - 1] = 0;
+
+  // Set explicit representation for local identifier
+  // "bound" currently exactly represents the divisions expression
+  setDivision(getNumLocalIds() - 1, bound, divisor);
+
+  // dividend - q * divisor >= 0
   bound[getNumIds() - 1] = -divisor;
   addInequality(bound);
 
   // -dividend +qdivisor * q + divisor - 1 >= 0
   std::transform(bound.begin(), bound.end(), bound.begin(),
                  std::negate<int64_t>());
   bound[bound.size() - 1] += divisor - 1;
   addInequality(bound);
void FlatAffineConstraints::swapId(unsigned posA, unsigned posB) {
 
   if (posA == posB)
     return;
 
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
     std::swap(atIneq(r, posA), atIneq(r, posB));
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
     std::swap(atEq(r, posA), atEq(r, posB));
+  for (unsigned r = 0, e = getNumLocalIds(); r < e; r++)
+    std::swap(atDiv(r, posA), atDiv(r, posB));
+
   std::swap(getId(posA), getId(posB));
+
+  // Swap division representation if both local identifiers
+  int64_t localIdOffset = getNumDimAndSymbolIds();
+  if (posA >= localIdOffset && posB >= localIdOffset) {
+    std::swap(divDenom[posA - localIdOffset], divDenom[posB - localIdOffset]);
+    divisions.swapRows(posA - localIdOffset, posB - localIdOffset);
+  }
 }
 
 void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= numDims + numSymbols &&
          "invalid separation position");
   numDims = numDims + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
     os << "= 0\n";
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     for (unsigned j = 0, f = getNumCols(); j < f; ++j) {
       os << atIneq(i, j) << " ";
     }
     os << ">= 0\n";
   }
-  os << '\n';
+  os << "Local Ids:\n";
+
+  for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) {
+    if (!atDivDenom(i)) {
+      os << "Local Id with Unknown explicit representation\n";
+      continue;
+    }
+
+    os << "[";
+    for (unsigned j = 0, f = getNumCols() - 1; j < f; ++j)
+      os << atDiv(i, j) << " + ";
+    os << atDiv(i, getNumCols() - 1);
+
+    os << "] / ";
+    os << *atDivDenom(i);
+    os << "\n";
+  }
+
+  os << "\n";
 }
 
 void FlatAffineConstraints::dump() const { print(llvm::errs()); }
 
 /// Removes duplicate constraints, trivially true constraints, and constraints
 /// that can be detected as redundant as a result of differing only in their
 /// constant term part. A constraint of the form <non-negative constant> >= 0 is
 /// considered trivially true.

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, clearConstraints) {
 
   fac.clearConstraints();
 
   fac.addInequality({1, 0});
   EXPECT_EQ(fac.atIneq(0, 0), 1);
   EXPECT_EQ(fac.atIneq(0, 1), 0);
 }
 
+TEST(FlatAffineConstraintsTest, testDivisions) {
+  FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
+
+  // Add inequality x >= 0
+  fac.addInequality({1, 0});
+  EXPECT_EQ(fac.atIneq(0, 0), 1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+
+  // Add division q = [(x)/4]
+  fac.addLocalFloorDiv({1, 0}, 4);
+
+  EXPECT_EQ(*fac.atDivDenom(0), 4);
+  EXPECT_EQ(fac.atDiv(0, 0), 1);
+  EXPECT_EQ(fac.atDiv(0, 1), 0);
+
+  // Add local identifier p
+  fac.addLocalId(1);
+  EXPECT_EQ(fac.atDivDenom(1), llvm::None);
+
+  // Test appending checks divisions
+  FlatAffineConstraints fac2 = makeFACFromConstraints(1, {}, {});
+  fac2.addInequality({-1, 0});
+  EXPECT_EQ(fac.atIneq(0, 0), 1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+
+  fac2.addLocalFloorDiv({1, 0}, 4);
+
+  EXPECT_EQ(*fac2.atDivDenom(0), 4);
+  EXPECT_EQ(fac2.atDiv(0, 0), 1);
+  EXPECT_EQ(fac2.atDiv(0, 1), 0);
+
+  fac2.addLocalId(1);
+  EXPECT_EQ(*fac2.atDivDenom(0), 4);
+  EXPECT_EQ(fac2.atDivDenom(1), llvm::None);
+
+  fac.append(fac2);
+  EXPECT_EQ(*fac.atDivDenom(0), 4);
+  EXPECT_EQ(fac.atDiv(0, 0), 1);
+  EXPECT_EQ(fac.atDiv(0, 1), 0);
+  EXPECT_EQ(fac.atDivDenom(1), llvm::None);
+
+  // Test swapping swaps divisions
+  fac.swapId(1, 2);
+  EXPECT_EQ(fac.atDivDenom(0), llvm::None);
+  EXPECT_EQ(*fac.atDivDenom(1), 4);
+  EXPECT_EQ(fac.atDiv(1, 0), 1);
+  EXPECT_EQ(fac.atDiv(1, 1), 0);
+
+  // Test removal of division
+  fac.removeId(2);
+  EXPECT_EQ(fac.getNumLocalIds(), (unsigned)1);
+
+  fac.clearDivision(0);
+  EXPECT_EQ(fac.atDivDenom(0), llvm::None);
+}
+
 } // namespace mlir