d8b84be107
Introduces a Transform-dialect SMT-extension so that we can have an op to express constrains on Transform-dialect params, in particular when these params are knobs -- see transform.tune.knob -- and can hence be seen as symbolic variables. This op allows expressing joint constraints over multiple params/knobs together. While the op's semantics are clearly defined, per SMTLIB, the interpreted semantics -- i.e. the `apply()` method -- for now just defaults to failure. In the future we should support attaching an implementation so that users can Bring Your Own Solver and thereby control performance of interpreting the op. For now the main usage is to walk schedule IR and collect these constraints so that knobs can be rewritten to constants that satisfy the constraints.
51 lines
2.0 KiB
Python
51 lines
2.0 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir import ir
|
|
from mlir.dialects import transform, smt
|
|
from mlir.dialects.transform import smt as transform_smt
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
with ir.Context(), ir.Location.unknown():
|
|
module = ir.Module.create()
|
|
with ir.InsertionPoint(module.body):
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate,
|
|
[],
|
|
transform.AnyOpType.get(),
|
|
)
|
|
with ir.InsertionPoint(sequence.body):
|
|
f(sequence.bodyTarget)
|
|
transform.YieldOp()
|
|
print(module)
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testConstrainParamsOp
|
|
@run
|
|
def testConstrainParamsOp(target):
|
|
dummy_value = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 42)
|
|
# CHECK: %[[PARAM_AS_PARAM:.*]] = transform.param.constant
|
|
symbolic_value = transform.ParamConstantOp(
|
|
transform.AnyParamType.get(), dummy_value
|
|
)
|
|
# CHECK: transform.smt.constrain_params(%[[PARAM_AS_PARAM]])
|
|
constrain_params = transform_smt.ConstrainParamsOp(
|
|
[symbolic_value], [smt.IntType.get()]
|
|
)
|
|
# CHECK-NEXT: ^bb{{.*}}(%[[PARAM_AS_SMT_SYMB:.*]]: !smt.int):
|
|
with ir.InsertionPoint(constrain_params.body):
|
|
# CHECK: %[[C0:.*]] = smt.int.constant 0
|
|
c0 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0))
|
|
# CHECK: %[[C43:.*]] = smt.int.constant 43
|
|
c43 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 43))
|
|
# CHECK: %[[LB:.*]] = smt.int.cmp le %[[C0]], %[[PARAM_AS_SMT_SYMB]]
|
|
lb = smt.IntCmpOp(smt.IntPredicate.le, c0, constrain_params.body.arguments[0])
|
|
# CHECK: %[[UB:.*]] = smt.int.cmp le %[[PARAM_AS_SMT_SYMB]], %[[C43]]
|
|
ub = smt.IntCmpOp(smt.IntPredicate.le, constrain_params.body.arguments[0], c43)
|
|
# CHECK: %[[BOUNDED:.*]] = smt.and %[[LB]], %[[UB]]
|
|
bounded = smt.AndOp([lb, ub])
|
|
# CHECK: smt.assert %[[BOUNDED:.*]]
|
|
smt.AssertOp(bounded)
|