[MLIR][Transform] apply_registered_pass op's options as a dict (#143159)

Improve ApplyRegisteredPassOp's support for taking options by taking
them as a dict (vs a list of string-valued key-value pairs).

Values of options are provided as either static attributes or as params
(which pass in attributes at interpreter runtime). In either case, the
keys and value attributes are converted to strings and a single
options-string, in the format used on the commandline, is constructed to
pass to the `addToPipeline`-pass API.

Author: rolfmorel

mlir/include/mlir/Dialect/Transform/IR/CMakeLists.txt

@@ -20,6 +20,10 @@ mlir_tablegen(TransformDialectEnums.h.inc -gen-enum-decls)
 mlir_tablegen(TransformDialectEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRTransformDialectEnumIncGen)
 add_dependencies(mlir-headers MLIRTransformDialectEnumIncGen)
+mlir_tablegen(TransformAttrs.h.inc -gen-attrdef-decls)
+mlir_tablegen(TransformAttrs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRTransformDialectAttributesIncGen)
+add_dependencies(mlir-headers MLIRTransformDialectAttributesIncGen)
 
 add_mlir_dialect(TransformOps transform)
 add_mlir_doc(TransformOps TransformOps Dialects/ -gen-op-doc -dialect=transform)

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

@@ -17,4 +17,7 @@
 
 #include "mlir/Dialect/Transform/IR/TransformDialectEnums.h.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformAttrs.h.inc"
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMATTRS_H

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

@@ -10,6 +10,14 @@
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMATTRS
 
 include "mlir/IR/EnumAttr.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+
+class Transform_Attr<string name, string attrMnemonic,
+                     list<Trait> traits = [],
+                     string baseCppClass = "::mlir::Attribute">
+    : AttrDef<Transform_Dialect, name, traits, baseCppClass> {
+  let mnemonic = attrMnemonic;
+}
 
 def PropagateFailuresCase : I32EnumAttrCase<"Propagate", 1, "propagate">;
 def SuppressFailuresCase : I32EnumAttrCase<"Suppress", 2, "suppress">;
@@ -33,4 +41,15 @@ def MatchCmpIPredicateAttr : I32EnumAttr<
   let cppNamespace = "::mlir::transform";
 }
 
+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` `=` $index `>`";
+}
+
 #endif  // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMATTRS

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

@@ -19,6 +19,7 @@ def Transform_Dialect : Dialect {
   let cppNamespace = "::mlir::transform";
 
   let hasOperationAttrVerify = 1;
+  let useDefaultAttributePrinterParser = 1;
   let extraClassDeclaration = [{
     /// Symbol name for the default entry point "named sequence".
     constexpr const static ::llvm::StringLiteral

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

@@ -405,10 +405,23 @@ def ApplyRegisteredPassOp : TransformDialectOp<"apply_registered_pass",
   let description = [{
     This transform applies the specified pass or pass pipeline to the targeted
     ops. The name of the pass/pipeline is specified as a string attribute, as
-    set during pass/pipeline registration. Optionally, pass options may be
-    specified as (space-separated) string attributes with the option to pass
-    these attributes via params. The pass options syntax is identical to the one
-    used with "mlir-opt".
+    set during pass/pipeline registration.
+
+    Optionally, pass options may be specified via a DictionaryAttr. This
+    dictionary is converted to a string -- formatted `key=value ...` -- which
+    is expected to be in the exact format used by the pass on the commandline.
+    Values are either attributes or (SSA-values of) Transform Dialect params.
+    For example:
+
+    ```mlir
+    transform.apply_registered_pass "canonicalize"
+        with options = { "top-down" = false,
+                         "max-iterations" = %max_iter,
+                         "test-convergence" = true,
+                         "max-num-rewrites" =  %max_rewrites }
+        to %module
+    : (!transform.any_param, !transform.any_param, !transform.any_op) -> !transform.any_op
+    ```
 
     This op first looks for a pass pipeline with the specified name. If no such
     pipeline exists, it looks for a pass with the specified name. If no such
@@ -422,7 +435,7 @@ def ApplyRegisteredPassOp : TransformDialectOp<"apply_registered_pass",
   }];
 
   let arguments = (ins StrAttr:$pass_name,
-                       DefaultValuedAttr<ArrayAttr, "{}">:$options,
+                       DefaultValuedAttr<DictionaryAttr, "{}">:$options,
                        Variadic<TransformParamTypeInterface>:$dynamic_options,
                        TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$result);

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

@@ -8,17 +8,22 @@
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Analysis/CallGraph.h"
+#include "mlir/Dialect/Transform/IR/TransformAttrs.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/Dialect/Transform/IR/Utils.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/SCCIterator.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformAttrs.cpp.inc"
+
 #ifndef NDEBUG
 void transform::detail::checkImplementsTransformOpInterface(
     StringRef name, MLIRContext *context) {
@@ -66,6 +71,10 @@ void transform::TransformDialect::initialize() {
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
       >();
   initializeTypes();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/Transform/IR/TransformAttrs.cpp.inc"
+      >();
   initializeLibraryModule();
 }