[MLIR][Transform] Introduce transform.tune.knob op (#146732)

A new transform op to represent that an attribute is to be chosen from a
set of alternatives and that this choice is made available as a
`!transform.param`. When a `selected` argument is provided, the op's
`apply()` semantics is that of just making this selected attribute
available as the result. When `selected` is not provided, `apply()`
complains that nothing has resolved the non-determinism that the op is
representing.

Author: rolfmorel

mlir/python/CMakeLists.txt

@@ -180,6 +180,15 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   DIALECT_NAME transform
   EXTENSION_NAME transform_debug_extension)
 
+declare_mlir_dialect_extension_python_bindings(
+ADD_TO_PARENT MLIRPythonSources.Dialects
+ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/TransformTuneExtensionOps.td
+  SOURCES
+    dialects/transform/tune.py
+  DIALECT_NAME transform
+  EXTENSION_NAME transform_tune_extension)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"

mlir/python/mlir/dialects/TransformTuneExtensionOps.td

@@ -0,0 +1,19 @@
+//===-- TransformTuneExtensionOps.td - Binding entry point -*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Entry point of the generated Python bindings for the Tune extension of the
+// Transform dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_TRANSFORM_DEBUG_EXTENSION_OPS
+#define PYTHON_BINDINGS_TRANSFORM_DEBUG_EXTENSION_OPS
+
+include "mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td"
+
+#endif // PYTHON_BINDINGS_TRANSFORM_DEBUG_EXTENSION_OPS

mlir/python/mlir/dialects/transform/tune.py

@@ -0,0 +1,82 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import Optional, Sequence
+
+from ...ir import (
+    Type,
+    Attribute,
+    ArrayAttr,
+    StringAttr,
+    F64Type,
+    IntegerType,
+    IntegerAttr,
+    FloatAttr,
+    BoolAttr,
+)
+from .._transform_tune_extension_ops_gen import *
+from .._transform_tune_extension_ops_gen import _Dialect
+
+try:
+    from .._ods_common import _cext as _ods_cext
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+from typing import Union
+
+
+@_ods_cext.register_operation(_Dialect, replace=True)
+class KnobOp(KnobOp):
+    def __init__(
+        self,
+        result: Type,  # !transform.any_param or !transform.param<Type>
+        name: Union[StringAttr, str],
+        options: Union[
+            ArrayAttr, Sequence[Union[Attribute, bool, int, float, str]], Attribute
+        ],
+        *,
+        selected: Optional[Attribute] = None,
+        loc=None,
+        ip=None,
+    ):
+        if isinstance(name, str):
+            name = StringAttr.get(name)
+
+        def map_to_attr(value):
+            if isinstance(value, bool):
+                return BoolAttr.get(value)
+            if isinstance(value, int):
+                return IntegerAttr.get(IntegerType.get_signless(64), value)
+            if isinstance(value, float):
+                return FloatAttr.get(F64Type.get(), value)
+            if isinstance(value, str):
+                return StringAttr.get(value)
+            assert isinstance(value, Attribute)
+            return value
+
+        if isinstance(options, Sequence) and not isinstance(options, ArrayAttr):
+            options = ArrayAttr.get([map_to_attr(opt) for opt in options])
+
+        super().__init__(
+            result,
+            name,
+            options,
+            selected=selected and map_to_attr(selected),
+            loc=loc,
+            ip=ip,
+        )
+
+
+def knob(
+    result: Type,  # !transform.any_param or !transform.param<Type>
+    name: Union[StringAttr, str],
+    options: Union[
+        ArrayAttr, Sequence[Union[Attribute, bool, int, float, str]], Attribute
+    ],
+    *,
+    selected: Optional[Attribute] = None,
+    loc=None,
+    ip=None,
+):
+    return KnobOp(result, name, options, selected=selected, loc=loc, ip=ip)