transform.tune.pick: tune.select but without forgetting the options

Author: rolfmorel

include/TPP/Dialect/Tune/TuneTransformOps.td

@@ -1,6 +1,7 @@
 #ifndef TUNE_TRANSFORM_OPS
 #define TUNE_TRANSFORM_OPS
 
+include "mlir/IR/CommonAttrConstraints.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -17,7 +18,25 @@ def TuneSelectOp : Op<Transform_Dialect, "tune.select", [
 
   let arguments = (ins SymbolRefAttr:$name,
                        ArrayAttr:$options);
-  let results = (outs TransformParamTypeInterface:$selected);
+  let results = (outs TransformParamTypeInterface:$result);
+  let assemblyFormat =
+      "$name `from` $options attr-dict `:` type(results)";
+}
+
+
+def TunePickOp : Op<Transform_Dialect, "tune.pick", [
+  DeclareOpInterfaceMethods<TransformOpInterface>,
+  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+]> {
+  let summary = "Non-deterministically select a value from a set of values";
+  let description = [{
+    TODO
+  }];
+
+  let arguments = (ins SymbolRefAttr:$name,
+                       ArrayAttr:$options,
+                       OptionalAttr<AnyAttr>:$selected);
+  let results = (outs TransformParamTypeInterface:$result);
   let assemblyFormat =
       "$name `from` $options attr-dict `:` type(results)";
 }

lib/TPP/Dialect/Tune/TransformOps/TuneTransformOps.cpp

@@ -29,6 +29,34 @@ transform::TuneSelectOp::apply(transform::TransformRewriter &rewriter,
          << "this op does not resolve non-deterministic choice!";
 }
 
+//===----------------------------------------------------------------------===//
+// TunePickOp
+//===----------------------------------------------------------------------===//
+
+void transform::TunePickOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  producesHandle(getOperation()->getOpResults(), effects);
+  onlyReadsPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+transform::TunePickOp::apply(transform::TransformRewriter &rewriter,
+                             transform::TransformResults &results,
+                             transform::TransformState &state) {
+  if (getSelected()) {
+    results.setParams(getOperation()->getOpResults()[0], *getSelected());
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  if (getOptions().size() == 1) {
+    results.setParams(getOperation()->getOpResults()[0], getOptions()[0]);
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  return emitDefiniteFailure() << "non-deterministic choice is only resolved "
+                                  "through providing a `selected` attr!";
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

python/mlir/dialects/transform/tune.py

@@ -13,14 +13,14 @@
     IntegerType,
     BoolAttr,
 )
-from .._tune_transform_ops_gen import TuneSelectOp
+from .._tune_transform_ops_gen import *
 
 from collections.abc import Sequence
-from typing import Union
+from typing import Union, Optional
 
 
 def select(
-    selected: Type,  # transform.any_param or transform.param<...>
+    result: Type,  # transform.any_param or transform.param<...>
     name: Union[str, Attribute],
     options: Union[ArrayAttr, Sequence[Union[Attribute, str, int, bool]]],
     loc=None,
@@ -44,9 +44,60 @@ def select(
         options = ArrayAttr.get(option_attrs)
 
     return TuneSelectOp(
-        selected=selected,
+        result=result,
+        name=name,
+        options=options,
+        loc=loc,
+        ip=ip,
+    )
+
+
+def pick(
+    result: Type,  # transform.any_param or transform.param<...>
+    name: Union[str, Attribute],
+    options: Union[ArrayAttr, Sequence[Union[Attribute, str, int, bool]]],
+    *,
+    selected: Optional[Union[Attribute, str, int, bool]] = None,
+    loc=None,
+    ip=None,
+) -> TunePickOp:
+    if isinstance(name, str):
+        name = SymbolRefAttr.get([name])
+
+    if not isinstance(options, ArrayAttr):
+        option_attrs = []
+        for option in options:
+            if isinstance(option, str):
+                option_attrs.append(StringAttr.get(option))
+            elif isinstance(option, int):
+                int_type = IntegerType.get_signless(64)
+                option_attrs.append(IntegerAttr.get(int_type, option))
+            elif isinstance(option, bool):
+                option_attrs.append(BoolAttr.get(option))
+            elif isinstance(option, Attribute):
+                option_attrs.append(option)
+            else:
+                assert False
+        options = ArrayAttr.get(option_attrs)
+
+
+    if selected is None:
+        pass
+    elif isinstance(selected, str):
+        selected = StringAttr.get(selected)
+    elif isinstance(selected, int):
+        int_type = IntegerType.get_signless(64)
+        selected = IntegerAttr.get(int_type, selected)
+    elif isinstance(selected, bool):
+        selected = BoolAttr.get(selected)
+    elif not isinstance(selected, Attribute):
+        assert False
+
+    return TunePickOp(
+        result=result,
         name=name,
         options=options,
+        selected=selected,
         loc=loc,
         ip=ip,
     )

python/mlir/tpp/sched/bundles.py

@@ -2,7 +2,7 @@
 
 from mlir import ir
 from mlir.dialects import transform
-from .common import apply_registered_pass, match, select
+from .common import apply_registered_pass, match, select, pick
 from .utils import GpuBackend, PipelineInterrupt
 
 from ..xsmm import utils as xsmm_utils
@@ -46,7 +46,7 @@ def tpp_mapping(
         m_vals, n_vals, k_vals = pack_block_factors
         m = select("m", m_vals if isinstance(m_vals, Sequence) else [m_vals])
         n = select("n", n_vals if isinstance(n_vals, Sequence) else [n_vals])
-        k = select("k", k_vals if isinstance(k_vals, Sequence) else [k_vals])
+        k = pick("k", k_vals if isinstance(k_vals, Sequence) else [k_vals])
         options = {"block-factors": [m, n, k]}
     func = apply_registered_pass(func, "pack-matmul", options=options)
     apply_registered_pass(func, "pack-vnni")

python/mlir/tpp/sched/common.py

@@ -3,13 +3,18 @@
 
 
 # Wrapper to addresss verbosity.
-def apply_registered_pass(*args, **kwargs):
-    return transform.apply_registered_pass(transform.AnyOpType.get(), *args, **kwargs)
+def select(*args, **kwargs):
+    return tune.select(transform.AnyParamType.get(), *args, **kwargs)
 
 
 # Wrapper to addresss verbosity.
-def select(*args, **kwargs):
-    return tune.select(transform.AnyParamType.get(), *args, **kwargs)
+def pick(*args, **kwargs):
+    return tune.pick(transform.AnyParamType.get(), *args, **kwargs)
+
+
+# Wrapper to addresss verbosity.
+def apply_registered_pass(*args, **kwargs):
+    return transform.apply_registered_pass(transform.AnyOpType.get(), *args, **kwargs)
 
 
 # Wrapper to addresss verbosity.

tools/tpp-tune/tpp-tune.py

@@ -49,6 +49,10 @@ def choices_finder(op):
             if op.name in choices:
                 raise RuntimeError(f"options name collision: {op.name} used twice")
             choices[op.name] = tuple(op.options)
+        elif isinstance(op, transform_tune.TunePickOp):
+            if op.name in choices:
+                raise RuntimeError(f"options name collision: {op.name} used twice")
+            choices[op.name] = tuple(op.options)
 
     choices_finder(schedule.operation)
 
@@ -63,6 +67,8 @@ def selected_rewriter(op: Union[ir.OpView, ir.Operation]):
                 )
                 for use in op.result.uses:
                     use.owner.operands[use.operand_number] = param
+        elif isinstance(op, transform_tune.TunePickOp):
+            op.attributes["selected"] = selected[op.name]
 
     selected_rewriter(schedule.operation)