154 lines
6.3 KiB
Python
154 lines
6.3 KiB
Python
|
import unittest
|
||
|
|
from contextlib import contextmanager
|
||
|
|
|
||
|
|
from llvmlite import ir
|
||
|
|
|
||
|
|
from numba.core import types, typing, callconv, cpu, cgutils
|
||
|
|
from numba.core.registry import cpu_target
|
||
|
|
|
||
|
|
|
||
|
|
class TestCompileCache(unittest.TestCase):
|
||
|
|
'''
|
||
|
|
Tests that the caching in BaseContext.compile_internal() works correctly by
|
||
|
|
checking the state of the cache when it is used by the CPUContext.
|
||
|
|
'''
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def _context_builder_sig_args(self):
|
||
|
|
typing_context = cpu_target.typing_context
|
||
|
|
context = cpu_target.target_context
|
||
|
|
lib = context.codegen().create_library('testing')
|
||
|
|
with context.push_code_library(lib):
|
||
|
|
module = ir.Module("test_module")
|
||
|
|
|
||
|
|
sig = typing.signature(types.int32, types.int32)
|
||
|
|
llvm_fnty = context.call_conv.get_function_type(sig.return_type,
|
||
|
|
sig.args)
|
||
|
|
function = cgutils.get_or_insert_function(module, llvm_fnty,
|
||
|
|
'test_fn')
|
||
|
|
args = context.call_conv.get_arguments(function)
|
||
|
|
assert function.is_declaration
|
||
|
|
entry_block = function.append_basic_block('entry')
|
||
|
|
builder = ir.IRBuilder(entry_block)
|
||
|
|
|
||
|
|
yield context, builder, sig, args
|
||
|
|
|
||
|
|
def test_cache(self):
|
||
|
|
def times2(i):
|
||
|
|
return 2*i
|
||
|
|
|
||
|
|
def times3(i):
|
||
|
|
return i*3
|
||
|
|
|
||
|
|
with self._context_builder_sig_args() as (
|
||
|
|
context, builder, sig, args,
|
||
|
|
):
|
||
|
|
initial_cache_size = len(context.cached_internal_func)
|
||
|
|
# Ensure the cache is empty to begin with
|
||
|
|
self.assertEqual(initial_cache_size + 0,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
# After one compile, it should contain one entry
|
||
|
|
context.compile_internal(builder, times2, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 1,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
# After a second compilation of the same thing, it should still contain
|
||
|
|
# one entry
|
||
|
|
context.compile_internal(builder, times2, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 1,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
# After compilation of another function, the cache should have grown by
|
||
|
|
# one more.
|
||
|
|
context.compile_internal(builder, times3, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 2,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
sig2 = typing.signature(types.float64, types.float64)
|
||
|
|
llvm_fnty2 = context.call_conv.get_function_type(sig2.return_type,
|
||
|
|
sig2.args)
|
||
|
|
function2 = cgutils.get_or_insert_function(builder.module,
|
||
|
|
llvm_fnty2, 'test_fn_2')
|
||
|
|
args2 = context.call_conv.get_arguments(function2)
|
||
|
|
assert function2.is_declaration
|
||
|
|
entry_block2 = function2.append_basic_block('entry')
|
||
|
|
builder2 = ir.IRBuilder(entry_block2)
|
||
|
|
|
||
|
|
# Ensure that the same function with a different signature does not
|
||
|
|
# reuse an entry from the cache in error
|
||
|
|
context.compile_internal(builder2, times3, sig2, args2)
|
||
|
|
self.assertEqual(initial_cache_size + 3,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
def test_closures(self):
|
||
|
|
"""
|
||
|
|
Caching must not mix up closures reusing the same code object.
|
||
|
|
"""
|
||
|
|
def make_closure(x, y):
|
||
|
|
def f(z):
|
||
|
|
return y + z
|
||
|
|
return f
|
||
|
|
|
||
|
|
with self._context_builder_sig_args() as (
|
||
|
|
context, builder, sig, args,
|
||
|
|
):
|
||
|
|
# Closures with distinct cell contents must each be compiled.
|
||
|
|
clo11 = make_closure(1, 1)
|
||
|
|
clo12 = make_closure(1, 2)
|
||
|
|
clo22 = make_closure(2, 2)
|
||
|
|
initial_cache_size = len(context.cached_internal_func)
|
||
|
|
res1 = context.compile_internal(builder, clo11, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 1,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
res2 = context.compile_internal(builder, clo12, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 2,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
# Same cell contents as above (first parameter isn't captured)
|
||
|
|
res3 = context.compile_internal(builder, clo22, sig, args)
|
||
|
|
self.assertEqual(initial_cache_size + 2,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
def test_error_model(self):
|
||
|
|
"""
|
||
|
|
Caching must not mix up different error models.
|
||
|
|
"""
|
||
|
|
def inv(x):
|
||
|
|
return 1.0 / x
|
||
|
|
|
||
|
|
inv_sig = typing.signature(types.float64, types.float64)
|
||
|
|
|
||
|
|
def compile_inv(context):
|
||
|
|
return context.compile_subroutine(builder, inv, inv_sig)
|
||
|
|
|
||
|
|
with self._context_builder_sig_args() as (
|
||
|
|
context, builder, sig, args,
|
||
|
|
):
|
||
|
|
py_error_model = callconv.create_error_model('python', context)
|
||
|
|
np_error_model = callconv.create_error_model('numpy', context)
|
||
|
|
|
||
|
|
py_context1 = context.subtarget(error_model=py_error_model)
|
||
|
|
py_context2 = context.subtarget(error_model=py_error_model)
|
||
|
|
np_context = context.subtarget(error_model=np_error_model)
|
||
|
|
|
||
|
|
initial_cache_size = len(context.cached_internal_func)
|
||
|
|
|
||
|
|
# Note the parent context's cache is shared by subtargets
|
||
|
|
self.assertEqual(initial_cache_size + 0,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
# Compiling with the same error model reuses the same cache slot
|
||
|
|
compile_inv(py_context1)
|
||
|
|
self.assertEqual(initial_cache_size + 1,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
compile_inv(py_context2)
|
||
|
|
self.assertEqual(initial_cache_size + 1,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
# Compiling with another error model creates a new cache slot
|
||
|
|
compile_inv(np_context)
|
||
|
|
self.assertEqual(initial_cache_size + 2,
|
||
|
|
len(context.cached_internal_func))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|