[mlir][transform] Add notifyPayloadOperationReplaced to TransformRewriter

[springerm]

Summary

This function allows users to update payload op mappings in cases where such replacements cannot be performed automatically by the rewriter/listener interface.

[ftynse]

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

994

Nit: s/manually/automatically

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

771–772

I'd rather return definite failure here.

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

883

Nit: use for (auto &&[original, replacement] : zip) instead of obnoxious std::get below.

[springerm]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

771–772

If this fails it means that tileAndFuseFirstExtractUse was updated without updating the remaining code base. (It only fails if the original op is not tracked.) I think this can stay as an assertion.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

protected:
 
   /// Return the single op that defines all given values (if any).
   static Operation *getCommonDefiningOp(ValueRange values);
 
   /// Return the transform op in which this TrackingListener is used.
   TransformOpInterface getTransformOp() const { return transformOp; }
 
 private:
+  friend class TransformRewriter;
+
   void notifyOperationRemoved(Operation *op) override;
 
   void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
+  using Listener::notifyOperationReplaced;
 
   /// The transform op in which this TrackingListener is used.
   TransformOpInterface transformOp;
 
   /// The handles that are consumed by the transform op.
   DenseSet<Value> consumedHandles;
 };
 

 
 public:
   /// Return "true" if the tracking listener had failures.
   bool hasTrackingFailures() const;
 
   /// Silence all tracking failures that have been encountered so far.
   void silenceTrackingFailure();
 
+  /// Notify the transform dialect interpreter that the given op has been
+  /// replaced with another op and that the mapping between handles and payload
+  /// ops/values should be updated. This function should be called before the
+  /// original op is erased. It fails if the operation could not be replaced,
+  /// e.g., because the original operation is not tracked.
+  ///
+  /// Note: As long as IR modifications are performed through this rewriter,
+  /// the transform state is usually updated automatically. This function should
+  /// be used when unsupported rewriter API is used; e.g., updating all uses of
+  /// a tracked operation one-by-one instead of using `RewriterBase::replaceOp`.
+  LogicalResult notifyPayloadOperationReplaced(Operation *op,
+                                               Operation *replacement);
+
 private:
   ErrorCheckingTrackingListener *const listener;
 };
 
 /// This trait is supposed to be attached to Transform dialect operations that
 /// can be standalone top-level transforms. Such operations typically contain
 /// other Transform dialect operations that can be executed following some
 /// control flow logic specific to the current operation. The operations with

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

static Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
   return fusedOp;
 }
 
 bool transform::FuseIntoContainingOp::allowsRepeatedHandleOperands() {
   // Allow repeated handles since we are fusing everything anyway.
   return true;
 }
 
-namespace {
-/// Unsafely exposes an internal protected method of TransformState::Extension
-/// as public.
-///
-/// MUST NOT be used directly.
-class UnsafeOpReplacementStateExtension : public TransformState::Extension {
-public:
-  UnsafeOpReplacementStateExtension(TransformState &state)
-      : TransformState::Extension(state) {}
-
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
-      UnsafeOpReplacementStateExtension)
-
-  LogicalResult doReplacePayloadOp(Operation *op, Operation *replacement) {
-    return replacePayloadOp(op, replacement);
-  }
-};
-} // namespace
-
-/// Replaces `payload` with `replacement` in all handles stored in the state.
-/// MUST NOT be used except for the case immediately below.
-static void forciblyReplaceReferencedPayloadOperation(TransformState &state,
-                                                      Operation *payload,
-                                                      Operation *replacement) {
-  UnsafeOpReplacementStateExtension extension(state);
-  // This may return failure if the payload is not associated with any handle,
-  // ignore that.
-  (void)extension.doReplacePayloadOp(payload, replacement);
-}
-
 DiagnosedSilenceableFailure
 transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
                                        transform::TransformResults &results,
                                        transform::TransformState &state) {
   SmallVector<Operation *> fusedOps;
   auto producerOps = state.getPayloadOps(getProducerOp());
   auto containingOps = state.getPayloadOps(getContainingOp());
   if (!llvm::hasSingleElement(containingOps)) {
while (!remainingProducers.empty()) {
     // Futhermore, producers should then be traversed according to a
     // topological sorting.
     auto [tiledOps, newContainingOp] =
         tileAndFuseFirstExtractUse(rewriter, diag, producerOp, containingOp);
     if (!tiledOps.empty()) {
       LLVM_DEBUG(DBGS() << "\nFused a direct extract use\n" << *containingOp);
       fusedOps.append(tiledOps);
       if (newContainingOp) {
+        // Update handles associated with the containing op so we don't need to
+        // invalidate them. This is a hack to support better composability
+        // between tiling and fusion while a proper mechanism is being
+        // investigated.
+        //
+        // DO NOT replicate this elsewhere unless you understand what you are
+        // doing.
+        LogicalResult replacementStatus =
+            rewriter.notifyPayloadOperationReplaced(containingOp,
+                                                    newContainingOp);
+        (void)replacementStatus;
+        assert(succeeded(replacementStatus) &&
+               "unable to update transform state mapping");
         rewriter.eraseOp(containingOp);
         containingOp = newContainingOp;
       }
       continue;
     }
 
     SmallVector<Operation *> tiledContainingOpOperand =
         tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
while (!remainingProducers.empty()) {
     if (cloned) {
       LLVM_DEBUG(DBGS() << "\nFused an use by cloning\n" << *containingOp);
       fusedOps.push_back(cloned);
       continue;
     }
     return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
   }
 
-  // Update handles associated with the containing op so we don't need to
-  // invalidate them. This is a hack to support better composability between
-  // tiling and fusion while a proper mechanism is being investigated.
-  //
-  // DO NOT replicate this elsewhere unless you understand what you are doing.
-  forciblyReplaceReferencedPayloadOperation(state, *containingOps.begin(),
-                                            containingOp);
-
   results.set(cast<OpResult>(getFusedOp()), fusedOps);
   results.set(cast<OpResult>(getNewContainingOp()), {containingOp});
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::FuseIntoContainingOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getProducerOp(), effects);

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

 /// Silence all tracking failures that have been encountered so far.
 void transform::TransformRewriter::silenceTrackingFailure() {
   if (hasTrackingFailures()) {
     DiagnosedSilenceableFailure status = listener->checkAndResetError();
     (void)status.silence();
   }
 }
 
+LogicalResult transform::TransformRewriter::notifyPayloadOperationReplaced(
+    Operation *op, Operation *replacement) {
+  return listener->replacePayloadOp(op, replacement);
+}
+
 //===----------------------------------------------------------------------===//
 // Utilities for TransformEachOpTrait.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 transform::detail::checkNestedConsumption(Location loc,
                                           ArrayRef<Operation *> targets) {
   for (auto &&[position, parent] : llvm::enumerate(targets)) {

mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir

-// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s -verify-diagnostics | FileCheck %s
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
 
 module {
   // CHECK-LABEL: func.func @fuse_tileable_op
   //  CHECK-SAME:   %[[CHUNK_SIZE:[0-9a-z]+]]: index
%0:2 = linalg.generic {
         linalg.yield %d, %e : f32, f32
     } -> (tensor<?xf32>, tensor<?xf32>)
     %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>
 
     %1 = affine.apply #map0()[%d0, %idx]
 
     // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])
     // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {
+    // expected-remark @below{{new containing op}}
     %2 = scf.forall (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {
       // CHECK: %[[I0:.*]] = affine.apply {{.*}}
       %3 = affine.apply #map1(%i)[%idx]
       // CHECK: %[[I1:.*]] = affine.min {{.*}}
       %4 = affine.min #map2(%i)[%d0, %idx]
       %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
 
       // CHECK: %[[T1:.*]]:2 = linalg.generic {{.*}}
module {
   }
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.generic is tileable. The op is tiled and fused.
-    transform.structured.fuse_into_containing_op %0 into %1
+    %fused, %containing = transform.structured.fuse_into_containing_op %0 into %1
       : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
+    test_print_remark_at_operand %containing, "new containing op" : !transform.any_op
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>

mlir/test/Dialect/Transform/test-interpreter.mlir

^bb1(%arg0: !transform.any_op):
   %0 = transform.structured.match ops{["test.annotate_me"]} in %arg0 : (!transform.any_op) -> !transform.any_op
   %1 = transform.test_produce_param_with_number_of_test_ops %0 : !transform.any_op
   transform.annotate %0 "new_attr" = %1 : !transform.any_op, !transform.test_dialect_param
 
   %2 = transform.param.constant 2 -> !transform.param<i64>
   transform.annotate %0 "broadcast_attr" = %2 : !transform.any_op, !transform.param<i64>
   transform.annotate %0 "unit_attr" : !transform.any_op
 }
+
+// -----
+
+func.func @notify_payload_op_replaced(%arg0: index, %arg1: index) {
+  %0 = arith.muli %arg0, %arg1 {original} : index
+  // expected-remark @below{{updated handle}}
+  %1 = arith.muli %arg0, %arg1 {replacement} : index
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match attributes{original} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match attributes{replacement} in %arg1 : (!transform.any_op) -> !transform.any_op
+  test_notify_payload_op_replaced %0, %1 : (!transform.any_op, !transform.any_op) -> ()
+  test_print_remark_at_operand %0, "updated handle" : !transform.any_op
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 
 LogicalResult mlir::test::TestReEnterRegionOp::verify() {
   if (getNumOperands() != getBody().front().getNumArguments()) {
     return emitOpError() << "expects as many operands as block arguments";
   }
   return success();
 }
 
+DiagnosedSilenceableFailure mlir::test::TestNotifyPayloadOpReplacedOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
+  auto originalOps = state.getPayloadOps(getOriginal());
+  auto replacementOps = state.getPayloadOps(getReplacement());
+  if (llvm::range_size(originalOps) != llvm::range_size(replacementOps))
+    return emitSilenceableError() << "expected same number of original and "
+                                     "replacement payload operations";
+  for (const auto &[original, replacement] :
+       llvm::zip(originalOps, replacementOps)) {
+    if (failed(
+            rewriter.notifyPayloadOperationReplaced(original, replacement))) {
+      auto diag = emitSilenceableError()
+                  << "unable to replace payload op in transform mapping";
+      diag.attachNote(original->getLoc()) << "original payload op";
+      diag.attachNote(replacement->getLoc()) << "replacement payload op";
+      return diag;
+    }
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestNotifyPayloadOpReplacedOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getOriginal(), effects);
+  transform::onlyReadsHandle(getReplacement(), effects);
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

: Op<Transform_Dialect, "test_re_enter_region",
         DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins Variadic<AnyType>:$args);
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat = "($args^ `:` type($args))? attr-dict-with-keyword regions";
   let cppNamespace = "::mlir::test";
   let hasVerifier = 1;
 }
 
+def TestNotifyPayloadOpReplacedOp
+  : Op<Transform_Dialect, "test_notify_payload_op_replaced",
+       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$original,
+                       TransformHandleTypeInterface:$replacement);
+  let results = (outs);
+  let assemblyFormat = "$original `,` $replacement attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD