209 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4 5from mlir.dialects import builtin6from mlir.dialects import cf7from mlir.dialects import func8from mlir.ir import *9 10 11def run(f):12 print("\nTEST:", f.__name__)13 f()14 gc.collect()15 assert Context._get_live_count() == 016 return f17 18 19# CHECK-LABEL: TEST: testBlockCreation20# CHECK: func @test(%[[ARG0:.*]]: i32 loc("arg0"), %[[ARG1:.*]]: i16 loc("arg1"))21# CHECK: cf.br ^bb1(%[[ARG1]] : i16)22# CHECK: ^bb1(%[[PHI0:.*]]: i16 loc("middle")):23# CHECK: cf.br ^bb2(%[[ARG0]] : i32)24# CHECK: ^bb2(%[[PHI1:.*]]: i32 loc("successor")):25# CHECK: return26@run27def testBlockCreation():28 with Context() as ctx, Location.unknown():29 module = builtin.ModuleOp()30 with InsertionPoint(module.body):31 f_type = FunctionType.get(32 [IntegerType.get_signless(32), IntegerType.get_signless(16)], []33 )34 f_op = func.FuncOp("test", f_type)35 entry_block = f_op.add_entry_block(36 [Location.name("arg0"), Location.name("arg1")]37 )38 i32_arg, i16_arg = entry_block.arguments39 successor_block = entry_block.create_after(40 i32_arg.type, arg_locs=[Location.name("successor")]41 )42 with InsertionPoint(successor_block) as successor_ip:43 assert successor_ip.block == successor_block44 func.ReturnOp([])45 middle_block = successor_block.create_before(46 i16_arg.type, arg_locs=[Location.name("middle")]47 )48 49 with InsertionPoint(entry_block) as entry_ip:50 assert entry_ip.block == entry_block51 cf.BranchOp([i16_arg], dest=middle_block)52 53 with InsertionPoint(middle_block) as middle_ip:54 assert middle_ip.block == middle_block55 cf.BranchOp([i32_arg], dest=successor_block)56 57 module.print(enable_debug_info=True)58 # Ensure region back references are coherent.59 assert entry_block.region == middle_block.region == successor_block.region60 61 assert len(entry_block.predecessors) == 062 63 assert len(entry_block.successors) == 164 assert middle_block == entry_block.successors[0]65 assert len(middle_block.predecessors) == 166 assert entry_block == middle_block.predecessors[0]67 68 assert len(middle_block.successors) == 169 assert successor_block == middle_block.successors[0]70 assert len(successor_block.predecessors) == 171 assert middle_block == successor_block.predecessors[0]72 73 assert len(successor_block.successors) == 074 75 76# CHECK-LABEL: TEST: testBlockCreationArgLocs77@run78def testBlockCreationArgLocs():79 with Context() as ctx:80 ctx.allow_unregistered_dialects = True81 f32 = F32Type.get()82 op = Operation.create("test", regions=1, loc=Location.unknown())83 blocks = op.regions[0].blocks84 85 with Location.name("default_loc"):86 blocks.append(f32)87 blocks.append()88 # CHECK: ^bb0(%{{.+}}: f32 loc("default_loc")):89 # CHECK-NEXT: ^bb1:90 op.print(enable_debug_info=True)91 92 try:93 blocks.append(f32)94 except RuntimeError as err:95 # CHECK: Missing loc: An MLIR function requires a Location but none was provided96 print("Missing loc:", err)97 98 try:99 blocks.append(f32, f32, arg_locs=[Location.unknown()])100 except ValueError as err:101 # CHECK: Wrong loc count: Expected 2 locations, got: 1102 print("Wrong loc count:", err)103 104 105# CHECK-LABEL: TEST: testFirstBlockCreation106# CHECK: func @test(%{{.*}}: f32 loc("arg_loc"))107# CHECK: return108@run109def testFirstBlockCreation():110 with Context() as ctx, Location.unknown():111 module = builtin.ModuleOp()112 f32 = F32Type.get()113 with InsertionPoint(module.body):114 f = func.FuncOp("test", ([f32], []))115 entry_block = Block.create_at_start(116 f.operation.regions[0], [f32], [Location.name("arg_loc")]117 )118 with InsertionPoint(entry_block):119 func.ReturnOp([])120 121 module.print(enable_debug_info=True)122 assert module.verify()123 assert f.body.blocks[0] == entry_block124 125 126# CHECK-LABEL: TEST: testBlockMove127# CHECK: %0 = "realop"() ({128# CHECK: ^bb0([[ARG0:%.+]]: f32):129# CHECK: "ret"([[ARG0]]) : (f32) -> ()130# CHECK: }) : () -> f32131@run132def testBlockMove():133 with Context() as ctx, Location.unknown():134 ctx.allow_unregistered_dialects = True135 module = Module.create()136 f32 = F32Type.get()137 with InsertionPoint(module.body):138 dummy = Operation.create("dummy", regions=1)139 block = Block.create_at_start(dummy.operation.regions[0], [f32])140 with InsertionPoint(block):141 ret_op = Operation.create("ret", operands=[block.arguments[0]])142 realop = Operation.create(143 "realop", results=[r.type for r in ret_op.operands], regions=1144 )145 block.append_to(realop.operation.regions[0])146 dummy.operation.erase()147 print(module)148 149 150# CHECK-LABEL: TEST: testBlockHash151@run152def testBlockHash():153 with Context() as ctx, Location.unknown():154 ctx.allow_unregistered_dialects = True155 module = Module.create()156 f32 = F32Type.get()157 with InsertionPoint(module.body):158 dummy = Operation.create("dummy", regions=1)159 block1 = Block.create_at_start(dummy.operation.regions[0], [f32])160 block2 = Block.create_at_start(dummy.operation.regions[0], [f32])161 assert hash(block1) != hash(block2)162 163 164# CHECK-LABEL: TEST: testBlockAddArgs165@run166def testBlockAddArgs():167 with Context() as ctx, Location.unknown(ctx) as loc:168 ctx.allow_unregistered_dialects = True169 f32 = F32Type.get()170 op = Operation.create("test", regions=1, loc=Location.unknown())171 blocks = op.regions[0].blocks172 blocks.append()173 # CHECK: ^bb0:174 op.print(enable_debug_info=True)175 blocks[0].add_argument(f32, loc)176 # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):177 op.print(enable_debug_info=True)178 179 180# CHECK-LABEL: TEST: testBlockEraseArgs181@run182def testBlockEraseArgs():183 with Context() as ctx, Location.unknown(ctx) as loc:184 ctx.allow_unregistered_dialects = True185 f32 = F32Type.get()186 op = Operation.create("test", regions=1, loc=Location.unknown())187 blocks = op.regions[0].blocks188 blocks.append(f32)189 # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):190 op.print(enable_debug_info=True)191 blocks[0].erase_argument(0)192 # CHECK: ^bb0:193 op.print(enable_debug_info=True)194 195 196# CHECK-LABEL: TEST: testBlockArgSetLocation197# CHECK: ^bb0(%{{.+}}: f32 loc("new_loc")):198@run199def testBlockArgSetLocation():200 with Context() as ctx, Location.unknown(ctx) as loc:201 ctx.allow_unregistered_dialects = True202 f32 = F32Type.get()203 op = Operation.create("test", regions=1, loc=Location.unknown())204 blocks = op.regions[0].blocks205 blocks.append(f32)206 arg = blocks[0].arguments[0]207 arg.set_location(Location.name("new_loc"))208 op.print(enable_debug_info=True)209