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

[ftynse]

Summary

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.

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

 transform::SplitHandlesOp::apply(transform::TransformResults &results,
                                  transform::TransformState &state) {
   int64_t numResultHandles =
       getHandle() ? state.getPayloadOps(getHandle()).size() : 0;
   int64_t expectedNumResultHandles = getNumResultHandles();
   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.
   for (const auto &en : llvm::enumerate(state.getPayloadOps(getHandle())))
     results.set(getResults()[en.index()].cast<OpResult>(), en.value());
   return DiagnosedSilenceableFailure::success();
 }
 

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

 
 transform.sequence failures(propagate) {
 ^bb1(%fun: !pdl.operation):
   %muli = transform.structured.match ops{["arith.muli"]} in %fun : (!pdl.operation) -> !pdl.operation
   %h:2 = split_handles %muli in [2] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
   // 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)
 }
 
 // -----
 
 func.func @split_handles(%a: index, %b: index, %c: index) {
   %0 = arith.muli %a, %b : index
   %1 = arith.muli %a, %c : index
func.func @split_handles(%a: index, %b: index, %c: index) {
   %0 = arith.muli %a, %b : index
   %1 = arith.muli %a, %c : index
   return
 }
 
 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.
   yield %fun : !pdl.operation
 }
 
 // -----
 
+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
   // expected-remark @below {{0 : i32}}
   transform.test_print_param %0 : !transform.test_dialect_param
 }
 
 // -----