1*d7e49736SMaksim Levental# RUN: %PYTHON %s | FileCheck %s 2*d7e49736SMaksim Levental 3*d7e49736SMaksim Leventalfrom mlir.ir import * 4*d7e49736SMaksim Leventalfrom mlir.dialects import cf 5*d7e49736SMaksim Levental 6*d7e49736SMaksim Levental 7*d7e49736SMaksim Leventaldef constructAndPrintInModule(f): 8*d7e49736SMaksim Levental print("\nTEST:", f.__name__) 9*d7e49736SMaksim Levental with Context() as ctx, Location.unknown(): 10*d7e49736SMaksim Levental ctx.allow_unregistered_dialects = True 11*d7e49736SMaksim Levental module = Module.create() 12*d7e49736SMaksim Levental with InsertionPoint(module.body): 13*d7e49736SMaksim Levental f() 14*d7e49736SMaksim Levental return f 15*d7e49736SMaksim Levental 16*d7e49736SMaksim Levental 17*d7e49736SMaksim Levental# CHECK-LABEL: TEST: testBranchAndSetSuccessor 18*d7e49736SMaksim Levental@constructAndPrintInModule 19*d7e49736SMaksim Leventaldef testBranchAndSetSuccessor(): 20*d7e49736SMaksim Levental op1 = Operation.create("custom.op1", regions=1) 21*d7e49736SMaksim Levental 22*d7e49736SMaksim Levental block0 = op1.regions[0].blocks.append() 23*d7e49736SMaksim Levental ip = InsertionPoint(block0) 24*d7e49736SMaksim Levental Operation.create("custom.terminator", ip=ip) 25*d7e49736SMaksim Levental 26*d7e49736SMaksim Levental block1 = op1.regions[0].blocks.append() 27*d7e49736SMaksim Levental ip = InsertionPoint(block1) 28*d7e49736SMaksim Levental br1 = cf.BranchOp([], block1, ip=ip) 29*d7e49736SMaksim Levental # CHECK: ^bb1: // pred: ^bb1 30*d7e49736SMaksim Levental # CHECK: cf.br ^bb1 31*d7e49736SMaksim Levental print(br1.successors[0]) 32*d7e49736SMaksim Levental # CHECK: num_successors 1 33*d7e49736SMaksim Levental print("num_successors", len(br1.successors)) 34*d7e49736SMaksim Levental 35*d7e49736SMaksim Levental block2 = op1.regions[0].blocks.append() 36*d7e49736SMaksim Levental ip = InsertionPoint(block2) 37*d7e49736SMaksim Levental br2 = cf.BranchOp([], block1, ip=ip) 38*d7e49736SMaksim Levental # CHECK: ^bb1: // 2 preds: ^bb1, ^bb2 39*d7e49736SMaksim Levental # CHECK: cf.br ^bb1 40*d7e49736SMaksim Levental print(br2.successors[0]) 41*d7e49736SMaksim Levental # CHECK: num_successors 1 42*d7e49736SMaksim Levental print("num_successors", len(br2.successors)) 43*d7e49736SMaksim Levental 44*d7e49736SMaksim Levental br1.successors[0] = block2 45*d7e49736SMaksim Levental # CHECK: ^bb2: // pred: ^bb1 46*d7e49736SMaksim Levental # CHECK: cf.br ^bb1 47*d7e49736SMaksim Levental print(br1.successors[0]) 48*d7e49736SMaksim Levental # CHECK: ^bb1: // pred: ^bb2 49*d7e49736SMaksim Levental # CHECK: cf.br ^bb2 50*d7e49736SMaksim Levental print(br2.operation.successors[0]) 51