[MLIR] FlatAffineValueConstraints: Fix bug in mergeSymbolIds

[Groverkss]

Summary

This patch fixes a bug in implementation mergeSymbolIds where symbol
identifiers were not unique after merging them. Asserts for checking uniqueness
before and after the merge are also added. The asserts checking uniqueness
after the merge fail without the fix on existing test cases.

[grosser]

Is there a way to add a test case for this?

[Groverkss]

In D111958#3068845, @grosser wrote:

Is there a way to add a test case for this?

The existing test cases fail on adding the assert. Do you prefer a test case that fails without the asserts too? I do not have a test case for that as of now, but I can try to find it.

[grosser]

The update of the commit message is enough, from my side.

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

451–454

This may stop working if the variable layout changes. I would prefer this implementation to not depend on that.

You could change this to take e.g. start and end indices or offset and length parameters. You could also write an additional function e.g. areSymbolIdsUnique that calls into this for your use case.

[Groverkss]

• Address Arjun's comments

[Groverkss]

Addressed @arjunp's comments

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

465

Here and below: I think the end of the sentence is missing. Check what?

472

Same here.

[Groverkss]

• Fix doc comments

[Groverkss]

Fixed doc comments.

[arjunp]

Seems good to me. Minor comments below.

mlir/lib/Analysis/AffineStructures.cpp

457

Maybe assert that end < maybeValues.size()

609

Maybe not for this patch but it would be nice to also make this not depend on the variable layout, maybe calling into some function to get an IdKind from a location.

[Groverkss]

• Add asserts for areIdsUnique
• Use range based for loop in areIdsUnique

[Groverkss]

Addressed comments.

mlir/include/mlir/Analysis/AffineStructures.h

void setValues(unsigned start, unsigned end, ArrayRef<Value> values) {
     assert((start < numIds || end == start) && "invalid start position");
     assert(end <= numIds && "invalid end position");
     assert(values.size() == end - start);
     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 with Value as `None` are considered
-  /// to be inequal to all other symbols.
+  /// 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:
   /// 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 override;
 

mlir/lib/Analysis/AffineStructures.cpp

 
 /// Calls areIdsAligned to check if two constraint systems have the same set
 /// of identifiers in the same order.
 bool FlatAffineValueConstraints::areIdsAlignedWithOther(
     const FlatAffineValueConstraints &other) {
   return areIdsAligned(*this, other);
 }
 
-/// Checks if the SSA values associated with `cst`'s identifiers are unique.
-static bool LLVM_ATTRIBUTE_UNUSED
-areIdsUnique(const FlatAffineValueConstraints &cst) {
+/// 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) {
   SmallPtrSet<Value, 8> uniqueIds;
-  for (auto val : cst.getMaybeValues()) {
+  ArrayRef<Optional<Value>> maybeValues = cst.getMaybeValues();
+  for (unsigned i = start; i < end; ++i) {
+    Optional<Value> val = maybeValues[i];
     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 FlatAffineConstraints &cst) {
+  return areIdsUnique(cst, 0, cst.getNumIds());
+}
+
+/// Checks if the SSA values associated with `cst`'s identifiers of kind `kind`
+/// are unique.
+static bool LLVM_ATTRIBUTE_UNUSED areIdsUnique(
+    const FlatAffineValueConstraints &cst, FlatAffineConstraints::IdKind kind) {
+
+  if (kind == FlatAffineConstraints::IdKind::Dimension)
+    return areIdsUnique(cst, 0, cst.getNumDimIds());
+  if (kind == FlatAffineConstraints::IdKind::Symbol)
+    return areIdsUnique(cst, cst.getNumDimIds(), cst.getNumDimAndSymbolIds());
+  if (kind == FlatAffineConstraints::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 of each system are by design separate/local
 /// and are placed one after other (A's followed by B's).
 //  E.g.: Input: A has ((%i, %j) [%M, %N]) and B has (%k, %j) [%P, %N, %M])
static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value id) {
   if (cst->findId(id, &pos) && pos >= cst->getNumDimIds() &&
       pos < cst->getNumDimAndSymbolIds()) {
     cst->swapId(pos, cst->getNumDimIds());
     cst->setDimSymbolSeparation(cst->getNumSymbolIds() - 1);
   }
 }
 
 /// Merge and align symbols of `this` and `other` such that both get union of
-/// of symbols that are unique. Symbols with Value as `None` are considered
-/// to be inequal to all other symbols.
+/// 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 FlatAffineValueConstraints::mergeSymbolIds(
     FlatAffineValueConstraints &other) {
+
+  assert(areIdsUnique(*this, IdKind::Symbol) && "Symbol ids are not unique");
+  assert(areIdsUnique(other, IdKind::Symbol) && "Symbol ids are not unique");
+
   SmallVector<Value, 4> aSymValues;
   getValues(getNumDimIds(), getNumDimAndSymbolIds(), &aSymValues);
 
   // Merge symbols: merge symbols into `other` first from `this`.
   unsigned s = other.getNumDimIds();
   for (Value aSymValue : aSymValues) {
     unsigned loc;
     // If the id is a symbol in `other`, then align it, otherwise assume that
     // it is a new symbol
     if (other.findId(aSymValue, &loc) && loc >= other.getNumDimIds() &&
-        loc < getNumDimAndSymbolIds())
+        loc < other.getNumDimAndSymbolIds())
       other.swapId(s, loc);
     else
       other.insertSymbolId(s - other.getNumDimIds(), aSymValue);
     s++;
   }
 
   // Symbols that are in other, but not in this, are added at the end.
   for (unsigned t = other.getNumDimIds() + getNumSymbolIds(),
                 e = other.getNumDimAndSymbolIds();
        t < e; t++)
     insertSymbolId(getNumSymbolIds(), other.getValue(t));
 
   assert(getNumSymbolIds() == other.getNumSymbolIds() &&
          "expected same number of symbols");
+  assert(areIdsUnique(*this, IdKind::Symbol) && "Symbol ids are not unique");
+  assert(areIdsUnique(other, IdKind::Symbol) && "Symbol ids are not unique");
 }
 
 // Changes all symbol identifiers which are loop IVs to dim identifiers.
 void FlatAffineValueConstraints::convertLoopIVSymbolsToDims() {
   // Gather all symbols which are loop IVs.
   SmallVector<Value, 4> loopIVs;
   for (unsigned i = getNumDimIds(), e = getNumDimAndSymbolIds(); i < e; i++) {
     if (hasValue(i) && getForInductionVarOwner(getValue(i)))