xref: /llvm-project/mlir/test/python/dialects/cf.py (revision d7e49736e61a6b8c622bdb05d4b808c2e3257e3d) 
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import cf
5
6
7def constructAndPrintInModule(f):
8    print("\nTEST:", f.__name__)
9    with Context() as ctx, Location.unknown():
10        ctx.allow_unregistered_dialects = True
11        module = Module.create()
12        with InsertionPoint(module.body):
13            f()
14    return f
15
16
17# CHECK-LABEL: TEST: testBranchAndSetSuccessor
18@constructAndPrintInModule
19def 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: ^bb1
30    # CHECK:   cf.br ^bb1
31    print(br1.successors[0])
32    # CHECK: num_successors 1
33    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, ^bb2
39    # CHECK:   cf.br ^bb1
40    print(br2.successors[0])
41    # CHECK: num_successors 1
42    print("num_successors", len(br2.successors))
43
44    br1.successors[0] = block2
45    # CHECK: ^bb2:  // pred: ^bb1
46    # CHECK:   cf.br ^bb1
47    print(br1.successors[0])
48    # CHECK: ^bb1:  // pred: ^bb2
49    # CHECK:   cf.br ^bb2
50    print(br2.operation.successors[0])
51