[mlir][transform] Clean up SplitHandlesOp

[springerm]

Summary

• Rename to SplitHandleOp: it splits a single handle.
• Drop num_result_handles attribute: it is redundant and can be inferred from the number of results.
• Improve documentation and minor code cleanups.

Depends On: D149929

[ftynse]

While we are at it, could we generalize this operation a bit? I'd like to see additional behavior, potentially controlled by attributes:

• if there are more ops than handles, put all the remaining ops into the first/last handle instead of failing (depending on the attribute);
• if there are less ops than handles, keep trailing handles empty (default, can be disabled by attribute);
• maybe some wraparound where each n-th op goes to n-th handle.

I have always found this op barely usable and very narrow as we rarely know the exact number of operations associated with a handle in the general case.

Also, https://reviews.llvm.org/D149652.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

526–529

This bit of the documentation is outdated.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

::mlir::ArrayAttr getCallableArgAttrs() {
       return getArgAttrs().value_or(nullptr);
     }
     ::mlir::ArrayAttr getCallableResAttrs() {
       return getResAttrs().value_or(nullptr);
     }
   }];
 }
 
-def SplitHandlesOp : TransformDialectOp<"split_handles",
+def SplitHandleOp : TransformDialectOp<"split_handle",
     [FunctionalStyleTransformOpTrait,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let summary = "Splits handles from a union of payload ops to a list";
+  let summary = "Splits a handle of payload ops into handles with a single op";
   let description = [{
-    Creates `num_result_handles` transform IR handles extracted from the
-    `handle` operand. The resulting Payload IR operation handles are listed
-    in the same order as the operations appear in the source `handle`.
-    This is useful for ensuring a statically known number of operations are
-    tracked by the source `handle` and to extract them into individual handles
-    that can be further manipulated in isolation.
-
-    This operation succeeds and returns `num_result_handles` if the statically
-    specified `num_result_handles` corresponds to the dynamic number of
-    operations contained in the source `handle`. Otherwise it silently fails.
+    Splits `handle` into one or multiple handles, as specified by the number
+    of results of this operation. `handle` should be mapped to as many payload
+    ops as there are results. Otherwise, this transform will fail silently.
+    Each result handle is mapped to exactly one payload op. The order
+    of the payload ops is preserved, i.e., the i-th payload op is mapped to the
+    i-th result handle.
+
+    This operation is useful for ensuring a statically known number of
+    operations are tracked by the source `handle` and to extract them into
+    individual handles that can be further manipulated in isolation.
+
+    If `handle` is empty, this transform will succeed and all result handles
+    are empty.
   }];
 
-  let arguments = (ins TransformHandleTypeInterface:$handle,
-                       I64Attr:$num_result_handles);
+  let arguments = (ins TransformHandleTypeInterface:$handle);
   let results = (outs Variadic<TransformHandleTypeInterface>:$results);
 
   let builders = [
     OpBuilder<(ins "Value":$handle, "int64_t":$numResultHandles)>
   ];
 
   let assemblyFormat = [{
-    $handle `in` `[` $num_result_handles `]`
-    attr-dict `:` functional-type(operands, results)
+    $handle attr-dict `:` functional-type(operands, results)
   }];
 }
 
 def ParamConstantOp : Op<Transform_Dialect, "param.constant", [
     MatchOpInterface,
     DeclareOpInterfaceMethods<TransformOpInterface>,
     MemoryEffectsOpInterface,
     ParamProducerTransformOpTrait]> {

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

 }
 
 LogicalResult transform::NamedSequenceOp::verify() {
   // Actual verification happens in a separate function for reusability.
   return verifyNamedSequenceOp(*this).checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//
-// SplitHandlesOp
+// SplitHandleOp
 //===----------------------------------------------------------------------===//
 
-void transform::SplitHandlesOp::build(OpBuilder &builder,
-                                      OperationState &result, Value target,
-                                      int64_t numResultHandles) {
+void transform::SplitHandleOp::build(OpBuilder &builder, OperationState &result,
+                                     Value target, int64_t numResultHandles) {
   result.addOperands(target);
-  result.addAttribute(SplitHandlesOp::getNumResultHandlesAttrName(result.name),
-                      builder.getI64IntegerAttr(numResultHandles));
   auto pdlOpType = pdl::OperationType::get(builder.getContext());
   result.addTypes(SmallVector<pdl::OperationType>(numResultHandles, pdlOpType));
 }
 
 DiagnosedSilenceableFailure
-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) {
-      for (OpResult result : getResults())
-        results.set(result, {});
-      return DiagnosedSilenceableFailure::success();
-    }
+transform::SplitHandleOp::apply(transform::TransformResults &results,
+  int64_t numPayloadOps = state.getPayloadOps(getHandle()).size();
+
+  // Empty handle corner case: all result handles are empty.
+  if (numPayloadOps == 0) {
 
-    // If the input handle was not empty and the number of result handles does
-    // not match, this is a legit silenceable error.
+  // If the input handle was not empty and the number of payload ops does not
+  // match, this is a legit silenceable error.
+  if (numPayloadOps != getNumResults())
     return emitSilenceableError()
-           << getHandle() << " expected to contain " << expectedNumResultHandles
-           << " operation handles but it contains " << numResultHandles
-           << " handles";
-  }
-  // Normal successful case.
+           << getHandle() << " expected to contain " << getNumResults()
+           << " payload ops but it contains " << numPayloadOps
+           << " payload ops";
+
   for (const auto &en : llvm::enumerate(state.getPayloadOps(getHandle())))
     results.set(getResults()[en.index()].cast<OpResult>(), en.value());
+
   return DiagnosedSilenceableFailure::success();
 }
 
-void transform::SplitHandlesOp::getEffects(
+void transform::SplitHandleOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getHandle(), effects);
   producesHandle(getResults(), effects);
   // There are no effects on the Payload IR as this is only a handle
   // manipulation.
 }
 
 //===----------------------------------------------------------------------===//

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

^bb1(%arg1: !pdl.operation):
   %muli = transform.structured.match ops{["arith.muli"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   // expected-error @below {{result number overflow}}
   %bbarg = get_consumers_of_result %muli[1] : (!pdl.operation) -> !pdl.operation
 
 }
 
 // -----
 
-func.func @split_handles(%a: index, %b: index, %c: index) {
+func.func @split_handle(%a: index, %b: index, %c: index) {
   %0 = arith.muli %a, %b : index
   %1 = arith.muli %a, %c : index
   return
 }
 
 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)
+  %h:2 = split_handle %muli : (!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 contains 2 handles}}
-  %h_2:3 = split_handles %muli_2 in [3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
+  // expected-error @below {{expected to contain 3 payload ops but it contains 2 payload ops}}
+  %h_2:3 = split_handle %muli_2 : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // -----
 
-func.func @split_handles(%a: index, %b: index, %c: index) {
+func.func @split_handle(%a: index, %b: index, %c: index) {
   %0 = arith.muli %a, %b : index
   %1 = arith.muli %a, %c : index
   return
 }
 
 transform.sequence failures(suppress) {
 ^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)
+  %h:2 = split_handle %muli : (!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
   // Silenceable failure and all handles are now empty.
-  %h_2:3 = split_handles %muli_2 in [3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
+  %h_2:3 = split_handle %muli_2 : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
   // expected-remark @below {{0}}
   transform.test_print_number_of_associated_payload_ir_ops %h_2#0
 }
 
 // -----
 
 "test.some_op"() : () -> ()
 "other_dialect.other_op"() : () -> ()
pdl.pattern @some : benefit(1) {
     pdl.rewrite %0 with "transform.dialect"
   }
 }
 
 "test.some_op"() : () -> ()
 
 // -----
 
-func.func @split_handles(%a: index, %b: index, %c: index) {
+func.func @split_handle(%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 contains 2 handles}}
-  %h_2:3 = split_handles %muli in [3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
+  // expected-error @below {{expected to contain 3 payload ops but it contains 2 payload ops}}
+  %h_2:3 = split_handle %muli : (!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)
+  %0:3 = split_handle %muli : (!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):