db7888ca9a
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.
73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
# 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
|
|
@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)
|
|
|
|
# CHECK: transform.debug.emit_param_as_remark %[[HEADS]]
|
|
debug.emit_param_as_remark(heads)
|
|
|
|
# CHECK: transform.tune.knob<"range_as_a_dict"> = 4 : i64 from options = {start = 2 : i64, step = 2 : i64, stop = 16 : i64} -> !transform.any_param
|
|
# NB: Membership of `selected` in non-ArrayAttr `options` is _not_ verified.
|
|
i64 = IntegerType.get_signless(64)
|
|
tune.knob(
|
|
any_param,
|
|
"range_as_a_dict",
|
|
DictAttr.get(
|
|
{
|
|
"start": IntegerAttr.get(i64, 2),
|
|
"stop": IntegerAttr.get(i64, 16),
|
|
"step": IntegerAttr.get(i64, 2),
|
|
}
|
|
),
|
|
selected=4,
|
|
)
|