Address @adam-smnk's review

Author: rolfmorel

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

@@ -41,17 +41,15 @@ def MatchCmpIPredicateAttr : I32EnumAttr<
   let cppNamespace = "::mlir::transform";
 }
 
-def ParamOperandIndexAttr : Transform_Attr<"ParamOperandIndex",
-                                           "param_operand_index" > {
-  let mnemonic = "param_operand_index";
+def ParamOperandAttr : Transform_Attr<"ParamOperand", "param_operand"> {
   let description = [{
     Used to refer to a specific param-operand (via its index) from within an
     attribute on a transform operation.
   }];
   let parameters = (ins
     "IntegerAttr":$index
   );
-  let assemblyFormat = "`<` $index `>`";
+  let assemblyFormat = "`<` `index` `=` $index `>`";
 }
 
 #endif  // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMATTRS

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

@@ -800,8 +800,8 @@ transform::ApplyRegisteredPassOp::apply(transform::TransformRewriter &rewriter,
     optionsStream << "="; // And the key-value separator.
 
     Attribute valueAttrToAppend;
-    if (auto paramOperandIndex = dyn_cast<transform::ParamOperandIndexAttr>(
-            namedAttribute.getValue())) {
+    if (auto paramOperandIndex =
+            dyn_cast<transform::ParamOperandAttr>(namedAttribute.getValue())) {
       // The corresponding value attribute is passed in via a param.
       // Obtain the param-operand via its specified index.
       size_t dynamicOptionIdx = paramOperandIndex.getIndex().getInt();
@@ -906,16 +906,20 @@ static ParseResult parseApplyRegisteredPassOptions(
         return parser.emitError(parser.getCurrentLocation())
                << "expected a valid attribute or operand as value associated "
                << "to key '" << key << "'";
+      // To make use of the operand, we need to store it in the options dict.
+      // As SSA-values cannot occur in attributes, what we do instead is store
+      // an attribute in its place that contains the index of the param-operand,
+      // so that an attr-value associated to the param can be resolved later on.
       dynamicOptions.push_back(operand);
       auto wrappedIndex = IntegerAttr::get(
           IntegerType::get(parser.getContext(), 64), dynamicOptionsIdx++);
-      valueAttr = transform::ParamOperandIndexAttr::get(parser.getContext(),
-                                                        wrappedIndex);
+      valueAttr =
+          transform::ParamOperandAttr::get(parser.getContext(), wrappedIndex);
     } else if (failed(parsedValueAttr.value())) {
       return failure(); // NB: Attempted parse should have output error message.
-    } else if (isa<transform::ParamOperandIndexAttr>(valueAttr)) {
+    } else if (isa<transform::ParamOperandAttr>(valueAttr)) {
       return parser.emitError(parser.getCurrentLocation())
-             << "the param_operand_index attribute is a marker reserved for "
+             << "the param_operand attribute is a marker reserved for "
              << "indicating a value will be passed via params and is only used "
              << "in the generic print format";
     }
@@ -951,7 +955,8 @@ static void printApplyRegisteredPassOptions(OpAsmPrinter &printer,
   llvm::interleaveComma(options, printer, [&](NamedAttribute namedAttribute) {
     printer << namedAttribute.getName() << " = ";
     Attribute value = namedAttribute.getValue();
-    if (auto indexAttr = dyn_cast<transform::ParamOperandIndexAttr>(value)) {
+    if (auto indexAttr = dyn_cast<transform::ParamOperandAttr>(value)) {
+      // Resolve index of param-operand to its actual SSA-value and print that.
       printer.printOperand(dynamicOptions[indexAttr.getIndex().getInt()]);
     } else {
       printer.printAttribute(value);
@@ -966,9 +971,9 @@ LogicalResult transform::ApplyRegisteredPassOp::verify() {
 
   auto dynamicOptions = SmallVector<Value>(getDynamicOptions());
   for (NamedAttribute namedAttr : getOptions())
-    if (auto paramOperandIndex =
-            dyn_cast<transform::ParamOperandIndexAttr>(namedAttr.getValue())) {
-      size_t dynamicOptionIdx = paramOperandIndex.getIndex().getInt();
+    if (auto paramOperand =
+            dyn_cast<transform::ParamOperandAttr>(namedAttr.getValue())) {
+      size_t dynamicOptionIdx = paramOperand.getIndex().getInt();
       if (dynamicOptionIdx < 0 || dynamicOptionIdx >= dynamicOptions.size())
         return emitOpError()
                << "dynamic option index " << dynamicOptionIdx
@@ -983,7 +988,7 @@ LogicalResult transform::ApplyRegisteredPassOp::verify() {
   for (Value dynamicOption : dynamicOptions)
     if (dynamicOption)
       return emitOpError() << "a param operand does not have a corresponding "
-                           << "param_operand_index attr in the options dict";
+                           << "param_operand attr in the options dict";
 
   return success();
 }

mlir/python/mlir/dialects/transform/__init__.py

@@ -21,9 +21,9 @@
 from typing import Dict, Optional, Sequence, Union, NewType
 
 
-@register_attribute_builder("ParamOperandIndexAttr")
-def _paramOperandIndexAttr(x: int, context) -> Attribute:
-    return Attribute.parse(f"#transform.param_operand_index<{x}>", context=context)
+@register_attribute_builder("ParamOperandAttr")
+def _paramOperandAttr(x: int, context) -> Attribute:
+    return Attribute.parse(f"#transform.param_operand<index={x}>", context=context)
 
 
 @_ods_cext.register_operation(_Dialect, replace=True)
@@ -239,7 +239,7 @@ def __init__(
         options_dict = {}
         dynamic_options = []
 
-        ParamOperandIndexAttr = AttrBuilder.get("ParamOperandIndexAttr")
+        ParamOperandAttr = AttrBuilder.get("ParamOperandAttr")
         context = (loc and loc.context) or Context.current
 
         cur_param_operand_idx = 0
@@ -249,7 +249,7 @@ def __init__(
 
             if isinstance(value, (Value, Operation, OpView)):
                 dynamic_options.append(_get_op_result_or_value(value))
-                options_dict[key] = ParamOperandIndexAttr(
+                options_dict[key] = ParamOperandAttr(
                     cur_param_operand_idx, context
                 )
                 cur_param_operand_idx += 1

mlir/test/Dialect/Transform/test-pass-application.mlir

@@ -204,9 +204,9 @@ func.func @invalid_options_due_to_reserved_attr() {
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg1: !transform.any_op) {
     %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-    // expected-error @+2 {{the param_operand_index attribute is a marker reserved for indicating a value will be passed via params and is only used in the generic print format}}
+    // expected-error @+2 {{the param_operand attribute is a marker reserved for indicating a value will be passed via params and is only used in the generic print format}}
     %2 = transform.apply_registered_pass "canonicalize"
-        with options = { "top-down" = #transform.param_operand_index<0> } to %1 : (!transform.any_op) -> !transform.any_op
+        with options = { "top-down" = #transform.param_operand<index=0> } to %1 : (!transform.any_op) -> !transform.any_op
     transform.yield
   }
 }
@@ -306,7 +306,7 @@ module attributes {transform.with_named_sequence} {
 // Check that the following cases are caugh in the generic format. //
 /////////////////////////////////////////////////////////////////////
 
-// Invalid due to param_operand_index occurences in options dict not being
+// Invalid due to param_operand occurences in options dict not being
 // one-to-one with the dynamic options provided as params:
 //   param_operand_index out of bounds w.r.t. the number of options provided via params.
 
@@ -317,7 +317,7 @@ module attributes {transform.with_named_sequence} {
     %1 = "transform.param.constant"() <{value = 10 : i64}> : () -> !transform.any_param
     // expected-error @below {{dynamic option index 1 is out of bounds for the number of dynamic options: 1}}
     %2 = "transform.apply_registered_pass"(%1, %0) <{
-      options = {"max-iterations" = #transform.param_operand_index<1 : i64>,
+      options = {"max-iterations" = #transform.param_operand<index=1 : i64>,
                  "test-convergence" = true,
                  "top-down" = false},
       pass_name = "canonicalize"}>
@@ -328,10 +328,10 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// Invalid due to param_operand_index occurences in options dict not being
+// Invalid due to param_operand occurences in options dict not being
 // one-to-one with the dynamic options provided as params:
 //   the first option-param is referred to twice and the second one not at all.
-// (The pretty-printed format supports this by passing in the same param twice.)
+// (In the pretty-printed format, if you want to refer to a param SSA-value twice, it counts as two param arguments.)
 
 "builtin.module"() ({
   "transform.named_sequence"() <{function_type = (!transform.any_op) -> (), sym_name = "__transform_main"}> ({
@@ -341,8 +341,8 @@ module attributes {transform.with_named_sequence} {
     %2 = "transform.param.constant"() <{value = 1 : i64}> : () -> !transform.any_param
     // expected-error @below {{dynamic option index 0 is already used in options}}
     %3 = "transform.apply_registered_pass"(%1, %2, %0) <{
-      options = {"max-iterations" = #transform.param_operand_index<0 : i64>,
-                 "max-num-rewrites" = #transform.param_operand_index<0 : i64>,
+      options = {"max-iterations" = #transform.param_operand<index=0 : i64>,
+                 "max-num-rewrites" = #transform.param_operand<index=0 : i64>,
                  "test-convergence" = true,
                  "top-down" = false},
       pass_name = "canonicalize"}>
@@ -353,7 +353,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// Invalid due to param_operand_index occurences in options dict not being
+// Invalid due to param_operand occurences in options dict not being
 // one-to-one with the dynamic options provided as params:
 //   two option-params are provide though only the first one is referred to from the options-dict.
 
@@ -363,9 +363,9 @@ module attributes {transform.with_named_sequence} {
     %0 = "transform.structured.match"(%arg0) <{ops = ["func.func"]}> : (!transform.any_op) -> !transform.any_op
     %1 = "transform.param.constant"() <{value = 10 : i64}> : () -> !transform.any_param
     %2 = "transform.param.constant"() <{value = 1 : i64}> : () -> !transform.any_param
-    // expected-error @below {{a param operand does not have a corresponding param_operand_index attr in the options dict}}
+    // expected-error @below {{a param operand does not have a corresponding param_operand attr in the options dict}}
     %3 = "transform.apply_registered_pass"(%1, %2, %0) <{
-      options = {"max-iterations" = #transform.param_operand_index<0 : i64>,
+      options = {"max-iterations" = #transform.param_operand<index=0 : i64>,
                  "test-convergence" = true,
                  "top-down" = false},
       pass_name = "canonicalize"}>