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

[Rolf Morel]

================
@@ -0,0 +1,56 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import transform
+from mlir.dialects.transform import tune, debug
+
+
+def run(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            sequence = transform.SequenceOp(
+                transform.FailurePropagationMode.Propagate,
+                [],
+                transform.AnyOpType.get(),
+            )
+            with InsertionPoint(sequence.body):
+                f(sequence.bodyTarget)
+                transform.YieldOp()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: TEST: testKnobOp
+ at run
+def testKnobOp(target):
+    any_param = transform.AnyParamType.get()
+
+    # CHECK: %[[HEADS_OR_TAILS:.*]] = transform.tune.knob<"coin"> options = [true, false] -> !transform.any_param
+    heads_or_tails = tune.KnobOp(
+        result=any_param, name=StringAttr.get("coin"), options=[True, False]
+    )
+    # CHECK: transform.tune.knob<"animal"> options = ["cat", "dog", unit] -> !transform.any_param
+    tune.KnobOp(any_param, name="animal", options=["cat", "dog", UnitAttr.get()])
+    # CHECK: transform.tune.knob<"tile_size"> options = [2, 4, 8, 16, 24, 32] -> !transform.any_param
+    tune.KnobOp(any_param, "tile_size", [2, 4, 8, 16, 24, 32])
+    # CHECK: transform.tune.knob<"magic_value"> options = [2.000000e+00, 2.250000e+00, 2.500000e+00, 2.750000e+00, 3.000000e+00] -> !transform.any_param
+    tune.knob(any_param, "magic_value", [2.0, 2.25, 2.5, 2.75, 3.0])
+
+    # CHECK: transform.debug.emit_param_as_remark %[[HEADS_OR_TAILS]]
+    debug.emit_param_as_remark(heads_or_tails)
+
+    # CHECK: %[[HEADS:.*]] = transform.tune.knob<"coin"> = true from options = [true, false] -> !transform.any_param
+    heads = tune.KnobOp(any_param, "coin", options=[True, False], selected=True)
+    # CHECK: transform.tune.knob<"animal"> = "dog" from options = ["cat", "dog", unit] -> !transform.any_param
+    tune.KnobOp(
+        any_param, name="animal", options=["cat", "dog", UnitAttr.get()], selected="dog"
+    )
+    # CHECK: transform.tune.knob<"tile_size"> = 8 : i64 from options = [2, 4, 8, 16, 24, 32] -> !transform.any_param
+    tune.KnobOp(any_param, "tile_size", [2, 4, 8, 16, 24, 32], selected=8)
+    # CHECK: transform.tune.knob<"magic_value"> = 2.500000e+00 : f64 from options = [2.000000e+00, 2.250000e+00, 2.500000e+00, 2.750000e+00, 3.000000e+00] -> !transform.any_param
+    tune.knob(any_param, "magic_value", [2.0, 2.25, 2.5, 2.75, 3.0], selected=2.5)
----------------
rolfmorel wrote:

👍 

Added an example to the Python tests.

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