1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4import io
5import itertools
6from mlir.ir import *
7from mlir.dialects import builtin
8# Note: std dialect needed for terminators.
9from mlir.dialects import std
10
11
12def run(f):
13  print("\nTEST:", f.__name__)
14  f()
15  gc.collect()
16  assert Context._get_live_count() == 0
17  return f
18
19
20# CHECK-LABEL: TEST: testBlockCreation
21# CHECK: func @test(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i16)
22# CHECK:   br ^bb1(%[[ARG1]] : i16)
23# CHECK: ^bb1(%[[PHI0:.*]]: i16):
24# CHECK:   br ^bb2(%[[ARG0]] : i32)
25# CHECK: ^bb2(%[[PHI1:.*]]: i32):
26# CHECK:   return
27@run
28def testBlockCreation():
29  with Context() as ctx, Location.unknown():
30    module = Module.create()
31    with InsertionPoint(module.body):
32      f_type = FunctionType.get(
33          [IntegerType.get_signless(32),
34           IntegerType.get_signless(16)], [])
35      f_op = builtin.FuncOp("test", f_type)
36      entry_block = f_op.add_entry_block()
37      i32_arg, i16_arg = entry_block.arguments
38      successor_block = entry_block.create_after(i32_arg.type)
39      with InsertionPoint(successor_block) as successor_ip:
40        assert successor_ip.block == successor_block
41        std.ReturnOp([])
42      middle_block = successor_block.create_before(i16_arg.type)
43
44      with InsertionPoint(entry_block) as entry_ip:
45        assert entry_ip.block == entry_block
46        std.BranchOp([i16_arg], dest=middle_block)
47
48      with InsertionPoint(middle_block) as middle_ip:
49        assert middle_ip.block == middle_block
50        std.BranchOp([i32_arg], dest=successor_block)
51    print(module.operation)
52    # Ensure region back references are coherent.
53    assert entry_block.region == middle_block.region == successor_block.region
54