Adds offset argument to specified range of ids know to be aligned when calling mergeA...

Supports use case where FlatAffineConstraints::composeMap adds dim identifiers with no SSA values (because the identifiers are the result of an AffineValueMap which is not materialized in the IR and thus has no SSA Value results).

PiperOrigin-RevId: 237145506

diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 543ae1400d4747c305660a561a7047303eeb5fbc..7575468d455c8ce7402a85518106afee2f494bea 100644 (file)
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -519,28 +519,29 @@ static void swapId(FlatAffineConstraints *A, unsigned posA, unsigned posB) {
   std::swap(A->getId(posA), A->getId(posB));
 }
 
-/// Merge and align the identifiers of A and B 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).
+/// 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).
 //  Eg: 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]
 //
 // TODO(mlir-team): expose this function at some point.
-static void mergeAndAlignIds(FlatAffineConstraints *A,
+static void mergeAndAlignIds(unsigned offset, FlatAffineConstraints *A,
                              FlatAffineConstraints *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 id values aren't unique");
   assert(areIdsUnique(*B) && "B's id values aren't unique");
 
-  assert(std::all_of(A->getIds().begin(),
+  assert(std::all_of(A->getIds().begin() + offset,
                      A->getIds().begin() + A->getNumDimAndSymbolIds(),
                      [](Optional<Value *> id) { return id.hasValue(); }));
 
-  assert(std::all_of(B->getIds().begin(),
+  assert(std::all_of(B->getIds().begin() + offset,
                      B->getIds().begin() + B->getNumDimAndSymbolIds(),
                      [](Optional<Value *> id) { return id.hasValue(); }));
 
@@ -554,14 +555,15 @@ static void mergeAndAlignIds(FlatAffineConstraints *A,
   }
 
   SmallVector<Value *, 4> aDimValues, aSymValues;
-  A->getIdValues(0, A->getNumDimIds(), &aDimValues);
+  A->getIdValues(offset, A->getNumDimIds(), &aDimValues);
   A->getIdValues(A->getNumDimIds(), A->getNumDimAndSymbolIds(), &aSymValues);
   {
     // Merge dims from A into B.
-    unsigned d = 0;
+    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");
         swapId(B, d, loc);
@@ -623,7 +625,7 @@ bool FlatAffineConstraints::composeMap(AffineValueMap *vMap) {
                                    vMap->getOperands().end());
     localCst.setIdValues(0, localCst.getNumDimAndSymbolIds(), values);
     // Align localCst and this.
-    mergeAndAlignIds(&localCst, this);
+    mergeAndAlignIds(/*offset=*/0, &localCst, this);
     // Finally, append localCst to this constraint set.
     append(localCst);
   }
@@ -772,7 +774,7 @@ bool FlatAffineConstraints::addAffineForOpDomain(
     }
     // Merge and align with localVarCst.
     if (localVarCst.getNumLocalIds() > 0) {
-      mergeAndAlignIds(this, &localVarCst);
+      mergeAndAlignIds(/*offset=*/0, this, &localVarCst);
       append(localVarCst);
     }
 
@@ -1721,7 +1723,7 @@ bool FlatAffineConstraints::addSliceBounds(ArrayRef<Value *> values,
           }
         }
       }
-      mergeAndAlignIds(this, &localVarCst);
+      mergeAndAlignIds(/*offset=*/0, this, &localVarCst);
       append(localVarCst);
     }
 
@@ -2692,7 +2694,7 @@ bool FlatAffineConstraints::unionBoundingBox(
   Optional<FlatAffineConstraints> otherCopy;
   if (!areIdsAligned(*this, otherCst)) {
     otherCopy.emplace(FlatAffineConstraints(otherCst));
-    mergeAndAlignIds(this, &otherCopy.getValue());
+    mergeAndAlignIds(/*offset=*/numDims, this, &otherCopy.getValue());
   }
 
   const auto &other = otherCopy ? *otherCopy : otherCst;