69 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir import ir4from mlir.dialects import transform, smt5from mlir.dialects.transform import smt as transform_smt6 7 8def run(f):9 print("\nTEST:", f.__name__)10 with ir.Context(), ir.Location.unknown():11 module = ir.Module.create()12 with ir.InsertionPoint(module.body):13 sequence = transform.SequenceOp(14 transform.FailurePropagationMode.Propagate,15 [],16 transform.AnyOpType.get(),17 )18 with ir.InsertionPoint(sequence.body):19 f(sequence.bodyTarget)20 transform.YieldOp()21 print(module)22 return f23 24 25# CHECK-LABEL: TEST: testConstrainParamsOp26@run27def testConstrainParamsOp(target):28 c42_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 42)29 # CHECK: %[[PARAM_AS_PARAM:.*]] = transform.param.constant30 symbolic_value_as_param = transform.ParamConstantOp(31 transform.AnyParamType.get(), c42_attr32 )33 # CHECK: transform.smt.constrain_params(%[[PARAM_AS_PARAM]])34 constrain_params = transform_smt.ConstrainParamsOp(35 [], [symbolic_value_as_param], [smt.IntType.get()]36 )37 # CHECK-NEXT: ^bb{{.*}}(%[[PARAM_AS_SMT_SYMB:.*]]: !smt.int):38 with ir.InsertionPoint(constrain_params.body):39 symbolic_value_as_smt_var = constrain_params.body.arguments[0]40 # CHECK: %[[C0:.*]] = smt.int.constant 041 c0 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0))42 # CHECK: %[[C43:.*]] = smt.int.constant 4343 c43 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 43))44 # CHECK: %[[LB:.*]] = smt.int.cmp le %[[C0]], %[[PARAM_AS_SMT_SYMB]]45 lb = smt.IntCmpOp(smt.IntPredicate.le, c0, symbolic_value_as_smt_var)46 # CHECK: %[[UB:.*]] = smt.int.cmp le %[[PARAM_AS_SMT_SYMB]], %[[C43]]47 ub = smt.IntCmpOp(smt.IntPredicate.le, symbolic_value_as_smt_var, c43)48 # CHECK: %[[BOUNDED:.*]] = smt.and %[[LB]], %[[UB]]49 bounded = smt.AndOp([lb, ub])50 # CHECK: smt.assert %[[BOUNDED:.*]]51 smt.AssertOp(bounded)52 smt.YieldOp([])53 54 # CHECK: transform.smt.constrain_params(%[[PARAM_AS_PARAM]])55 compute_with_params = transform_smt.ConstrainParamsOp(56 [transform.ParamType.get(ir.IntegerType.get_signless(32))],57 [symbolic_value_as_param],58 [smt.IntType.get()],59 )60 # CHECK-NEXT: ^bb{{.*}}(%[[SMT_SYMB:.*]]: !smt.int):61 with ir.InsertionPoint(compute_with_params.body):62 symbolic_value_as_smt_var = compute_with_params.body.arguments[0]63 # CHECK: %[[TWICE:.*]] = smt.int.add %[[SMT_SYMB]], %[[SMT_SYMB]]64 twice_symb = smt.IntAddOp(65 [symbolic_value_as_smt_var, symbolic_value_as_smt_var]66 )67 # CHECK: smt.yield %[[TWICE]]68 smt.YieldOp([twice_symb])69