Modify TransformRewriter listener to get the match failure remark and use it to test failure in the op.

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

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

@@ -1074,10 +1074,18 @@ class ErrorCheckingTrackingListener : public TrackingListener {
   /// resets the error state to "success".
   DiagnosedSilenceableFailure checkAndResetError();
 
+  /// Return the latest match notification message.
+  std::string getLatestMatchFailureMessage();
+
   /// Return "true" if this tracking listener had a failure.
   bool failed() const;
 
 protected:
+
+  void
+  notifyMatchFailure(Location loc,
+    function_ref<void(Diagnostic &)> reasonCallback) override;
+
   void
   notifyPayloadReplacementNotFound(Operation *op, ValueRange values,
                                    DiagnosedSilenceableFailure &&diag) override;
@@ -1089,6 +1097,9 @@ class ErrorCheckingTrackingListener : public TrackingListener {
 
   /// The number of errors that have been encountered.
   int64_t errorCounter = 0;
+
+  /// Latest message from match failure notification.
+  std::string matchFailureMsg = "";
 };
 
 /// This is a special rewriter to be used in transform op implementations,

mlir/include/mlir/Transforms/RegionUtils.h

@@ -72,7 +72,8 @@ SmallVector<Value> makeRegionIsolatedFromAbove(
 
 /// Move SSA values used within an operation before an insertion point,
 /// so that the operation itself (or its replacement) can be moved to
-/// the insertion point.
+/// the insertion point. Current support is only for movement of
+/// dependencies of `op` before `insertionPoint` in the same basic block.
 LogicalResult moveOperationDependencies(RewriterBase &rewriter, Operation *op,
                                         Operation *insertionPoint,
                                         DominanceInfo &dominance);

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

@@ -1390,6 +1390,17 @@ void transform::ErrorCheckingTrackingListener::notifyPayloadReplacementNotFound(
   ++errorCounter;
 }
 
+std::string transform::ErrorCheckingTrackingListener::getLatestMatchFailureMessage() {
+  return matchFailureMsg;
+}
+
+void transform::ErrorCheckingTrackingListener::notifyMatchFailure(
+  Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
+  Diagnostic diag(loc, DiagnosticSeverity::Remark);
+  reasonCallback(diag);
+  matchFailureMsg = diag.str();
+}
+
 //===----------------------------------------------------------------------===//
 // TransformRewriter
 //===----------------------------------------------------------------------===//

mlir/lib/Transforms/Utils/RegionUtils.cpp

@@ -1073,13 +1073,21 @@ LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
         op, "unsupported caes where operation and insertion point are not in "
             "the same basic block");
   }
+  // If `insertionPoint` does not dominate `op`, do nothing
+  if (!dominance.properlyDominates(insertionPoint, op)) {
+    return rewriter.notifyMatchFailure(op,
+                                       "insertion point does not dominate op");
+  }
 
   // Find the backward slice of operation for each `Value` the operation
   // depends on. Prune the slice to only include operations not already
   // dominated by the `insertionPoint`
   BackwardSliceOptions options;
   options.inclusive = true;
   options.omitUsesFromAbove = false;
+  // Since current support is to only move within a same basic block,
+  // the slices dont need to look past block arguments.
+  options.omitBlockArguments = true;
   options.filter = [&](Operation *sliceBoundaryOp) {
     return !dominance.properlyDominates(sliceBoundaryOp, insertionPoint);
   };

mlir/test/Transforms/move-operation-deps.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt --allow-unregistered-dialect --transform-interpreter --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --allow-unregistered-dialect --transform-interpreter --split-input-file --verify-diagnostics %s | FileCheck %s
 
 // Check simple move of dependent operation before insertion.
 func.func @simple_move() -> f32 {
@@ -95,7 +95,7 @@ module attributes {transform.with_named_sequence} {
 func.func @move_region_dependencies() -> f32 {
   %0 = "before"() : () -> (f32)
   %1 = "moved_op_1"() : () -> (f32)
-  %2 = "moved_op"() ({
+  %2 = "moved_op_2"() ({
     "yield"(%1) : (f32) -> ()
   }) : () -> (f32)
   %3 = "foo"() ({
@@ -139,6 +139,7 @@ module attributes {transform.with_named_sequence} {
         : (!transform.any_op) -> !transform.any_op
     %op2 = transform.structured.match ops{["before"]} in %arg0
         : (!transform.any_op) -> !transform.any_op
+    // expected-remark@+1{{cannot move dependencies before operation in backward slice of op}}
     transform.test.move_operand_deps %op1 before %op2
         : !transform.any_op, !transform.any_op
     transform.yield
@@ -147,7 +148,9 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-func.func @move_region_dependencies() -> f32 {
+// Fail when the "before" operation is part of the operation slice (computed
+// when looking through implicitly captured values).
+func.func @do_not_move_slice() -> f32 {
   %0 = "before"() : () -> (f32)
   %1 = "moved_op"() ({
     "yield"(%0) : (f32) -> ()
@@ -164,6 +167,32 @@ module attributes {transform.with_named_sequence} {
         : (!transform.any_op) -> !transform.any_op
     %op2 = transform.structured.match ops{["before"]} in %arg0
         : (!transform.any_op) -> !transform.any_op
+    // expected-remark@+1{{cannot move dependencies before operation in backward slice of op}}
+    transform.test.move_operand_deps %op1 before %op2
+        : !transform.any_op, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// Dont move ops when insertion point does not dominate the op
+func.func @do_not_move() -> f32 {
+  %1 = "moved_op"() : () -> (f32)
+  %2 = "foo"() ({
+    "yield"(%1) : (f32) -> ()
+  }) : () -> (f32)
+  %3 = "before"() : () -> f32
+  return %2 : f32
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0 : !transform.any_op {transform.readonly}) {
+    %op1 = transform.structured.match ops{["foo"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    %op2 = transform.structured.match ops{["before"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    // expected-remark@+1{{insertion point does not dominate op}}
     transform.test.move_operand_deps %op1 before %op2
         : !transform.any_op, !transform.any_op
     transform.yield

mlir/test/lib/Transforms/TestTransformsOps.cpp

@@ -33,8 +33,8 @@ transform::TestMoveOperandDeps::apply(TransformRewriter &rewriter,
   Operation *moveBefore = *state.getPayloadOps(getInsertionPoint()).begin();
   if (failed(moveOperationDependencies(rewriter, op, moveBefore))) {
     auto listener = cast<ErrorCheckingTrackingListener>(rewriter.getListener());
-    std::string errorMsg = listener->checkAndResetError().getMessage();
-    return emitSilenceableFailure(op, errorMsg);
+    std::string errorMsg = listener->getLatestMatchFailureMessage();
+    (void)emitRemark(errorMsg);
   }
   return DiagnosedSilenceableFailure::success();
 }