6d8dd3da4b
This allows these containers to be used in `match` statements, which allows extracting properties and asserting a shape at the same time. It seems to be only possible, to match as _either_ a `Mapping` _or_ a `Sequence`, so the `OpAttributeMap` is only a `Mapping`. I couldn't find a way to make these C++ based types properly inherit from `Sequence` or `Mapping`, so the Mixins are not provided (nanobind only allows C++ parent classes, modifying `__base__` complains about differing destructors). `OpAttributeMap` was lacking the `get` method, so I simply copied it from `collections.abc.Mapping`. When writing the tests i ran into the error, that I wrote `func.FuncOp(body=[Block(...)])` instead of `func.FuncOp(body=Region(blocks=[Block(...)]))`. So maybe also turning `Region` itself into a Sequence would be a good addition as well? Would extend the Scope of this PR, though. makslevental You suggested I make the PR, so i'm tagging you here as a potential reviewer. I hope that is ok with you. :) --------- Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
229 lines
7.8 KiB
Python
229 lines
7.8 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import gc
|
|
|
|
from mlir.dialects import builtin
|
|
from mlir.dialects import cf
|
|
from mlir.dialects import func
|
|
from mlir.ir import *
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
gc.collect()
|
|
assert Context._get_live_count() == 0
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockCreation
|
|
# CHECK: func @test(%[[ARG0:.*]]: i32 loc("arg0"), %[[ARG1:.*]]: i16 loc("arg1"))
|
|
# CHECK: cf.br ^bb1(%[[ARG1]] : i16)
|
|
# CHECK: ^bb1(%[[PHI0:.*]]: i16 loc("middle")):
|
|
# CHECK: cf.br ^bb2(%[[ARG0]] : i32)
|
|
# CHECK: ^bb2(%[[PHI1:.*]]: i32 loc("successor")):
|
|
# CHECK: return
|
|
@run
|
|
def testBlockCreation():
|
|
with Context() as ctx, Location.unknown():
|
|
module = builtin.ModuleOp()
|
|
with InsertionPoint(module.body):
|
|
f_type = FunctionType.get(
|
|
[IntegerType.get_signless(32), IntegerType.get_signless(16)], []
|
|
)
|
|
f_op = func.FuncOp("test", f_type)
|
|
entry_block = f_op.add_entry_block(
|
|
[Location.name("arg0"), Location.name("arg1")]
|
|
)
|
|
i32_arg, i16_arg = entry_block.arguments
|
|
successor_block = entry_block.create_after(
|
|
i32_arg.type, arg_locs=[Location.name("successor")]
|
|
)
|
|
with InsertionPoint(successor_block) as successor_ip:
|
|
assert successor_ip.block == successor_block
|
|
func.ReturnOp([])
|
|
middle_block = successor_block.create_before(
|
|
i16_arg.type, arg_locs=[Location.name("middle")]
|
|
)
|
|
|
|
with InsertionPoint(entry_block) as entry_ip:
|
|
assert entry_ip.block == entry_block
|
|
cf.BranchOp([i16_arg], dest=middle_block)
|
|
|
|
with InsertionPoint(middle_block) as middle_ip:
|
|
assert middle_ip.block == middle_block
|
|
cf.BranchOp([i32_arg], dest=successor_block)
|
|
|
|
module.print(enable_debug_info=True)
|
|
# Ensure region back references are coherent.
|
|
assert entry_block.region == middle_block.region == successor_block.region
|
|
|
|
assert len(entry_block.predecessors) == 0
|
|
|
|
assert len(entry_block.successors) == 1
|
|
assert middle_block == entry_block.successors[0]
|
|
assert len(middle_block.predecessors) == 1
|
|
assert entry_block == middle_block.predecessors[0]
|
|
|
|
assert len(middle_block.successors) == 1
|
|
assert successor_block == middle_block.successors[0]
|
|
assert len(successor_block.predecessors) == 1
|
|
assert middle_block == successor_block.predecessors[0]
|
|
|
|
assert len(successor_block.successors) == 0
|
|
|
|
# Same checks but using structural pattern matching.
|
|
match entry_block:
|
|
case Block(
|
|
predecessors=[],
|
|
successors=[
|
|
Block(
|
|
predecessors=[matched_entry_block],
|
|
successors=[
|
|
Block(predecessors=[matched_middle_block], successors=[]),
|
|
],
|
|
),
|
|
],
|
|
) if (
|
|
entry_block == matched_entry_block
|
|
and middle_block == matched_middle_block
|
|
):
|
|
assert True
|
|
case _:
|
|
assert False
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockCreationArgLocs
|
|
@run
|
|
def testBlockCreationArgLocs():
|
|
with Context() as ctx:
|
|
ctx.allow_unregistered_dialects = True
|
|
f32 = F32Type.get()
|
|
op = Operation.create("test", regions=1, loc=Location.unknown())
|
|
blocks = op.regions[0].blocks
|
|
|
|
with Location.name("default_loc"):
|
|
blocks.append(f32)
|
|
blocks.append()
|
|
# CHECK: ^bb0(%{{.+}}: f32 loc("default_loc")):
|
|
# CHECK-NEXT: ^bb1:
|
|
op.print(enable_debug_info=True)
|
|
|
|
try:
|
|
blocks.append(f32)
|
|
except RuntimeError as err:
|
|
# CHECK: Missing loc: An MLIR function requires a Location but none was provided
|
|
print("Missing loc:", err)
|
|
|
|
try:
|
|
blocks.append(f32, f32, arg_locs=[Location.unknown()])
|
|
except ValueError as err:
|
|
# CHECK: Wrong loc count: Expected 2 locations, got: 1
|
|
print("Wrong loc count:", err)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testFirstBlockCreation
|
|
# CHECK: func @test(%{{.*}}: f32 loc("arg_loc"))
|
|
# CHECK: return
|
|
@run
|
|
def testFirstBlockCreation():
|
|
with Context() as ctx, Location.unknown():
|
|
module = builtin.ModuleOp()
|
|
f32 = F32Type.get()
|
|
with InsertionPoint(module.body):
|
|
f = func.FuncOp("test", ([f32], []))
|
|
entry_block = Block.create_at_start(
|
|
f.operation.regions[0], [f32], [Location.name("arg_loc")]
|
|
)
|
|
with InsertionPoint(entry_block):
|
|
func.ReturnOp([])
|
|
|
|
module.print(enable_debug_info=True)
|
|
assert module.verify()
|
|
assert f.body.blocks[0] == entry_block
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockMove
|
|
# CHECK: %0 = "realop"() ({
|
|
# CHECK: ^bb0([[ARG0:%.+]]: f32):
|
|
# CHECK: "ret"([[ARG0]]) : (f32) -> ()
|
|
# CHECK: }) : () -> f32
|
|
@run
|
|
def testBlockMove():
|
|
with Context() as ctx, Location.unknown():
|
|
ctx.allow_unregistered_dialects = True
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
with InsertionPoint(module.body):
|
|
dummy = Operation.create("dummy", regions=1)
|
|
block = Block.create_at_start(dummy.operation.regions[0], [f32])
|
|
with InsertionPoint(block):
|
|
ret_op = Operation.create("ret", operands=[block.arguments[0]])
|
|
realop = Operation.create(
|
|
"realop", results=[r.type for r in ret_op.operands], regions=1
|
|
)
|
|
block.append_to(realop.operation.regions[0])
|
|
dummy.operation.erase()
|
|
print(module)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockHash
|
|
@run
|
|
def testBlockHash():
|
|
with Context() as ctx, Location.unknown():
|
|
ctx.allow_unregistered_dialects = True
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
with InsertionPoint(module.body):
|
|
dummy = Operation.create("dummy", regions=1)
|
|
block1 = Block.create_at_start(dummy.operation.regions[0], [f32])
|
|
block2 = Block.create_at_start(dummy.operation.regions[0], [f32])
|
|
assert hash(block1) != hash(block2)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockAddArgs
|
|
@run
|
|
def testBlockAddArgs():
|
|
with Context() as ctx, Location.unknown(ctx) as loc:
|
|
ctx.allow_unregistered_dialects = True
|
|
f32 = F32Type.get()
|
|
op = Operation.create("test", regions=1, loc=Location.unknown())
|
|
blocks = op.regions[0].blocks
|
|
blocks.append()
|
|
# CHECK: ^bb0:
|
|
op.print(enable_debug_info=True)
|
|
blocks[0].add_argument(f32, loc)
|
|
# CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
|
|
op.print(enable_debug_info=True)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockEraseArgs
|
|
@run
|
|
def testBlockEraseArgs():
|
|
with Context() as ctx, Location.unknown(ctx) as loc:
|
|
ctx.allow_unregistered_dialects = True
|
|
f32 = F32Type.get()
|
|
op = Operation.create("test", regions=1, loc=Location.unknown())
|
|
blocks = op.regions[0].blocks
|
|
blocks.append(f32)
|
|
# CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
|
|
op.print(enable_debug_info=True)
|
|
blocks[0].erase_argument(0)
|
|
# CHECK: ^bb0:
|
|
op.print(enable_debug_info=True)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBlockArgSetLocation
|
|
# CHECK: ^bb0(%{{.+}}: f32 loc("new_loc")):
|
|
@run
|
|
def testBlockArgSetLocation():
|
|
with Context() as ctx, Location.unknown(ctx) as loc:
|
|
ctx.allow_unregistered_dialects = True
|
|
f32 = F32Type.get()
|
|
op = Operation.create("test", regions=1, loc=Location.unknown())
|
|
blocks = op.regions[0].blocks
|
|
blocks.append(f32)
|
|
arg = blocks[0].arguments[0]
|
|
arg.set_location(Location.name("new_loc"))
|
|
op.print(enable_debug_info=True)
|