[MLIR] FlatAffineConstraints: Ensure dimensionalities match when calling mergeLocalIds

[Groverkss]

Summary

This patch reorders mergeLocalIds usage to merge locals only after number of
dimensions and symbols are same. This does not change any functionality
because it does not matter in what order identifiers are merged, since
the reason to do it is to ensure that two FACs are aligned.

The order ensured in this patch simplifies a subsequent patch to improve
mergeLocalIds which requires dimensions and symbols to be aligned.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1877–1884

Do you rule out a use case for merging local ids where the number of dim and symbol ids do not match? Note that this is just merge the local id's in the system as opposed to touching the constraints.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1877–1884

Yes, that case is not supported by mergeLocalIds after this patch. Although, I don't see any use case as of now where the implementation cannot be changed to match dim and symbol ids before matching local ids. The reason to ensure this dimension/symbol matching is shown more in this child revision: https://reviews.llvm.org/D112867.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1877–1884

I too tend to agree - I don't see such uses cases. It's in general systematic to merge/align dims/symbols first and then merge local IDs. So this patch looks good.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   void removeRedundantConstraints();
 
   /// Converts identifiers in the column range [idStart, idLimit) to local
   /// variables.
   void convertDimToLocal(unsigned dimStart, unsigned dimLimit);
 
   /// 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.
+  /// of its local ids. Number of dimension and symbol ids should match in
+  /// `this` and `other`.
   void mergeLocalIds(FlatAffineConstraints &other);
 
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   void print(raw_ostream &os) const;
   void dump() const;
 

mlir/lib/Analysis/AffineAnalysis.cpp

for (unsigned i = 0, e = domain.getNumDimIds(); i < e; ++i) {
     } else {
       rel.insertDomainId(i);
       rel.setValue(i, domain.getValue(i));
     }
   }
 
   // Append domain constraints to `ret`.
   domainRel.appendRangeId(rel.getNumRangeDims());
-  domainRel.mergeLocalIds(rel);
   domainRel.mergeSymbolIds(rel);
+  domainRel.mergeLocalIds(rel);
   rel.append(domainRel);
 
   return success();
 }
 
 // Populates 'accessMap' with composition of AffineApplyOps reachable from
 // indices of MemRefAccess.
 void MemRefAccess::getAccessMap(AffineValueMap *accessMap) const {

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(); }));
 
-  // Bring A and B to common local space
-  a->mergeLocalIds(*b);
-
   SmallVector<Value, 4> aDimValues;
   a->getValues(offset, a->getNumDimIds(), &aDimValues);
 
   {
     // Merge dims from A into B.
     unsigned d = offset;
     for (auto aDimValue : aDimValues) {
       unsigned loc;
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 and align symbols of A and B
   a->mergeSymbolIds(*b);
+  // Merge and align local ids of A and B
+  a->mergeLocalIds(*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);
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.
+/// of its local ids. Number of dimension and symbol ids should match in
+/// `this` and `other`.
 void FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
+  assert(getNumDimIds() == other.getNumDimIds() &&
+         "Number of dimension ids should match");
+  assert(getNumSymbolIds() == other.getNumSymbolIds() &&
+         "Number of symbol ids should match");
   unsigned initLocals = getNumLocalIds();
   insertLocalId(getNumLocalIds(), other.getNumLocalIds());
   other.insertLocalId(0, initLocals);
 }
 
 /// Removes local variables using equalities. Each equality is checked if it
 /// can be reduced to the form: `e = affine-expr`, where `e` is a local
 /// variable and `affine-expr` is an affine expression not containing `e`.
 /// If an equality satisfies this form, the local variable is replaced in

 void FlatAffineRelation::compose(const FlatAffineRelation &other) {
   assert(getNumDomainDims() == other.getNumRangeDims() &&
          "Domain of this and range of other do not match");
   assert(std::equal(values.begin(), values.begin() + getNumDomainDims(),
                     other.values.begin() + other.getNumDomainDims()) &&
          "Domain of this and range of other do not match");
 
   FlatAffineRelation rel = other;
+
+  // Convert `rel` from
+  //    [otherDomain] -> [otherRange]
+  // to
+  //    [otherDomain] -> [otherRange thisRange]
+  // and `this` from
+  //    [thisDomain] -> [thisRange]
+  // to
+  //    [otherDomain thisDomain] -> [thisRange].
+  unsigned removeDims = rel.getNumRangeDims();
+  insertDomainId(0, rel.getNumDomainDims());
+  rel.appendRangeId(getNumRangeDims());
+
+  // Merge symbol and local identifiers.
   mergeSymbolIds(rel);
   mergeLocalIds(rel);
 
-  // Convert domain of `this` and range of `rel` to local identifiers.
-  convertDimToLocal(0, getNumDomainDims());
-  rel.convertDimToLocal(rel.getNumDomainDims(), rel.getNumDimIds());
-  // Add dimensions such that both relations become `domainRel -> rangeThis`.
-  appendDomainId(rel.getNumDomainDims());
-  rel.appendRangeId(getNumRangeDims());
+  // Convert `rel` from [otherDomain] -> [otherRange thisRange] to
+  // [otherDomain] -> [thisRange] by converting first otherRange range ids
+  // to local ids.
+  rel.convertDimToLocal(rel.getNumDomainDims(),
+                        rel.getNumDomainDims() + removeDims);
+  // Convert `this` from [otherDomain thisDomain] -> [thisRange] to
+  // [otherDomain] -> [thisRange] by converting last thisDomain domain ids
+  // to local ids.
+  convertDimToLocal(getNumDomainDims() - removeDims, getNumDomainDims());
 
   auto thisMaybeValues = getMaybeDimValues();
   auto relMaybeValues = rel.getMaybeDimValues();
 
   // Add and match domain of `rel` to domain of `this`.
   for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
     if (relMaybeValues[i].hasValue())
       setValue(i, relMaybeValues[i].getValue());