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