[MLIR] Add mergeLocalIds and mergeSymbolIds

[Groverkss]

Summary

This patch adds mergeLocalIds andmergeSymbolIds as public functions
for FlatAffineConstraints and FlatAffineValueConstraints respectively.

mergeLocalIds is also required to support divisions in intersection,
subtraction, equality checks, and complement for PresburgerSet.

This patch is part of a series of patches aimed at generalizing affine
dependence analysis.

[bondhugula]

mlir/include/mlir/Analysis/AffineStructures.h

432

The name of this method doesn't reflect its behavior. Given the doc comment, shouldn't this be named: mergeLocalIds?

852

Likewise. Prefer to name the methods imperatively.

mlir/lib/Analysis/AffineStructures.cpp

561

Avoid auto here. As compact and more readable without auto.

565

findId -> llvm::is_contained

[Groverkss]

• Addressed bondhugula's comments

[Groverkss]

@bondhugula Thank you for the review!

mlir/lib/Analysis/AffineStructures.cpp

565

I don't think llvm::is_contained can be used here since findIf sets loc to the position of the id and is used for swapping.

mlir/include/mlir/Analysis/AffineStructures.h

public:
 
   /// Removes redundant constraints using Simplex. Although the algorithm can
   /// theoretically take exponential time in the worst case (rare), it is known
   /// to perform much better in the average case. If V is the number of vertices
   /// in the polytope and C is the number of constraints, the algorithm takes
   /// O(VC) time.
   void removeRedundantConstraints();
 
+  /// Merge local ids of `this` and `other`. This is done by appending local ids
+  /// of `other` to `this` and inserting local ids of `this` to `other` at start
+  /// of its local ids.
+  void mergeLocalIds(FlatAffineConstraints &other);
+
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Return the index at which the specified kind of id starts.
public:
   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.
+  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;
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes

mlir/lib/Analysis/AffineStructures.cpp

static void mergeAndAlignIds(unsigned offset, FlatAffineValueConstraints *a,
   assert(std::all_of(a->getMaybeValues().begin() + offset,
                      a->getMaybeValues().begin() + a->getNumDimAndSymbolIds(),
                      [](Optional<Value> id) { return id.hasValue(); }));
 
   assert(std::all_of(b->getMaybeValues().begin() + offset,
                      b->getMaybeValues().begin() + b->getNumDimAndSymbolIds(),
                      [](Optional<Value> id) { return id.hasValue(); }));
 
-  // Place local id's of A after local id's of B.
-  b->insertLocalId(/*pos=*/0, /*num=*/a->getNumLocalIds());
-  a->appendLocalId(/*num=*/b->getNumLocalIds() - a->getNumLocalIds());
+  // Bring A and B to common local space
+  a->mergeLocalIds(*b);
 
-  SmallVector<Value, 4> aDimValues, aSymValues;
+  SmallVector<Value, 4> aDimValues;
   a->getValues(offset, a->getNumDimIds(), &aDimValues);
-  a->getValues(a->getNumDimIds(), a->getNumDimAndSymbolIds(), &aSymValues);
 
   {
     // Merge dims from A into B.
     unsigned d = offset;
     for (auto aDimValue : aDimValues) {
       unsigned loc;
       if (b->findId(aDimValue, &loc)) {
         assert(loc >= offset && "A's dim appears in B's aligned range");
         assert(loc < b->getNumDimIds() &&
                "A's dim appears in B's non-dim position");
         b->swapId(d, loc);
       } else {
         b->insertDimId(d, aDimValue);
       }
       d++;
     }
     // Dimensions that are in B, but not in A, are added at the end.
     for (unsigned t = a->getNumDimIds(), e = b->getNumDimIds(); t < e; t++) {
       a->appendDimId(b->getValue(t));
     }
     assert(a->getNumDimIds() == b->getNumDimIds() &&
            "expected same number of dims");
   }
 
-  {
-    // Merge symbols: merge A's symbols into B first.
-    unsigned s = 0;
-    for (auto aSymValue : aSymValues) {
-      unsigned loc;
-      if (b->findId(aSymValue, &loc)) {
-        assert(loc >= b->getNumDimIds() && loc < b->getNumDimAndSymbolIds() &&
-               "A's symbol appears in B's non-symbol position");
-        b->swapId(s + b->getNumDimIds(), loc);
-      } else {
-        b->insertSymbolId(s, aSymValue);
-      }
-      s++;
-    }
-    // Symbols that are in B, but not in A, are added at the end.
-    for (unsigned t = a->getNumDimAndSymbolIds(),
-                  e = b->getNumDimAndSymbolIds();
-         t < e; t++) {
-      a->appendSymbolId(b->getValue(t));
-    }
-    assert(a->getNumDimAndSymbolIds() == b->getNumDimAndSymbolIds() &&
-           "expected same number of dims and symbols");
-  }
+  // Merge and align symbols of A and B
+  a->mergeSymbolIds(*b);
 
   assert(areIdsAligned(*a, *b) && "IDs expected to be aligned");
 }
 
 // Call 'mergeAndAlignIds' to align constraint systems of 'this' and 'other'.
 void FlatAffineValueConstraints::mergeAndAlignIdsWithOther(
     unsigned offset, FlatAffineValueConstraints *other) {
   mergeAndAlignIds(offset, this, other);
static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value id) {
   unsigned pos;
   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.
+void FlatAffineValueConstraints::mergeSymbolIds(
+    FlatAffineValueConstraints &other) {
+  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())
+      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");
+}
+
 // 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)))
       loopIVs.push_back(getValue(i));
   }
void FlatAffineConstraints::removeRedundantConstraints() {
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
   equalities.resizeVertically(pos);
 }
 
+/// Merge local ids of `this` and `other`. This is done by appending local ids
+/// of `other` to `this` and inserting local ids of `this` to `other` at start
+/// of its local ids.
+void FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
+  unsigned initLocals = getNumLocalIds();
+  insertLocalId(getNumLocalIds(), other.getNumLocalIds());
+  other.insertLocalId(0, initLocals);
+}
+
 std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
     unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
     ArrayRef<AffineExpr> localExprs, MLIRContext *context) const {
   assert(pos + offset < getNumDimIds() && "invalid dim start pos");
   assert(symStartPos >= (pos + offset) && "invalid sym start pos");
   assert(getNumLocalIds() == localExprs.size() &&
          "incorrect local exprs count");