[Mlir-commits] [mlir] 30ce0c9 - [mlir] properly handle an edge case in transform.split_handles

[Alex Zinenko]

Author: Alex Zinenko
Date: 2023-05-05T12:48:20Z
New Revision: 30ce0c9508f50f1b7bce4ff637234084bb7a571d

URL: https://github.com/llvm/llvm-project/commit/30ce0c9508f50f1b7bce4ff637234084bb7a571d
DIFF: https://github.com/llvm/llvm-project/commit/30ce0c9508f50f1b7bce4ff637234084bb7a571d.diff

LOG: [mlir] properly handle an edge case in transform.split_handles

The `transform.split_handles` transform op has an edge case that is
supposed to propagate empty handle state to all resulting handles
regardless of their number. Actually set empty payload for them instead
of leaving them unset.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D149652

Added: 
    

Modified: 
    mlir/lib/Dialect/Transform/IR/TransformOps.cpp
    mlir/test/Dialect/Transform/test-interpreter.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 6284b41302e2d..138f4ff69cfbc 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -1510,13 +1510,17 @@ transform::SplitHandlesOp::apply(transform::TransformResults &results,
   if (numResultHandles != expectedNumResultHandles) {
     // Empty input handle corner case: always propagates empty handles in both
     // suppress and propagate modes.
-    if (numResultHandles == 0)
+    if (numResultHandles == 0) {
+      for (OpResult result : getResults())
+        results.set(result, {});
       return DiagnosedSilenceableFailure::success();
+    }
+
     // If the input handle was not empty and the number of result handles does
     // not match, this is a legit silenceable error.
     return emitSilenceableError()
            << getHandle() << " expected to contain " << expectedNumResultHandles
-           << " operation handles but it only contains " << numResultHandles
+           << " operation handles but it contains " << numResultHandles
            << " handles";
   }
   // Normal successful case.

diff  --git a/mlir/test/Dialect/Transform/test-interpreter.mlir b/mlir/test/Dialect/Transform/test-interpreter.mlir
index b80ee22dd291b..3668abb1c3a42 100644
--- a/mlir/test/Dialect/Transform/test-interpreter.mlir
+++ b/mlir/test/Dialect/Transform/test-interpreter.mlir
@@ -831,7 +831,7 @@ transform.sequence failures(propagate) {
   // expected-remark @below {{1}}
   transform.test_print_number_of_associated_payload_ir_ops %h#0
   %muli_2 = transform.structured.match ops{["arith.muli"]} in %fun : (!pdl.operation) -> !pdl.operation
-  // expected-error @below {{expected to contain 3 operation handles but it only contains 2 handles}}
+  // expected-error @below {{expected to contain 3 operation handles but it contains 2 handles}}
   %h_2:3 = split_handles %muli_2 in [3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
@@ -975,7 +975,7 @@ func.func @split_handles(%a: index, %b: index, %c: index) {
 transform.sequence -> !pdl.operation failures(propagate) {
 ^bb1(%fun: !pdl.operation):
   %muli = transform.structured.match ops{["arith.muli"]} in %fun : (!pdl.operation) -> !pdl.operation
-  // expected-error @below {{expected to contain 3 operation handles but it only contains 2 handles}}
+  // expected-error @below {{expected to contain 3 operation handles but it contains 2 handles}}
   %h_2:3 = split_handles %muli in [3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
   /// Test that yield does not crash in the presence of silenceable error in
   /// propagate mode.
@@ -984,6 +984,17 @@ transform.sequence -> !pdl.operation failures(propagate) {
 
 // -----
 
+transform.sequence -> !transform.any_op failures(suppress) {
+^bb0(%arg0: !transform.any_op):
+  %muli = transform.structured.match ops{["arith.muli"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // Edge case propagating empty handles in splitting.
+  %0:3 = split_handles %muli in [3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+  // Test does not crash when accessing the empty handle.
+  yield %0#0 : !transform.any_op
+}
+
+// -----
+
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   %0 = transform.test_produce_integer_param_with_type i32 : !transform.test_dialect_param