[MLIR][Presburger] Attach values only to non-local identifiers in FAVC

[Groverkss]

Summary

This patch changes FlatAffineValueConstraints to only allow attaching
values to non-local identifiers.

The reasoning for this change is:

1. Information attached to local identifiers can be lost since local identifiers can be removed for output size optimizations.
2. There are no current use cases for attaching values to Local identifiers.
3. Attaching a value to a local identifier does not make sense since a local identifier represents existential quantification.

This patch also adds some additional asserts to the affected functions.

[arjunp]

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

419–421

These aren't equivalent. I agree that the start <= end assert is required, though. But is the ability to pass any value when end == start not used? I figure we could preserve that ability as well

456–458

Simliarly, can we just early exit for start >= end?

459

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1354–1358

It seems bad to silently "lose" the Value here. Should we maybe assert-fail if this happens? Is this used?

[Groverkss]

• Address Arjun's comments

[Groverkss]

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

419–421

I think the new assert covers all existing cases and adds another check that start <= end.

Could you point to a case that would pass before (and be a valid range, i.e. start <= end) but not now? I do not see one.

456–458

That wouldn't make sense to me. If start >= end and values.size() has to have the same size as the range, I don't see what the user would pass. I would prefer having this assertion to prevent bugs.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1354–1358

We swap local ids with dim/sym ids when projecting a dim/sym to a local variable. So, it is used.

[arjunp]

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

419–421

Previously start == end would do nothing and work even if start, end are too big. Just mention in the description that some additional asserts were added if we are adding additional constraints

[bondhugula]

LGTM - thanks. There aren't indeed any use cases to attach Values to local identifiers. Some minor comments.

This patch changes FlatAffineValueConstraints to only allow attaching non-local identifiers to it.

The first line of the commit summary has a typo somewhere.

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

36

It looks like we are missing updates to the values field in the class to document they are only for dims and symbols?

408

Nit: assertion messages like these aren't typically ended with a full stop.

446

Update doc comment to mention it's only for dim and symbol identifiers.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

306

Update doc comment on insertId?

[Groverkss]

• Address Uday's comments

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

 class AffineValueMap;
 class IntegerSet;
 class MLIRContext;
 class Value;
 class MemRefType;
 struct MutableAffineMap;
 
 /// FlatAffineValueConstraints represents an extension of IntegerPolyhedron
-/// where each identifier can have an SSA Value attached to it.
+/// where each non-local identifier can have an SSA Value attached to it.
 class FlatAffineValueConstraints : public presburger::IntegerPolyhedron {
 public:
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers.
   FlatAffineValueConstraints(unsigned numReservedInequalities,
                              unsigned numReservedEqualities,
                              unsigned numReservedCols, unsigned numDims,
                              unsigned numSymbols, unsigned numLocals,
                              ArrayRef<Optional<Value>> valArgs = {})
       : IntegerPolyhedron(numReservedInequalities, numReservedEqualities,
                           numReservedCols,
                           presburger::PresburgerSpace::getSetSpace(
                               numDims, numSymbols, numLocals)) {
     assert(numReservedCols >= getNumIds() + 1);
-    assert(valArgs.empty() || valArgs.size() == getNumIds());
+    assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolIds());
     values.reserve(numReservedCols);
     if (valArgs.empty())
-      values.resize(getNumIds(), None);
+      values.resize(getNumDimAndSymbolIds(), None);
     else
       values.append(valArgs.begin(), valArgs.end());
   }
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   FlatAffineValueConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                              unsigned numLocals = 0,
                              ArrayRef<Optional<Value>> valArgs = {})
       : FlatAffineValueConstraints(/*numReservedInequalities=*/0,
                                    /*numReservedEqualities=*/0,
                                    /*numReservedCols=*/numDims + numSymbols +
                                        numLocals + 1,
                                    numDims, numSymbols, numLocals, valArgs) {}
 
   FlatAffineValueConstraints(const IntegerPolyhedron &fac,
                              ArrayRef<Optional<Value>> valArgs = {})
       : IntegerPolyhedron(fac) {
-    assert(valArgs.empty() || valArgs.size() == getNumIds());
+    assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolIds());
     if (valArgs.empty())
-      values.resize(getNumIds(), None);
+      values.resize(getNumDimAndSymbolIds(), None);
     else
       values.append(valArgs.begin(), valArgs.end());
   }
 
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineValueConstraints(const AffineValueMap &avm);
   explicit FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef);
public:
   bool containsId(Value val) const;
 
   /// Swap the posA^th identifier with the posB^th identifier.
   void swapId(unsigned posA, unsigned posB) override;
 
   /// Insert 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. `vals` are the Values
-  /// corresponding to the identifiers. Return the absolute column position
-  /// (i.e., not relative to the kind of identifier) of the first added
-  /// identifier.
+  /// corresponding to the identifiers. Values should not be used with
+  /// IdKind::Local since values can only be attached to non-local identifiers.
+  /// Return the absolute column position (i.e., not relative to the kind of
+  /// identifier) of the first added identifier.
   ///
   /// Note: Empty Values are allowed in `vals`.
   unsigned insertDimId(unsigned pos, unsigned num = 1) {
     return insertId(IdKind::SetDim, pos, num);
   }
   unsigned insertSymbolId(unsigned pos, unsigned num = 1) {
     return insertId(IdKind::Symbol, pos, num);
   }
public:
   bool areIdsAlignedWithOther(const FlatAffineValueConstraints &other);
 
   /// Replaces the contents of this FlatAffineValueConstraints with `other`.
   void clearAndCopyFrom(const IntegerRelation &other) override;
 
   /// Returns the Value associated with the pos^th identifier. Asserts if
   /// no Value identifier was associated.
   inline Value getValue(unsigned pos) const {
+    assert(pos < getNumDimAndSymbolIds() && "Invalid position");
     assert(hasValue(pos) && "identifier's Value not set");
     return values[pos].getValue();
   }
 
   /// Returns true if the pos^th identifier has an associated Value.
-  inline bool hasValue(unsigned pos) const { return values[pos].hasValue(); }
+  inline bool hasValue(unsigned pos) const {
+    assert(pos < getNumDimAndSymbolIds() && "Invalid position");
+    return values[pos].hasValue();
+  }
 
   /// Returns true if at least one identifier has an associated Value.
   bool hasValues() const;
 
   /// Returns the Values associated with identifiers in range [start, end).
   /// Asserts if no Value was associated with one of these identifiers.
   inline void getValues(unsigned start, unsigned end,
                         SmallVectorImpl<Value> *values) const {
-    assert((start < getNumIds() || start == end) && "invalid start position");
-    assert(end <= getNumIds() && "invalid end position");
+    assert(end <= getNumDimAndSymbolIds() && "invalid end position");
+    assert(start <= end && "invalid start position");
     values->clear();
     values->reserve(end - start);
     for (unsigned i = start; i < end; i++)
       values->push_back(getValue(i));
   }
   inline void getAllValues(SmallVectorImpl<Value> *values) const {
-    getValues(0, getNumIds(), values);
+    getValues(0, getNumDimAndSymbolIds(), values);
   }
 
   inline ArrayRef<Optional<Value>> getMaybeValues() const {
     return {values.data(), values.size()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimValues() const {
     return {values.data(), getNumDimIds()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const {
     return {values.data() + getNumDimIds(), getNumSymbolIds()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const {
     return {values.data(), getNumDimIds() + getNumSymbolIds()};
   }
 
   /// Sets the Value associated with the pos^th identifier.
   inline void setValue(unsigned pos, Value val) {
-    assert(pos < getNumIds() && "invalid id position");
+    assert(pos < getNumDimAndSymbolIds() && "invalid id position");
     values[pos] = val;
   }
 
   /// Sets the Values associated with the identifiers in the range [start, end).
+  /// The range must contain only dim and symbol identifiers.
   void setValues(unsigned start, unsigned end, ArrayRef<Value> values) {
-    assert((start < getNumIds() || end == start) && "invalid start position");
     assert(end <= getNumIds() && "invalid end position");
-    assert(values.size() == end - start);
+    assert(start <= end && "invalid start position");
+    assert(values.size() == end - start &&
+           "value should be provided for each identifier in the range.");
     for (unsigned i = start; i < end; ++i)
       setValue(i, values[i - start]);
   }
 
   /// Merge and align symbols of `this` and `other` such that both get union of
   /// of symbols that are unique. Symbols in `this` and `other` should be
   /// unique. Symbols with Value as `None` are considered to be inequal to all
   /// other symbols.
   void mergeSymbolIds(FlatAffineValueConstraints &other);
protected:
   void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
                                bool *isResultIntegerExact = nullptr) override;
 
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// FlatAffineConstraints. Also, prints for each identifier whether there is
   /// an SSA Value attached to it.
   void printSpace(raw_ostream &os) const override;
 
-  /// 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 with any Value are set to
+  /// Values corresponding to the (column) non-local identifiers of this
+  /// constraint system appearing in the order the identifiers correspond to
+  /// columns. Identifiers that aren't associated with any Value are set to
   /// None.
   SmallVector<Optional<Value>, 8> values;
 };
 
 /// A FlatAffineRelation represents a set of ordered pairs (domain -> range)
 /// where "domain" and "range" are tuples of identifiers. The relation is
 /// represented as a FlatAffineValueConstraints with separation of dimension
 /// identifiers into domain and  range. The identifiers are stored as:

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 FlatAffineValueConstraints::FlatAffineValueConstraints(IntegerSet set)
     : IntegerPolyhedron(set.getNumInequalities(), set.getNumEqualities(),
                         set.getNumDims() + set.getNumSymbols() + 1,
                         PresburgerSpace::getSetSpace(set.getNumDims(),
                                                      set.getNumSymbols(),
                                                      /*numLocals=*/0)) {
 
   // Resize values.
-  values.resize(getNumIds(), None);
+  values.resize(getNumDimAndSymbolIds(), None);
 
   // Flatten expressions and add them to the constraint system.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatAffineValueConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
     assert(false && "flattening unimplemented for semi-affine integer sets");
     return;
   }

 unsigned FlatAffineValueConstraints::insertSymbolId(unsigned pos,
                                                     ValueRange vals) {
   return insertId(IdKind::Symbol, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
                                               unsigned num) {
   unsigned absolutePos = IntegerPolyhedron::insertId(kind, pos, num);
+
+  if (kind != IdKind::Local) {
     values.insert(values.begin() + absolutePos, num, None);
-  assert(values.size() == getNumIds());
+    assert(values.size() == getNumDimAndSymbolIds());
+  }
+
   return absolutePos;
 }
 
 unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
                                               ValueRange vals) {
-  assert(!vals.empty() && "expected ValueRange with Values");
+  assert(!vals.empty() && "expected ValueRange with Values.");
+  assert(kind != IdKind::Local &&
+         "values cannot be attached to local identifiers.");
   unsigned num = vals.size();
   unsigned absolutePos = IntegerPolyhedron::insertId(kind, pos, num);
 
   // If a Value is provided, insert it; otherwise use None.
   for (unsigned i = 0; i < num; ++i)
     values.insert(values.begin() + absolutePos + i,
                   vals[i] ? Optional<Value>(vals[i]) : None);
 
-  assert(values.size() == getNumIds());
+  assert(values.size() == getNumDimAndSymbolIds());
   return absolutePos;
 }
 
 bool FlatAffineValueConstraints::hasValues() const {
   return llvm::find_if(values, [](Optional<Value> id) {
            return id.hasValue();
          }) != values.end();
 }
bool FlatAffineValueConstraints::areIdsAlignedWithOther(
   return areIdsAligned(*this, other);
 }
 
 /// Checks if the SSA values associated with `cst`'s identifiers in range
 /// [start, end) are unique.
 static bool LLVM_ATTRIBUTE_UNUSED areIdsUnique(
     const FlatAffineValueConstraints &cst, unsigned start, unsigned end) {
 
-  assert(start <= cst.getNumIds() && "Start position out of bounds");
-  assert(end <= cst.getNumIds() && "End position out of bounds");
+  assert(start <= cst.getNumDimAndSymbolIds() &&
+         "Start position out of bounds");
+  assert(end <= cst.getNumDimAndSymbolIds() && "End position out of bounds");
 
   if (start >= end)
     return true;
 
   SmallPtrSet<Value, 8> uniqueIds;
   ArrayRef<Optional<Value>> maybeValues =
       cst.getMaybeValues().slice(start, end - start);
   for (Optional<Value> val : maybeValues) {
     if (val.hasValue() && !uniqueIds.insert(val.getValue()).second)
       return false;
   }
   return true;
 }
 
 /// Checks if the SSA values associated with `cst`'s identifiers are unique.
 static bool LLVM_ATTRIBUTE_UNUSED
 areIdsUnique(const FlatAffineValueConstraints &cst) {
-  return areIdsUnique(cst, 0, cst.getNumIds());
+  return areIdsUnique(cst, 0, cst.getNumDimAndSymbolIds());
 }
 
 /// Checks if the SSA values associated with `cst`'s identifiers of kind `kind`
 /// are unique.
 static bool LLVM_ATTRIBUTE_UNUSED
 areIdsUnique(const FlatAffineValueConstraints &cst, IdKind kind) {
 
   if (kind == IdKind::SetDim)
     return areIdsUnique(cst, 0, cst.getNumDimIds());
   if (kind == IdKind::Symbol)
     return areIdsUnique(cst, cst.getNumDimIds(), cst.getNumDimAndSymbolIds());
-  if (kind == IdKind::Local)
-    return areIdsUnique(cst, cst.getNumDimAndSymbolIds(), cst.getNumIds());
   llvm_unreachable("Unexpected IdKind");
 }
 
 /// Merge and align the identifiers of A and B starting at 'offset', so that
 /// both constraint systems get the union of the contained identifiers that is
 /// dimension-wise and symbol-wise unique; both constraint systems are updated
 /// so that they have the union of all identifiers, with A's original
 /// identifiers appearing first followed by any of B's identifiers that didn't
 /// appear in A. Local identifiers in B that have the same division
 /// representation as local identifiers in A are merged into one.
 //  E.g.: Input: A has ((%i, %j) [%M, %N]) and B has (%k, %j) [%P, %N, %M])
 //        Output: both A, B have (%i, %j, %k) [%M, %N, %P]
 static void mergeAndAlignIds(unsigned offset, FlatAffineValueConstraints *a,
                              FlatAffineValueConstraints *b) {
   assert(offset <= a->getNumDimIds() && offset <= b->getNumDimIds());
   // A merge/align isn't meaningful if a cst's ids aren't distinct.
   assert(areIdsUnique(*a) && "A's values aren't unique");
   assert(areIdsUnique(*b) && "B's values aren't unique");
 
   assert(std::all_of(a->getMaybeValues().begin() + offset,
-                     a->getMaybeValues().begin() + a->getNumDimAndSymbolIds(),
+                     a->getMaybeValues().end(),
                      [](Optional<Value> id) { return id.hasValue(); }));
 
   assert(std::all_of(b->getMaybeValues().begin() + offset,
-                     b->getMaybeValues().begin() + b->getNumDimAndSymbolIds(),
+                     b->getMaybeValues().end(),
                      [](Optional<Value> id) { return id.hasValue(); }));
 
   SmallVector<Value, 4> aDimValues;
   a->getValues(offset, a->getNumDimIds(), &aDimValues);
 
   {
     // Merge dims from A into B.
     unsigned d = offset;
void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
   // and align the IDs from both constraints, and then append the constraints
   // from the ifOp into the current one.
   mergeAndAlignIdsWithOther(0, &cst);
   append(cst);
 }
 
 bool FlatAffineValueConstraints::hasConsistentState() const {
   return IntegerPolyhedron::hasConsistentState() &&
-         values.size() == getNumIds();
+         values.size() == getNumDimAndSymbolIds();
 }
 
 void FlatAffineValueConstraints::removeIdRange(IdKind kind, unsigned idStart,
                                                unsigned idLimit) {
   IntegerPolyhedron::removeIdRange(kind, idStart, idLimit);
   unsigned offset = getIdKindOffset(kind);
+
+  if (kind != IdKind::Local) {
     values.erase(values.begin() + idStart + offset,
                  values.begin() + idLimit + offset);
   }
+}
 
 // Determine whether the identifier at 'pos' (say id_r) can be expressed as
 // modulo of another known identifier (say id_n) w.r.t a constant. For example,
 // if the following constraints hold true:
 // ```
 // 0 <= id_r <= divisor - 1
 // id_n - (divisor * q_expr) = id_r
 // ```

 bool FlatAffineValueConstraints::containsId(Value val) const {
   return llvm::any_of(values, [&](const Optional<Value> &mayBeId) {
     return mayBeId.hasValue() && mayBeId.getValue() == val;
   });
 }
 
 void FlatAffineValueConstraints::swapId(unsigned posA, unsigned posB) {
   IntegerPolyhedron::swapId(posA, posB);
+
+  if (getIdKindAt(posA) == IdKind::Local && getIdKindAt(posB) == IdKind::Local)
+    return;
+
+  // Treat value of a local identifer as None.
+  if (getIdKindAt(posA) == IdKind::Local)
+    values[posB] = None;
+  else if (getIdKindAt(posB) == IdKind::Local)
+    values[posA] = None;
+  else
     std::swap(values[posA], values[posB]);
 }
 
 void FlatAffineValueConstraints::addBound(BoundType type, Value val,
                                           int64_t value) {
   unsigned pos;
   if (!findId(val, &pos))
     // This is a pre-condition for this method.
     assert(0 && "id not found");
   addBound(type, pos, value);
 }
 
 void FlatAffineValueConstraints::printSpace(raw_ostream &os) const {
   IntegerPolyhedron::printSpace(os);
   os << "(";
-  for (unsigned i = 0, e = getNumIds(); i < e; i++) {
+  for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; i++) {
     if (hasValue(i))
       os << "Value ";
     else
       os << "None ";
   }
+  for (unsigned i = getIdKindOffset(IdKind::Local),
+                e = getIdKindEnd(IdKind::Local);
+       i < e; ++i)
+    os << "Local ";
   os << " const)\n";
 }
 
 void FlatAffineValueConstraints::clearAndCopyFrom(
     const IntegerRelation &other) {
 
   if (auto *otherValueSet =
           dyn_cast<const FlatAffineValueConstraints>(&other)) {
     *this = *otherValueSet;
   } else {
     *static_cast<IntegerRelation *>(this) = other;
     values.clear();
-    values.resize(getNumIds(), None);
+    values.resize(getNumDimAndSymbolIds(), None);
   }
 }
 
 void FlatAffineValueConstraints::fourierMotzkinEliminate(
     unsigned pos, bool darkShadow, bool *isResultIntegerExact) {
-  SmallVector<Optional<Value>, 8> newVals;
-  newVals.reserve(getNumIds() - 1);
-  newVals.append(values.begin(), values.begin() + pos);
-  newVals.append(values.begin() + pos + 1, values.end());
+  SmallVector<Optional<Value>, 8> newVals = values;
+  if (getIdKindAt(pos) != IdKind::Local)
+    newVals.erase(newVals.begin() + pos);
   // Note: Base implementation discards all associated Values.
   IntegerPolyhedron::fourierMotzkinEliminate(pos, darkShadow,
                                              isResultIntegerExact);
   values = newVals;
-  assert(values.size() == getNumIds());
+  assert(values.size() == getNumDimAndSymbolIds());
 }
 
 void FlatAffineValueConstraints::projectOut(Value val) {
   unsigned pos;
   bool ret = findId(val, &pos);
   assert(ret);
   (void)ret;
   fourierMotzkinEliminate(pos);