Fix SSA ordering issue, add test for it. Also, update to a different handling of the residual

Author: krzysz00

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -5199,7 +5199,11 @@ struct CancelLinearizeOfDelinearizePortion final
       return rewriter.notifyMatchFailure(
           linearizeOp, "no run of delinearize outputs to deal with");
 
-    SmallVector<std::tuple<Value, Value>> delinearizeReplacements;
+    // Record all the delinearize replacements so we can do them after creating
+    // the new linearization operation, since the new operation might use
+    // outputs of something we're replacing.
+    SmallVector<SmallVector<Value>> delinearizeReplacements;
+
     SmallVector<Value> newIndex;
     newIndex.reserve(numLinArgs);
     SmallVector<OpFoldResult> newBasis;
@@ -5212,18 +5216,26 @@ struct CancelLinearizeOfDelinearizePortion final
       // Update here so we don't forget this during early continues
       prevMatchEnd = m.linStart + m.length;
 
+      PatternRewriter::InsertionGuard g(rewriter);
+      rewriter.setInsertionPoint(m.delinearize);
+
+      ArrayRef<OpFoldResult> basisToMerge =
+          linBasisRef.slice(m.linStart, m.length);
       // We use the slice from the linearize's basis above because of the
       // "bounds inferred from `disjoint`" case above.
       OpFoldResult newSize =
-          computeProduct(linearizeOp.getLoc(), rewriter,
-                         linBasisRef.slice(m.linStart, m.length));
+          computeProduct(linearizeOp.getLoc(), rewriter, basisToMerge);
 
       // Trivial case where we can just skip past the delinearize all together
       if (m.length == m.delinearize.getNumResults()) {
         newIndex.push_back(m.delinearize.getLinearIndex());
         newBasis.push_back(newSize);
+        // Pad out set of replacements so we don't do anything with this one.
+        delinearizeReplacements.push_back(SmallVector<Value>());
         continue;
       }
+
+      SmallVector<Value> newDelinResults;
       SmallVector<OpFoldResult> newDelinBasis = m.delinearize.getPaddedBasis();
       newDelinBasis.erase(newDelinBasis.begin() + m.delinStart,
                           newDelinBasis.begin() + m.delinStart + m.length);
@@ -5232,31 +5244,39 @@ struct CancelLinearizeOfDelinearizePortion final
           m.delinearize.getLoc(), m.delinearize.getLinearIndex(),
           newDelinBasis);
 
+      // Since there may be other uses of the indices we just merged together,
+      // create a residual affine.delinearize_index that delinearizes the
+      // merged output into its component parts.
+      Value combinedElem = newDelinearize.getResult(m.delinStart);
+      auto residualDelinearize = rewriter.create<AffineDelinearizeIndexOp>(
+          m.delinearize.getLoc(), combinedElem, basisToMerge);
+
       // Swap all the uses of the unaffected delinearize outputs to the new
       // delinearization so that the old code can be removed if this
       // linearize_index is the only user of the merged results.
+      llvm::append_range(newDelinResults,
+                         newDelinearize.getResults().take_front(m.delinStart));
+      llvm::append_range(newDelinResults, residualDelinearize.getResults());
       llvm::append_range(
-          delinearizeReplacements,
-          llvm::zip_equal(
-              m.delinearize.getResults().take_front(m.delinStart),
-              newDelinearize.getResults().take_front(m.delinStart)));
-      llvm::append_range(
-          delinearizeReplacements,
-          llvm::zip_equal(
-              m.delinearize.getResults().drop_front(m.delinStart + m.length),
-              newDelinearize.getResults().drop_front(m.delinStart + 1)));
+          newDelinResults,
+          newDelinearize.getResults().drop_front(m.delinStart + 1));
 
-      Value newLinArg = newDelinearize.getResult(m.delinStart);
-      newIndex.push_back(newLinArg);
+      delinearizeReplacements.push_back(newDelinResults);
+      newIndex.push_back(combinedElem);
       newBasis.push_back(newSize);
     }
     llvm::append_range(newIndex, multiIndex.drop_front(prevMatchEnd));
     llvm::append_range(newBasis, linBasisRef.drop_front(prevMatchEnd));
     rewriter.replaceOpWithNewOp<AffineLinearizeIndexOp>(
         linearizeOp, newIndex, newBasis, linearizeOp.getDisjoint());
 
-    for (auto [from, to] : delinearizeReplacements)
-      rewriter.replaceAllUsesWith(from, to);
+    for (auto [m, newResults] :
+         llvm::zip_equal(matches, delinearizeReplacements)) {
+      if (newResults.empty())
+        continue;
+      rewriter.replaceOp(m.delinearize, newResults);
+    }
+
     return success();
   }
 };

mlir/test/Dialect/Affine/canonicalize.mlir

@@ -2172,6 +2172,27 @@ func.func @partial_cancel_linearize_delinearize_not_fully_permuted(%arg0: index,
 
 // -----
 
+// Ensure we don't get SSA errors when creating new `affine.delinearize` operations.
+// CHECK-LABEL: func @cancel_linearize_delinearize_placement
+// CHECK-SAME: (%[[ARG0:.+]]: index)
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[NEW_DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (8, 32) : index, index
+// CHECK-NEXT: %[[DELIN_PART:.+]]:2 = affine.delinearize_index %[[NEW_DELIN]]#1 into (8, 4) : index, index
+// CHECK-NEXT: %[[L1:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#1, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (4, 8, 4, 8)
+// CHECK-NEXT: %[[L2:.+]] = affine.linearize_index disjoint [%[[NEW_DELIN]]#1, %[[C0]], %[[C0]]] by (32, 8, 4)
+// CHECK-NEXT: %[[L3:.+]] = affine.linearize_index disjoint [%[[DELIN_PART]]#0, %[[NEW_DELIN]]#0, %[[C0]], %[[C0]]] by (8, 8, 4, 4)
+// CHECK-NEXT: return %[[L1]], %[[L2]], %[[L3]]
+func.func @cancel_linearize_delinearize_placement(%arg0: index) -> (index, index, index) {
+  %c0 = arith.constant 0 : index
+  %0:3 = affine.delinearize_index %arg0 into (8, 8, 4) : index, index, index
+  %1 = affine.linearize_index disjoint [%0#2, %0#0, %c0, %c0] by (4, 8, 4, 8) : index
+  %2 = affine.linearize_index disjoint [%0#1, %0#2, %c0, %c0] by (8, 4, 8, 4) : index
+  %3 = affine.linearize_index disjoint [%0#1, %0#0, %c0, %c0] by (8, 8, 4, 4) : index
+  return %1, %2, %3 : index, index, index
+}
+
+// -----
+
 // Won't cancel because the linearize and delinearize are using a different basis
 // CHECK-LABEL: func @no_cancel_linearize_delinearize_different_basis(
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,