527 lines
14 KiB
Python
527 lines
14 KiB
Python
|
import unittest
|
||
|
from llvmlite import ir
|
||
|
from llvmlite import binding as llvm
|
||
|
from llvmlite.tests import TestCase
|
||
|
|
||
|
from . import refprune_proto as proto
|
||
|
|
||
|
|
||
|
def _iterate_cases(generate_test):
|
||
|
def wrap(fn):
|
||
|
def wrapped(self):
|
||
|
return generate_test(self, fn)
|
||
|
wrapped.__doc__ = f"generated test for {fn.__module__}.{fn.__name__}"
|
||
|
return wrapped
|
||
|
|
||
|
for k, case_fn in proto.__dict__.items():
|
||
|
if k.startswith('case'):
|
||
|
yield f'test_{k}', wrap(case_fn)
|
||
|
|
||
|
|
||
|
class TestRefPrunePrototype(TestCase):
|
||
|
"""
|
||
|
Test that the prototype is working.
|
||
|
"""
|
||
|
def generate_test(self, case_gen):
|
||
|
nodes, edges, expected = case_gen()
|
||
|
got = proto.FanoutAlgorithm(nodes, edges).run()
|
||
|
self.assertEqual(expected, got)
|
||
|
|
||
|
# Generate tests
|
||
|
for name, case in _iterate_cases(generate_test):
|
||
|
locals()[name] = case
|
||
|
|
||
|
|
||
|
ptr_ty = ir.IntType(8).as_pointer()
|
||
|
|
||
|
|
||
|
class TestRefPrunePass(TestCase):
|
||
|
"""
|
||
|
Test that the C++ implementation matches the expected behavior as for
|
||
|
the prototype.
|
||
|
|
||
|
This generates a LLVM module for each test case, runs the pruner and checks
|
||
|
that the expected results are achieved.
|
||
|
"""
|
||
|
|
||
|
def make_incref(self, m):
|
||
|
fnty = ir.FunctionType(ir.VoidType(), [ptr_ty])
|
||
|
return ir.Function(m, fnty, name='NRT_incref')
|
||
|
|
||
|
def make_decref(self, m):
|
||
|
fnty = ir.FunctionType(ir.VoidType(), [ptr_ty])
|
||
|
return ir.Function(m, fnty, name='NRT_decref')
|
||
|
|
||
|
def make_switcher(self, m):
|
||
|
fnty = ir.FunctionType(ir.IntType(32), ())
|
||
|
return ir.Function(m, fnty, name='switcher')
|
||
|
|
||
|
def make_brancher(self, m):
|
||
|
fnty = ir.FunctionType(ir.IntType(1), ())
|
||
|
return ir.Function(m, fnty, name='brancher')
|
||
|
|
||
|
def generate_ir(self, nodes, edges):
|
||
|
# Build LLVM module for the CFG
|
||
|
m = ir.Module()
|
||
|
|
||
|
incref_fn = self.make_incref(m)
|
||
|
decref_fn = self.make_decref(m)
|
||
|
switcher_fn = self.make_switcher(m)
|
||
|
brancher_fn = self.make_brancher(m)
|
||
|
|
||
|
fnty = ir.FunctionType(ir.VoidType(), [ptr_ty])
|
||
|
fn = ir.Function(m, fnty, name='main')
|
||
|
[ptr] = fn.args
|
||
|
ptr.name = 'mem'
|
||
|
# populate the BB nodes
|
||
|
bbmap = {}
|
||
|
for bb in edges:
|
||
|
bbmap[bb] = fn.append_basic_block(bb)
|
||
|
# populate the BB
|
||
|
builder = ir.IRBuilder()
|
||
|
for bb, jump_targets in edges.items():
|
||
|
builder.position_at_end(bbmap[bb])
|
||
|
# Insert increfs and decrefs
|
||
|
for action in nodes[bb]:
|
||
|
if action == 'incref':
|
||
|
builder.call(incref_fn, [ptr])
|
||
|
elif action == 'decref':
|
||
|
builder.call(decref_fn, [ptr])
|
||
|
else:
|
||
|
raise AssertionError('unreachable')
|
||
|
|
||
|
# Insert the terminator.
|
||
|
# Switch base on the number of jump targets.
|
||
|
n_targets = len(jump_targets)
|
||
|
if n_targets == 0:
|
||
|
builder.ret_void()
|
||
|
elif n_targets == 1:
|
||
|
[dst] = jump_targets
|
||
|
builder.branch(bbmap[dst])
|
||
|
elif n_targets == 2:
|
||
|
[left, right] = jump_targets
|
||
|
sel = builder.call(brancher_fn, ())
|
||
|
builder.cbranch(sel, bbmap[left], bbmap[right])
|
||
|
elif n_targets > 2:
|
||
|
sel = builder.call(switcher_fn, ())
|
||
|
[head, *tail] = jump_targets
|
||
|
|
||
|
sw = builder.switch(sel, default=bbmap[head])
|
||
|
for i, dst in enumerate(tail):
|
||
|
sw.add_case(sel.type(i), bbmap[dst])
|
||
|
else:
|
||
|
raise AssertionError('unreachable')
|
||
|
|
||
|
return m
|
||
|
|
||
|
def apply_refprune(self, irmod):
|
||
|
mod = llvm.parse_assembly(str(irmod))
|
||
|
pm = llvm.ModulePassManager()
|
||
|
pm.add_refprune_pass()
|
||
|
pm.run(mod)
|
||
|
return mod
|
||
|
|
||
|
def check(self, mod, expected, nodes):
|
||
|
# preprocess incref/decref locations
|
||
|
d = {}
|
||
|
for k, vs in nodes.items():
|
||
|
n_incref = vs.count('incref')
|
||
|
n_decref = vs.count('decref')
|
||
|
d[k] = {'incref': n_incref, 'decref': n_decref}
|
||
|
for k, stats in d.items():
|
||
|
if expected.get(k):
|
||
|
stats['incref'] -= 1
|
||
|
for dec_bb in expected[k]:
|
||
|
d[dec_bb]['decref'] -= 1
|
||
|
|
||
|
# find the main function
|
||
|
for f in mod.functions:
|
||
|
if f.name == 'main':
|
||
|
break
|
||
|
# check each BB
|
||
|
for bb in f.blocks:
|
||
|
stats = d[bb.name]
|
||
|
text = str(bb)
|
||
|
n_incref = text.count('NRT_incref')
|
||
|
n_decref = text.count('NRT_decref')
|
||
|
self.assertEqual(stats['incref'], n_incref, msg=f'BB {bb}')
|
||
|
self.assertEqual(stats['decref'], n_decref, msg=f'BB {bb}')
|
||
|
|
||
|
def generate_test(self, case_gen):
|
||
|
nodes, edges, expected = case_gen()
|
||
|
irmod = self.generate_ir(nodes, edges)
|
||
|
outmod = self.apply_refprune(irmod)
|
||
|
self.check(outmod, expected, nodes)
|
||
|
|
||
|
# Generate tests
|
||
|
for name, case in _iterate_cases(generate_test):
|
||
|
locals()[name] = case
|
||
|
|
||
|
|
||
|
class BaseTestByIR(TestCase):
|
||
|
refprune_bitmask = 0
|
||
|
|
||
|
prologue = r"""
|
||
|
declare void @NRT_incref(i8* %ptr)
|
||
|
declare void @NRT_decref(i8* %ptr)
|
||
|
"""
|
||
|
|
||
|
def check(self, irmod, subgraph_limit=None):
|
||
|
mod = llvm.parse_assembly(f"{self.prologue}\n{irmod}")
|
||
|
pm = llvm.ModulePassManager()
|
||
|
if subgraph_limit is None:
|
||
|
pm.add_refprune_pass(self.refprune_bitmask)
|
||
|
else:
|
||
|
pm.add_refprune_pass(self.refprune_bitmask,
|
||
|
subgraph_limit=subgraph_limit)
|
||
|
before = llvm.dump_refprune_stats()
|
||
|
pm.run(mod)
|
||
|
after = llvm.dump_refprune_stats()
|
||
|
return mod, after - before
|
||
|
|
||
|
|
||
|
class TestPerBB(BaseTestByIR):
|
||
|
refprune_bitmask = llvm.RefPruneSubpasses.PER_BB
|
||
|
|
||
|
per_bb_ir_1 = r"""
|
||
|
define void @main(i8* %ptr) {
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_bb_1(self):
|
||
|
mod, stats = self.check(self.per_bb_ir_1)
|
||
|
self.assertEqual(stats.basicblock, 2)
|
||
|
|
||
|
per_bb_ir_2 = r"""
|
||
|
define void @main(i8* %ptr) {
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_bb_2(self):
|
||
|
mod, stats = self.check(self.per_bb_ir_2)
|
||
|
self.assertEqual(stats.basicblock, 4)
|
||
|
# not pruned
|
||
|
self.assertIn("call void @NRT_incref(i8* %ptr)", str(mod))
|
||
|
|
||
|
per_bb_ir_3 = r"""
|
||
|
define void @main(i8* %ptr, i8* %other) {
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %other)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_bb_3(self):
|
||
|
mod, stats = self.check(self.per_bb_ir_3)
|
||
|
self.assertEqual(stats.basicblock, 2)
|
||
|
# not pruned
|
||
|
self.assertIn("call void @NRT_decref(i8* %other)", str(mod))
|
||
|
|
||
|
per_bb_ir_4 = r"""
|
||
|
; reordered
|
||
|
define void @main(i8* %ptr, i8* %other) {
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %other)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_bb_4(self):
|
||
|
mod, stats = self.check(self.per_bb_ir_4)
|
||
|
self.assertEqual(stats.basicblock, 4)
|
||
|
# not pruned
|
||
|
self.assertIn("call void @NRT_decref(i8* %other)", str(mod))
|
||
|
|
||
|
|
||
|
class TestDiamond(BaseTestByIR):
|
||
|
refprune_bitmask = llvm.RefPruneSubpasses.DIAMOND
|
||
|
|
||
|
per_diamond_1 = r"""
|
||
|
define void @main(i8* %ptr) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br label %bb_B
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_diamond_1(self):
|
||
|
mod, stats = self.check(self.per_diamond_1)
|
||
|
self.assertEqual(stats.diamond, 2)
|
||
|
|
||
|
per_diamond_2 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
br label %bb_D
|
||
|
bb_C:
|
||
|
br label %bb_D
|
||
|
bb_D:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_diamond_2(self):
|
||
|
mod, stats = self.check(self.per_diamond_2)
|
||
|
self.assertEqual(stats.diamond, 2)
|
||
|
|
||
|
per_diamond_3 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
br label %bb_D
|
||
|
bb_C:
|
||
|
call void @NRT_decref(i8* %ptr) ; reject because of decref in diamond
|
||
|
br label %bb_D
|
||
|
bb_D:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_diamond_3(self):
|
||
|
mod, stats = self.check(self.per_diamond_3)
|
||
|
self.assertEqual(stats.diamond, 0)
|
||
|
|
||
|
per_diamond_4 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_incref(i8* %ptr) ; extra incref will not affect prune
|
||
|
br label %bb_D
|
||
|
bb_C:
|
||
|
br label %bb_D
|
||
|
bb_D:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_diamond_4(self):
|
||
|
mod, stats = self.check(self.per_diamond_4)
|
||
|
self.assertEqual(stats.diamond, 2)
|
||
|
|
||
|
per_diamond_5 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
br label %bb_D
|
||
|
bb_C:
|
||
|
br label %bb_D
|
||
|
bb_D:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_per_diamond_5(self):
|
||
|
mod, stats = self.check(self.per_diamond_5)
|
||
|
self.assertEqual(stats.diamond, 4)
|
||
|
|
||
|
|
||
|
class TestFanout(BaseTestByIR):
|
||
|
"""More complex cases are tested in TestRefPrunePass
|
||
|
"""
|
||
|
|
||
|
refprune_bitmask = llvm.RefPruneSubpasses.FANOUT
|
||
|
|
||
|
fanout_1 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
bb_C:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_1(self):
|
||
|
mod, stats = self.check(self.fanout_1)
|
||
|
self.assertEqual(stats.fanout, 3)
|
||
|
|
||
|
fanout_2 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
bb_C:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
br label %bb_B ; illegal jump to other decref
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_2(self):
|
||
|
mod, stats = self.check(self.fanout_2)
|
||
|
self.assertEqual(stats.fanout, 0)
|
||
|
|
||
|
fanout_3 = r"""
|
||
|
define void @main(i8* %ptr, i1 %cond) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
bb_C:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret void
|
||
|
}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_3(self):
|
||
|
mod, stats = self.check(self.fanout_3)
|
||
|
self.assertEqual(stats.fanout, 6)
|
||
|
|
||
|
def test_fanout_3_limited(self):
|
||
|
# With subgraph limit at 1, it is essentially turning off the fanout
|
||
|
# pruner.
|
||
|
mod, stats = self.check(self.fanout_3, subgraph_limit=1)
|
||
|
self.assertEqual(stats.fanout, 0)
|
||
|
|
||
|
|
||
|
class TestFanoutRaise(BaseTestByIR):
|
||
|
refprune_bitmask = llvm.RefPruneSubpasses.FANOUT_RAISE
|
||
|
|
||
|
fanout_raise_1 = r"""
|
||
|
define i32 @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret i32 0
|
||
|
bb_C:
|
||
|
store i8* null, i8** %excinfo, !numba_exception_output !0
|
||
|
ret i32 1
|
||
|
}
|
||
|
!0 = !{i1 true}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_raise_1(self):
|
||
|
mod, stats = self.check(self.fanout_raise_1)
|
||
|
self.assertEqual(stats.fanout_raise, 2)
|
||
|
|
||
|
fanout_raise_2 = r"""
|
||
|
define i32 @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret i32 0
|
||
|
bb_C:
|
||
|
store i8* null, i8** %excinfo, !numba_exception_typo !0 ; bad metadata
|
||
|
ret i32 1
|
||
|
}
|
||
|
|
||
|
!0 = !{i1 true}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_raise_2(self):
|
||
|
# This is ensuring that fanout_raise is not pruning when the metadata
|
||
|
# is incorrectly named.
|
||
|
mod, stats = self.check(self.fanout_raise_2)
|
||
|
self.assertEqual(stats.fanout_raise, 0)
|
||
|
|
||
|
fanout_raise_3 = r"""
|
||
|
define i32 @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
ret i32 0
|
||
|
bb_C:
|
||
|
store i8* null, i8** %excinfo, !numba_exception_output !0
|
||
|
ret i32 1
|
||
|
}
|
||
|
|
||
|
!0 = !{i32 1} ; ok; use i32
|
||
|
"""
|
||
|
|
||
|
def test_fanout_raise_3(self):
|
||
|
mod, stats = self.check(self.fanout_raise_3)
|
||
|
self.assertEqual(stats.fanout_raise, 2)
|
||
|
|
||
|
fanout_raise_4 = r"""
|
||
|
define i32 @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
ret i32 1 ; BAD; all tails are raising without decref
|
||
|
bb_C:
|
||
|
ret i32 1 ; BAD; all tails are raising without decref
|
||
|
}
|
||
|
|
||
|
!0 = !{i1 1}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_raise_4(self):
|
||
|
mod, stats = self.check(self.fanout_raise_4)
|
||
|
self.assertEqual(stats.fanout_raise, 0)
|
||
|
|
||
|
fanout_raise_5 = r"""
|
||
|
define i32 @main(i8* %ptr, i1 %cond, i8** %excinfo) {
|
||
|
bb_A:
|
||
|
call void @NRT_incref(i8* %ptr)
|
||
|
br i1 %cond, label %bb_B, label %bb_C
|
||
|
bb_B:
|
||
|
call void @NRT_decref(i8* %ptr)
|
||
|
br label %common.ret
|
||
|
bb_C:
|
||
|
store i8* null, i8** %excinfo, !numba_exception_output !0
|
||
|
br label %common.ret
|
||
|
common.ret:
|
||
|
%common.ret.op = phi i32 [ 0, %bb_B ], [ 1, %bb_C ]
|
||
|
ret i32 %common.ret.op
|
||
|
}
|
||
|
!0 = !{i1 1}
|
||
|
"""
|
||
|
|
||
|
def test_fanout_raise_5(self):
|
||
|
mod, stats = self.check(self.fanout_raise_5)
|
||
|
self.assertEqual(stats.fanout_raise, 2)
|
||
|
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
unittest.main()
|