51 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import cf5 6 7def constructAndPrintInModule(f):8 print("\nTEST:", f.__name__)9 with Context() as ctx, Location.unknown():10 ctx.allow_unregistered_dialects = True11 module = Module.create()12 with InsertionPoint(module.body):13 f()14 return f15 16 17# CHECK-LABEL: TEST: testBranchAndSetSuccessor18@constructAndPrintInModule19def testBranchAndSetSuccessor():20 op1 = Operation.create("custom.op1", regions=1)21 22 block0 = op1.regions[0].blocks.append()23 ip = InsertionPoint(block0)24 Operation.create("custom.terminator", ip=ip)25 26 block1 = op1.regions[0].blocks.append()27 ip = InsertionPoint(block1)28 br1 = cf.BranchOp([], block1, ip=ip)29 # CHECK: ^bb1: // pred: ^bb130 # CHECK: cf.br ^bb131 print(br1.successors[0])32 # CHECK: num_successors 133 print("num_successors", len(br1.successors))34 35 block2 = op1.regions[0].blocks.append()36 ip = InsertionPoint(block2)37 br2 = cf.BranchOp([], block1, ip=ip)38 # CHECK: ^bb1: // 2 preds: ^bb1, ^bb239 # CHECK: cf.br ^bb140 print(br2.successors[0])41 # CHECK: num_successors 142 print("num_successors", len(br2.successors))43 44 br1.successors[0] = block245 # CHECK: ^bb2: // pred: ^bb146 # CHECK: cf.br ^bb147 print(br1.successors[0])48 # CHECK: ^bb1: // pred: ^bb249 # CHECK: cf.br ^bb250 print(br2.operation.successors[0])51