150 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir import ir4from mlir.dialects import transform5from mlir.dialects.transform import tune, debug6 7 8def run(f):9 print("\n// TEST:", 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: testKnobOp26@run27def testKnobOp(target):28 any_param = transform.AnyParamType.get()29 30 # CHECK: %[[HEADS_OR_TAILS:.*]] = transform.tune.knob<"coin"> options = [true, false] -> !transform.any_param31 heads_or_tails = tune.KnobOp(32 result=any_param, name=ir.StringAttr.get("coin"), options=[True, False]33 )34 # CHECK: transform.tune.knob<"animal"> options = ["cat", "dog", unit] -> !transform.any_param35 tune.KnobOp(any_param, name="animal", options=["cat", "dog", ir.UnitAttr.get()])36 # CHECK: transform.tune.knob<"tile_size"> options = [2, 4, 8, 16, 24, 32] -> !transform.any_param37 tune.KnobOp(any_param, "tile_size", [2, 4, 8, 16, 24, 32])38 # CHECK: transform.tune.knob<"magic_value"> options = [2.000000e+00, 2.250000e+00, 2.500000e+00, 2.750000e+00, 3.000000e+00] -> !transform.any_param39 tune.knob(any_param, "magic_value", [2.0, 2.25, 2.5, 2.75, 3.0])40 41 # CHECK: transform.debug.emit_param_as_remark %[[HEADS_OR_TAILS]]42 debug.emit_param_as_remark(heads_or_tails)43 44 # CHECK: %[[HEADS:.*]] = transform.tune.knob<"coin"> = true from options = [true, false] -> !transform.any_param45 heads = tune.KnobOp(any_param, "coin", options=[True, False], selected=True)46 # CHECK: transform.tune.knob<"animal"> = "dog" from options = ["cat", "dog", unit] -> !transform.any_param47 tune.KnobOp(48 any_param,49 name="animal",50 options=["cat", "dog", ir.UnitAttr.get()],51 selected="dog",52 )53 # CHECK: transform.tune.knob<"tile_size"> = 8 : i64 from options = [2, 4, 8, 16, 24, 32] -> !transform.any_param54 tune.KnobOp(any_param, "tile_size", [2, 4, 8, 16, 24, 32], selected=8)55 # 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_param56 tune.knob(any_param, "magic_value", [2.0, 2.25, 2.5, 2.75, 3.0], selected=2.5)57 58 # CHECK: transform.debug.emit_param_as_remark %[[HEADS]]59 debug.emit_param_as_remark(heads)60 61 # CHECK: transform.tune.knob<"range_as_a_dict"> = 4 : i64 from options = {start = 2 : i64, step = 2 : i64, stop = 16 : i64} -> !transform.any_param62 # NB: Membership of `selected` in non-ArrayAttr `options` is _not_ verified.63 i64 = ir.IntegerType.get_signless(64)64 tune.knob(65 any_param,66 "range_as_a_dict",67 ir.DictAttr.get(68 {69 "start": ir.IntegerAttr.get(i64, 2),70 "stop": ir.IntegerAttr.get(i64, 16),71 "step": ir.IntegerAttr.get(i64, 2),72 }73 ),74 selected=4,75 )76 77 78# CHECK-LABEL: TEST: testAlternativesOp79@run80def testAlternativesOp(target):81 any_param = transform.AnyParamType.get()82 83 # CHECK: %[[LEFT_OR_RIGHT_OUTCOME:.*]] = transform.tune.alternatives<"left_or_right"> -> !transform.any_param {84 left_or_right = tune.AlternativesOp(85 [transform.AnyParamType.get()], "left_or_right", 286 )87 idx_for_left, idx_for_right = 0, 188 with ir.InsertionPoint(left_or_right.alternatives[idx_for_left].blocks[0]):89 # CHECK: %[[C0:.*]] = transform.param.constant 090 i32_0 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)91 c0 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_0)92 # CHECK: transform.yield %[[C0]]93 transform.yield_(c0)94 # CHECK-NEXT: }, {95 with ir.InsertionPoint(left_or_right.alternatives[idx_for_right].blocks[0]):96 # CHECK: %[[C1:.*]] = transform.param.constant 197 i32_1 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1)98 c1 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_1)99 # CHECK: transform.yield %[[C1]]100 transform.yield_(c1)101 # CHECK-NEXT: }102 outcome_of_left_or_right_decision = left_or_right.results[0]103 104 # CHECK: transform.tune.alternatives<"fork_in_the_road"> selected_region = 0 -> !transform.any_param {105 fork_in_the_road = tune.AlternativesOp(106 [transform.AnyParamType.get()], "fork_in_the_road", 2, selected_region=0107 )108 with ir.InsertionPoint(fork_in_the_road.alternatives[idx_for_left].blocks[0]):109 # CHECK: %[[C0:.*]] = transform.param.constant 0110 i32_0 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)111 c0 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_0)112 # CHECK: transform.yield %[[C0]]113 transform.yield_(c0)114 # CHECK-NEXT: }, {115 with ir.InsertionPoint(fork_in_the_road.alternatives[idx_for_right].blocks[0]):116 # CHECK: %[[C1:.*]] = transform.param.constant 1117 i32_1 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1)118 c1 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_1)119 # CHECK: transform.yield %[[C1]]120 transform.yield_(c1)121 # CHECK-NEXT: }122 123 # CHECK: transform.tune.alternatives<"left_or_right_as_before"> selected_region = %[[LEFT_OR_RIGHT_OUTCOME]] : !transform.any_param {124 left_or_right_as_before = tune.AlternativesOp(125 [],126 "left_or_right_as_before",127 2,128 selected_region=outcome_of_left_or_right_decision,129 )130 with ir.InsertionPoint(131 left_or_right_as_before.alternatives[idx_for_left].blocks[0]132 ):133 # CHECK: transform.param.constant 1337134 i32_1337 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1337)135 c1337 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_1337)136 # CHECK: transform.debug.emit_param_as_remark137 debug.emit_param_as_remark(c1337)138 transform.yield_([])139 # CHECK-NEXT: }, {140 with ir.InsertionPoint(141 left_or_right_as_before.alternatives[idx_for_right].blocks[0]142 ):143 # CHECK: transform.param.constant 42144 i32_42 = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 42)145 c42 = transform.ParamConstantOp(transform.AnyParamType.get(), i32_42)146 # CHECK: transform.debug.emit_param_as_remark147 debug.emit_param_as_remark(c42)148 transform.yield_([])149 # CHECK-NEXT: }150