246 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4from mlir.ir import *5 6 7def run(f):8 print("\nTEST:", f.__name__)9 f()10 gc.collect()11 assert Context._get_live_count() == 012 13 14# CHECK-LABEL: TEST: test_insert_at_block_end15def test_insert_at_block_end():16 ctx = Context()17 ctx.allow_unregistered_dialects = True18 with Location.unknown(ctx):19 module = Module.parse(20 r"""21 func.func @foo() -> () {22 "custom.op1"() : () -> ()23 }24 """25 )26 entry_block = module.body.operations[0].regions[0].blocks[0]27 ip = InsertionPoint(entry_block)28 assert ip.block == entry_block29 assert ip.ref_operation is None30 ip.insert(Operation.create("custom.op2"))31 # CHECK: "custom.op1"32 # CHECK: "custom.op2"33 module.operation.print()34 35 36run(test_insert_at_block_end)37 38 39# CHECK-LABEL: TEST: test_insert_before_operation40def test_insert_before_operation():41 ctx = Context()42 ctx.allow_unregistered_dialects = True43 with Location.unknown(ctx):44 module = Module.parse(45 r"""46 func.func @foo() -> () {47 "custom.op1"() : () -> ()48 "custom.op2"() : () -> ()49 }50 """51 )52 entry_block = module.body.operations[0].regions[0].blocks[0]53 ip = InsertionPoint(entry_block.operations[1])54 assert ip.block == entry_block55 assert ip.ref_operation == entry_block.operations[1]56 ip.insert(Operation.create("custom.op3"))57 # CHECK: "custom.op1"58 # CHECK: "custom.op3"59 # CHECK: "custom.op2"60 module.operation.print()61 62 63run(test_insert_before_operation)64 65 66# CHECK-LABEL: TEST: test_insert_after_operation67def test_insert_after_operation():68 ctx = Context()69 ctx.allow_unregistered_dialects = True70 with Location.unknown(ctx):71 module = Module.parse(72 r"""73 func.func @foo() -> () {74 "custom.op1"() : () -> ()75 "custom.op2"() : () -> ()76 }77 """78 )79 entry_block = module.body.operations[0].regions[0].blocks[0]80 custom_op1 = entry_block.operations[0]81 custom_op2 = entry_block.operations[1]82 InsertionPoint.after(custom_op1).insert(Operation.create("custom.op3"))83 InsertionPoint.after(custom_op2).insert(Operation.create("custom.op4"))84 # CHECK: "custom.op1"85 # CHECK: "custom.op3"86 # CHECK: "custom.op2"87 # CHECK: "custom.op4"88 module.operation.print()89 90 91run(test_insert_after_operation)92 93 94# CHECK-LABEL: TEST: test_insert_at_block_begin95def test_insert_at_block_begin():96 ctx = Context()97 ctx.allow_unregistered_dialects = True98 with Location.unknown(ctx):99 module = Module.parse(100 r"""101 func.func @foo() -> () {102 "custom.op2"() : () -> ()103 }104 """105 )106 entry_block = module.body.operations[0].regions[0].blocks[0]107 ip = InsertionPoint.at_block_begin(entry_block)108 assert ip.block == entry_block109 assert ip.ref_operation == entry_block.operations[0]110 ip.insert(Operation.create("custom.op1"))111 # CHECK: "custom.op1"112 # CHECK: "custom.op2"113 module.operation.print()114 115 116run(test_insert_at_block_begin)117 118 119# CHECK-LABEL: TEST: test_insert_at_block_begin_empty120def test_insert_at_block_begin_empty():121 # TODO: Write this test case when we can create such a situation.122 pass123 124 125run(test_insert_at_block_begin_empty)126 127 128# CHECK-LABEL: TEST: test_insert_at_terminator129def test_insert_at_terminator():130 ctx = Context()131 ctx.allow_unregistered_dialects = True132 with Location.unknown(ctx):133 module = Module.parse(134 r"""135 func.func @foo() -> () {136 "custom.op1"() : () -> ()137 return138 }139 """140 )141 entry_block = module.body.operations[0].regions[0].blocks[0]142 return_op = entry_block.operations[1]143 ip = InsertionPoint.at_block_terminator(entry_block)144 assert ip.block == entry_block145 assert ip.ref_operation == return_op146 custom_op2 = Operation.create("custom.op2")147 ip.insert(custom_op2)148 InsertionPoint.after(custom_op2).insert(Operation.create("custom.op3"))149 # CHECK: "custom.op1"150 # CHECK: "custom.op2"151 # CHECK: "custom.op3"152 module.operation.print()153 154 try:155 InsertionPoint.after(return_op).insert(Operation.create("custom.op4"))156 except IndexError as e:157 # CHECK: ERROR: Cannot insert operation at the end of a block that already has a terminator.158 print(f"ERROR: {e}")159 160 161run(test_insert_at_terminator)162 163 164# CHECK-LABEL: TEST: test_insert_at_block_terminator_missing165def test_insert_at_block_terminator_missing():166 ctx = Context()167 ctx.allow_unregistered_dialects = True168 with ctx:169 module = Module.parse(170 r"""171 func.func @foo() -> () {172 "custom.op1"() : () -> ()173 }174 """175 )176 entry_block = module.body.operations[0].regions[0].blocks[0]177 try:178 ip = InsertionPoint.at_block_terminator(entry_block)179 except ValueError as e:180 # CHECK: Block has no terminator181 print(e)182 else:183 assert False, "Expected exception"184 185 186run(test_insert_at_block_terminator_missing)187 188 189# CHECK-LABEL: TEST: test_insert_at_end_with_terminator_errors190def test_insert_at_end_with_terminator_errors():191 with Context() as ctx, Location.unknown():192 ctx.allow_unregistered_dialects = True193 module = Module.parse(194 r"""195 func.func @foo() -> () {196 return197 }198 """199 )200 entry_block = module.body.operations[0].regions[0].blocks[0]201 with InsertionPoint(entry_block):202 try:203 Operation.create("custom.op1", results=[], operands=[])204 except IndexError as e:205 # CHECK: ERROR: Cannot insert operation at the end of a block that already has a terminator.206 print(f"ERROR: {e}")207 208 209run(test_insert_at_end_with_terminator_errors)210 211 212# CHECK-LABEL: TEST: test_insertion_point_context213def test_insertion_point_context():214 ctx = Context()215 ctx.allow_unregistered_dialects = True216 with Location.unknown(ctx):217 module = Module.parse(218 r"""219 func.func @foo() -> () {220 "custom.op1"() : () -> ()221 }222 """223 )224 entry_block = module.body.operations[0].regions[0].blocks[0]225 with InsertionPoint(entry_block):226 Operation.create("custom.op2")227 with InsertionPoint.at_block_begin(entry_block):228 custom_opa = Operation.create("custom.opa")229 Operation.create("custom.opb")230 Operation.create("custom.op3")231 with InsertionPoint.after(custom_opa):232 Operation.create("custom.op4")233 Operation.create("custom.op5")234 235 # CHECK: "custom.opa"236 # CHECK: "custom.op4"237 # CHECK: "custom.op5"238 # CHECK: "custom.opb"239 # CHECK: "custom.op1"240 # CHECK: "custom.op2"241 # CHECK: "custom.op3"242 module.operation.print()243 244 245run(test_insertion_point_context)246