[mlir][Transform] Allow parameter types in merge_handles op

[qedawkins]

Summary

Similar to operation handles, merging parameters can be useful to
avoid repetition of common transformations across a set of parameters.
For example, forming a list of static values for comparison rather than
comparing the parameters one at a time.

[ftynse]

Note that AnyHandleOrParamType also allows for value handles, but the implementation doesn't support them. Please either restrict the type to not allow for value handle or add support for those in the implementation. Looks good otherwise.

[qedawkins]

Enable merging of value handles

[qedawkins]

In D153240#4431954, @ftynse wrote:

Note that AnyHandleOrParamType also allows for value handles, but the implementation doesn't support them. Please either restrict the type to not allow for value handle or add support for those in the implementation. Looks good otherwise.

Thanks, I've been a bit confused by the terminology surrounding handles and parameters. Related, transform.replicate seems to be missing support for value handles as well: https://github.com/llvm/llvm-project/blob/4b6d41cd1d73a72403679ec21732df3610fc0964/mlir/lib/Dialect/Transform/IR/TransformOps.cpp#L1472.

[ftynse]

Thanks!

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

1259

Nit: I don't remember if there is a check somewhere that we have more than one operand so .front() wouldn't crash here, better to double-check and return failure if there no verifier.

[ftynse]

In D153240#4433599, @qedawkins wrote:

In D153240#4431954, @ftynse wrote:

Note that AnyHandleOrParamType also allows for value handles, but the implementation doesn't support them. Please either restrict the type to not allow for value handle or add support for those in the implementation. Looks good otherwise.

Thanks, I've been a bit confused by the terminology surrounding handles and parameters. Related, transform.replicate seems to be missing support for value handles as well: https://github.com/llvm/llvm-project/blob/4b6d41cd1d73a72403679ec21732df3610fc0964/mlir/lib/Dialect/Transform/IR/TransformOps.cpp#L1472.

Value handles are newer than many operations and are not widely used, so there are probably more operations that could be updated but haven't because there was no need.

For terminology, parameters are a kind of "attribute handles" if you wish.

[qedawkins]

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

1259

The SameOperandsAndResultType OpTrait seems to check for at least one operand: https://github.com/llvm/llvm-project/blob/58669354bf1f8eef39979e31915b1e212a3985c9/mlir/lib/IR/Operation.cpp#L1024

(verified by trying it anyway) If wanted I can add a test.

error: unexpected error: 'transform.merge_handles' op expected 1 or more operands, but found 0
  transform.merge_handles : !transform.any_value

[qedawkins]

Value handles are newer than many operations and are not widely used, so there are probably more operations that could be updated but haven't because there was no need.

For terminology, parameters are a kind of "attribute handles" if you wish.

Ah makes sense, thanks for the note.

[ftynse]

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

1259

Thanks for checking! There's no need to test that.

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

let arguments = (ins TransformParamTypeInterface:$param,
                        MatchCmpIPredicateAttr:$predicate);
   let assemblyFormat =
       "$predicate $param `,` $reference attr-dict `:` type($param)";
 }
 
 def MergeHandlesOp : TransformDialectOp<"merge_handles",
     [DeclareOpInterfaceMethods<TransformOpInterface, ["allowsRepeatedHandleOperands"]>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     SameOperandsAndResultType]> {
+     MatchOpInterface, SameOperandsAndResultType]> {
   let summary = "Merges handles into one pointing to the union of payload ops";
   let description = [{
     Creates a new Transform IR handle value that points to the same Payload IR
-    operations as the operand handles. The Payload IR operations are listed
-    in the same order as they are in the operand handles, grouped by operand
-    handle, e.g., all Payload IR operations associated with the first handle
-    come first, then all Payload IR operations associated with the second handle
-    and so on. If `deduplicate` is set, do not add the given Payload IR
-    operation more than once to the final list regardless of it coming from the
+    operations/values/parameters as the operand handles. The Payload IR elements
+    are listed in the same order as they are in the operand handles, grouped by
+    operand handle, e.g., all Payload IR associated with the first handle comes
+    first, then all Payload IR associated with the second handle and so on. If
+    `deduplicate` is set, do not add the given Payload IR operation, value, or
+    parameter more than once to the final list regardless of it coming from the
     same or different handles. Consumes the operands and produces a new handle.
   }];
 
-  let arguments = (ins Variadic<TransformHandleTypeInterface>:$handles,
+  let arguments = (ins Variadic<Transform_AnyHandleOrParamType>:$handles,
                        UnitAttr:$deduplicate);
-  let results = (outs TransformHandleTypeInterface:$result);
+  let results = (outs Transform_AnyHandleOrParamType:$result);
   let assemblyFormat = "(`deduplicate` $deduplicate^)? $handles attr-dict `:` type($result)";
   let hasFolder = 1;
 }
 
 def NamedSequenceOp : TransformDialectOp<"named_sequence",
     [CallableOpInterface,
      FunctionOpInterface,
      IsolatedFromAbove,

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

 //===----------------------------------------------------------------------===//
 // MergeHandlesOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::MergeHandlesOp::apply(transform::TransformRewriter &rewriter,
                                  transform::TransformResults &results,
                                  transform::TransformState &state) {
+  ValueRange handles = getHandles();
+  if (isa<TransformHandleTypeInterface>(handles.front().getType())) {
     SmallVector<Operation *> operations;
-  for (Value operand : getHandles())
+    for (Value operand : handles)
       llvm::append_range(operations, state.getPayloadOps(operand));
     if (!getDeduplicate()) {
       results.set(llvm::cast<OpResult>(getResult()), operations);
       return DiagnosedSilenceableFailure::success();
     }
 
     SetVector<Operation *> uniqued(operations.begin(), operations.end());
     results.set(llvm::cast<OpResult>(getResult()), uniqued.getArrayRef());
     return DiagnosedSilenceableFailure::success();
   }
 
+  if (llvm::isa<TransformParamTypeInterface>(handles.front().getType())) {
+    SmallVector<Attribute> attrs;
+    for (Value attribute : handles)
+      llvm::append_range(attrs, state.getParams(attribute));
+    if (!getDeduplicate()) {
+      results.setParams(cast<OpResult>(getResult()), attrs);
+      return DiagnosedSilenceableFailure::success();
+    }
+
+    SetVector<Attribute> uniqued(attrs.begin(), attrs.end());
+    results.setParams(cast<OpResult>(getResult()), uniqued.getArrayRef());
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  assert(
+      llvm::isa<TransformValueHandleTypeInterface>(handles.front().getType()) &&
+      "expected value handle type");
+  SmallVector<Value> payloadValues;
+  for (Value value : handles)
+    llvm::append_range(payloadValues, state.getPayloadValues(value));
+  if (!getDeduplicate()) {
+    results.setValues(cast<OpResult>(getResult()), payloadValues);
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  SetVector<Value> uniqued(payloadValues.begin(), payloadValues.end());
+  results.setValues(cast<OpResult>(getResult()), uniqued.getArrayRef());
+  return DiagnosedSilenceableFailure::success();
+}
+
 bool transform::MergeHandlesOp::allowsRepeatedHandleOperands() {
   // Handles may be the same if deduplicating is enabled.
   return getDeduplicate();
 }
 
 void transform::MergeHandlesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getHandles(), effects);

mlir/test/Dialect/Linalg/match-ops-interpreter.mlir

 
 // -----
 
 module attributes { transform.with_named_sequence } {
   transform.named_sequence @do_nothing(%arg0: !transform.any_op {transform.readonly}) {
     transform.yield
   }
 
+  transform.named_sequence @print_dimension_size_match(%arg0: !transform.any_op {transform.readonly}) {
+    transform.test_print_remark_at_operand %arg0, "matched sizes" : !transform.any_op
+    transform.yield
+  }
+
   transform.named_sequence @match_dimension_capture(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op {
     // Capture multiple dimension values. Suppress failures so we can print them anyway after the capture.
     %0:9 = transform.match.structured failures(suppress) %arg0 
       : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, 
             !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>) {
     ^bb0(%arg1: !transform.any_op):
       // This also tests the positional specification used by other ops, which may not test it again.
       %1 = transform.match.structured.dim %arg1[all] : (!transform.any_op) -> !transform.param<i64>
transform.named_sequence @match_dimension_capture(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op {
     transform.test_print_param %0#4, "dimensions 0, 2:" at %0#0 : !transform.param<i64>, !transform.any_op
     transform.test_print_param %0#5, "dimensions 0, -1:" at %0#0 : !transform.param<i64>, !transform.any_op
     transform.test_print_param %0#6, "dimensions except -1:" at %0#0 : !transform.param<i64>, !transform.any_op
     transform.test_print_param %0#7, "dimensions except 0, -2:" at %0#0 : !transform.param<i64>, !transform.any_op
     transform.test_print_param %0#8, "dimensions 0, -3:" at %0#0 : !transform.param<i64>, !transform.any_op
     transform.yield %0#0 : !transform.any_op
   }
 
+  transform.named_sequence @match_dimension_sizes(%arg0: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+    %0 = transform.match.structured failures(propagate) %arg0 : (!transform.any_op) -> !transform.any_op {
+    ^bb0(%arg1: !transform.any_op):
+      %1 = transform.match.structured.dim %arg1[all] : (!transform.any_op) -> !transform.param<i64>
+      %c2 = transform.param.constant 2 : i64 -> !transform.param<i64>
+      %c3 = transform.param.constant 3 : i64 -> !transform.param<i64>
+      %c4 = transform.param.constant 4 : i64 -> !transform.param<i64>
+      %2 = transform.merge_handles %c2, %c3, %c4 : !transform.param<i64>
+      transform.match.param.cmpi eq %1, %2 : !transform.param<i64>
+
+      transform.match.structured.yield %arg1 : !transform.any_op
+    }
+    transform.yield %0 : !transform.any_op
+  }
+
   transform.sequence failures(propagate) attributes { transform.target_tag = "transform" } {
   ^bb0(%arg0: !transform.any_op):
-    transform.foreach_match in %arg0 @match_dimension_capture -> @do_nothing : (!transform.any_op) -> !transform.any_op
+    %0 = transform.foreach_match in %arg0 @match_dimension_capture -> @do_nothing : (!transform.any_op) -> !transform.any_op
+    %1 = transform.foreach_match in %0 @match_dimension_sizes -> @print_dimension_size_match : (!transform.any_op) -> !transform.any_op
   }
 
   func.func @payload(%lhs: tensor<2x4xf32>, %rhs: tensor<4x3xf32>, %out: tensor<2x3xf32>) attributes { transform.target_tag = "start_here" } {
     // The last does not emit anything because it fails to match 
     // due to 0 and -3 being the same dimension in the 3D case.
     // expected-remark @below {{dimensions all: 2 : i64, 3 : i64, 4 : i64}}
     // expected-remark @below {{dimension 0: 2 : i64}}
     // expected-remark @below {{dimension -1: 4 : i64}}
     // expected-remark @below {{dimensions 0, 2: 2 : i64, 4 : i64}}
     // expected-remark @below {{dimensions 0, -1: 2 : i64, 4 : i64}}
     // expected-remark @below {{dimensions except -1: 2 : i64, 3 : i64}}
     // expected-remark @below {{dimensions except 0, -2: 4 : i64}}
     // expected-remark @below {{dimensions 0, -3:}}
+    // expected-remark @below {{matched sizes}}
     linalg.generic {
       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
       iterator_types = ["parallel", "parallel", "reduction"]
     } ins(%lhs, %rhs: tensor<2x4xf32>, tensor<4x3xf32>) outs(%out: tensor<2x3xf32>) {
     ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
       %0 = arith.mulf %arg0, %arg1 : f32
       %1 = arith.addf %0, %arg2 : f32
       linalg.yield %1 : f32

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

^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["transform.test_dummy_payload_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-remark @below {{2 iterations}}
   transform.test_tracked_rewrite %0 : (!transform.any_op) -> ()
   // One replacement op (test.drop_mapping) is dropped from the mapping.
   // expected-remark @below {{2}}
   test_print_number_of_associated_payload_ir_ops %0 : !transform.any_op
 }
 
+// -----
+
+// Parameter deduplication happens by value
+
+module {
+
+  transform.sequence failures(propagate) {
+  ^bb0(%0: !transform.any_op):
+    %1 = transform.param.constant 1 -> !transform.param<i64>
+    %2 = transform.param.constant 1 -> !transform.param<i64>
+    %3 = transform.param.constant 2 -> !transform.param<i64>
+    %4 = transform.merge_handles %1, %2 { deduplicate } : !transform.param<i64>
+    // expected-remark @below {{1}}
+    test_print_number_of_associated_payload_ir_params %4 : !transform.param<i64>
+
+    %5 = transform.merge_handles %1, %1 { deduplicate } : !transform.param<i64>
+    // expected-remark @below {{1}}
+    test_print_number_of_associated_payload_ir_params %5 : !transform.param<i64>
+
+    %6 = transform.merge_handles %1, %3 { deduplicate } : !transform.param<i64>
+    // expected-remark @below {{2}}
+    test_print_number_of_associated_payload_ir_params %6 : !transform.param<i64>
+
+    %7 = transform.merge_handles %1, %1, %2, %3 : !transform.param<i64>
+    // expected-remark @below {{4}}
+    test_print_number_of_associated_payload_ir_params %7 : !transform.param<i64>
+  }
+}
 
 // -----
 
+%0:3 = "test.get_two_results"() : () -> (i32, i32, f32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.get_two_results"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %2 = test_produce_value_handle_to_result %1, 0 : (!transform.any_op) -> !transform.any_value
+  %3 = test_produce_value_handle_to_result %1, 1 : (!transform.any_op) -> !transform.any_value
+
+  %4 = transform.merge_handles %2, %2 { deduplicate } : !transform.any_value
+  // expected-remark @below {{1}}
+  test_print_number_of_associated_payload_ir_values %4 : !transform.any_value
+
+  %5 = transform.merge_handles %2, %3 { deduplicate } : !transform.any_value
+  // expected-remark @below {{2}}
+  test_print_number_of_associated_payload_ir_values %5 : !transform.any_value
+
+  %6 = test_produce_value_handle_to_result %1, 0 : (!transform.any_op) -> !transform.any_value
+  %7 = transform.merge_handles %2, %6 { deduplicate } : !transform.any_value
+  // expected-remark @below {{1}}
+  test_print_number_of_associated_payload_ir_values %6 : !transform.any_value
+
+  %8 = transform.merge_handles %2, %2, %3, %4 : !transform.any_value
+  // expected-remark @below {{4}}
+  test_print_number_of_associated_payload_ir_values %8 : !transform.any_value
+}
+// -----
+
 // CHECK-LABEL: func @test_annotation()
 //  CHECK-NEXT:   "test.annotate_me"()
 //  CHECK-SAME:                        broadcast_attr = 2 : i64
 //  CHECK-SAME:                        new_attr = 1 : i32
 //  CHECK-SAME:                        unit_attr
 //  CHECK-NEXT:   "test.annotate_me"()
 //  CHECK-SAME:                        broadcast_attr = 2 : i64
 //  CHECK-SAME:                        existing_attr = "test"

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

 }
 
 void mlir::test::TestPrintNumberOfAssociatedPayloadIROps::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getHandle(), effects);
 }
 
 DiagnosedSilenceableFailure
+mlir::test::TestPrintNumberOfAssociatedPayloadIRValues::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
+  if (!getValueHandle())
+    emitRemark() << 0;
+  emitRemark() << llvm::range_size(state.getPayloadValues(getValueHandle()));
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestPrintNumberOfAssociatedPayloadIRValues::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getValueHandle(), effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestPrintNumberOfAssociatedPayloadIRParams::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
+  if (!getParam())
+    emitRemark() << 0;
+  emitRemark() << llvm::range_size(state.getParams(getParam()));
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestPrintNumberOfAssociatedPayloadIRParams::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getParam(), effects);
+}
+
+DiagnosedSilenceableFailure
 mlir::test::TestCopyPayloadOp::apply(transform::TransformRewriter &rewriter,
                                      transform::TransformResults &results,
                                      transform::TransformState &state) {
   results.set(llvm::cast<OpResult>(getCopy()),
               state.getPayloadOps(getHandle()));
   return DiagnosedSilenceableFailure::success();
 }
 

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

def TestPrintNumberOfAssociatedPayloadIROps
   : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_ops",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$handle);
   let assemblyFormat = "$handle attr-dict `:` type($handle)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestPrintNumberOfAssociatedPayloadIRValues
+  : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_values",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformValueHandleTypeInterface:$value_handle);
+  let assemblyFormat = "$value_handle attr-dict `:` type($value_handle)";
+  let cppNamespace = "::mlir::test";
+}
+
+def TestPrintNumberOfAssociatedPayloadIRParams
+  : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_params",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformParamTypeInterface:$param);
+  let assemblyFormat = "$param attr-dict `:` type($param)";
+  let cppNamespace = "::mlir::test";
+}
+
 def TestCopyPayloadOp
   : Op<Transform_Dialect, "test_copy_payload",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$handle);
   let results = (outs TransformHandleTypeInterface:$copy);
   let cppNamespace = "::mlir::test";
   let assemblyFormat =