[Mlir-commits] [mlir] [MLIR][Transform] Introduce `transform.tune.knob` op (PR #146732)

[Rolf Morel]

================
@@ -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
+
+
+ at _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])
----------------
rolfmorel wrote:

Hmm. I don't think there's any general facility at the moment* to recursively convert nested Python values to their wrapped-up equivalents. That is, the registered builder for `ArrayAttr` does nothing to the provided (list) argument: https://github.com/llvm/llvm-project/blob/8740ff822d462844506134bb7c425e1778518b95/mlir/python/mlir/ir.py#L165-L167

There are some wrappers for things like `StrArrayAttr` that do convert the elements, but for this op the array can be heterogeneous and hence none of those builders would be suitable. I (and I think @nicolasvasilache as well) would be in favour of more of these general converter utilities to do automagical conversions, also when that doesn't map to a pre-specified nested type (like `StrArrayAttr`) on the C++ side. Preferably these would live in a central place and it would be amazing if these were called as part of the auto-generated `__init__` methods of the bindings. This seems pretty feasible to do.

*: something like the following, which does the recursive thing, but this one also deals with a special attribute: https://github.com/llvm/llvm-project/blob/8740ff822d462844506134bb7c425e1778518b95/mlir/python/mlir/dialects/transform/__init__.py#L247-L266

https://github.com/llvm/llvm-project/pull/146732