162 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir import ir4from mlir.dialects.transform import interpreter as interp5 6 7def test_in_context(f):8 with ir.Context(), ir.Location.unknown():9 f()10 return f11 12 13print_root_module = """14module attributes {transform.with_named_sequence} {15 transform.named_sequence @__transform_main(%root: !transform.any_op) {16 transform.print %root { name = \"from interpreter\" }: !transform.any_op17 transform.yield18 }19}"""20 21 22@test_in_context23def print_self():24 m = ir.Module.parse(print_root_module.replace("from interpreter", "print_self"))25 interp.apply_named_sequence(m, m.body.operations[0], m)26 27 28# CHECK-LABEL: print_self29# CHECK: transform.named_sequence @__transform_main30# CHECK: transform.print31# CHECK: transform.yield32 33 34@test_in_context35def print_self_via_apply_method():36 m = ir.Module.parse(37 print_root_module.replace("from interpreter", "print_self_via_apply_method")38 )39 m.body.operations[0].apply(m)40 41 42# CHECK-LABEL: print_self_via_apply_method43# CHECK: transform.named_sequence @__transform_main44# CHECK: transform.print45# CHECK: transform.yield46 47 48@test_in_context49def print_other():50 transform = ir.Module.parse(51 print_root_module.replace("from interpreter", "print_other")52 )53 payload = ir.Module.parse("module attributes { this.is.payload } {}")54 interp.apply_named_sequence(payload, transform.body.operations[0], transform)55 56 57# CHECK-LABEL: print_other58# CHECK-NOT: transform59# CHECK: this.is.payload60 61 62@test_in_context63def transform_options():64 options = interp.TransformOptions()65 options.expensive_checks = False66 options.enforce_single_top_level_transform_op = True67 m = ir.Module.parse(68 print_root_module.replace("from interpreter", "transform_options")69 )70 payload = ir.Module.parse("module attributes { this.is.payload } {}")71 interp.apply_named_sequence(payload, m.body.operations[0], m, options)72 73 74# CHECK-LABEL: transform_options75 76 77@test_in_context78def failed():79 payload = ir.Module.parse("module attributes { this.is.payload } {}")80 try:81 interp.apply_named_sequence(payload, payload, payload)82 except ValueError as e:83 assert (84 "must implement TransformOpInterface to be used as transform root" in str(e)85 )86 87 88print_root_via_include_module = """89module @print_root_via_include_module attributes {transform.with_named_sequence} {90 transform.named_sequence private @callee1(%root: !transform.any_op {transform.readonly})91 transform.named_sequence private @callee2(%root: !transform.any_op {transform.readonly})92 transform.named_sequence @__transform_main(%root: !transform.any_op) {93 transform.include @callee2 failures(propagate)94 (%root) : (!transform.any_op) -> ()95 transform.yield96 }97}"""98 99callee2_definition = """100module attributes {transform.with_named_sequence} {101 transform.named_sequence private @callee1(%root: !transform.any_op {transform.readonly})102 transform.named_sequence @callee2(%root: !transform.any_op {transform.readonly}) {103 transform.include @callee1 failures(propagate)104 (%root) : (!transform.any_op) -> ()105 transform.yield106 }107}108"""109 110callee1_definition = """111module attributes {transform.with_named_sequence} {112 transform.named_sequence @callee1(%root: !transform.any_op {transform.readonly}) {113 transform.print %root { name = \"from interpreter\" }: !transform.any_op114 transform.yield115 }116}117"""118 119 120@test_in_context121def include():122 main = ir.Module.parse(print_root_via_include_module)123 callee1 = ir.Module.parse(callee1_definition)124 callee2 = ir.Module.parse(callee2_definition)125 interp.copy_symbols_and_merge_into(main, callee1)126 interp.copy_symbols_and_merge_into(main, callee2)127 128 # CHECK: @print_root_via_include_module129 # CHECK: transform.named_sequence @__transform_main130 # CHECK: transform.include @callee2131 #132 # CHECK: transform.named_sequence @callee1133 # CHECK: transform.print134 #135 # CHECK: transform.named_sequence @callee2136 # CHECK: transform.include @callee1137 interp.apply_named_sequence(main, main.body.operations[0], main)138 139 140@test_in_context141def partial_include():142 main = ir.Module.parse(print_root_via_include_module)143 callee2 = ir.Module.parse(callee2_definition)144 interp.copy_symbols_and_merge_into(main, callee2)145 146 try:147 interp.apply_named_sequence(main, main.body.operations[0], main)148 except ValueError as e:149 assert "Failed to apply" in str(e)150 151 152@test_in_context153def repeated_include():154 main = ir.Module.parse(print_root_via_include_module)155 callee2 = ir.Module.parse(callee2_definition)156 interp.copy_symbols_and_merge_into(main, callee2)157 158 try:159 interp.copy_symbols_and_merge_into(main, callee2)160 except ValueError as e:161 assert "doubly defined symbol @callee2" in str(e)162