def test_code_cache_different_types_miss():
# Compilation of a distinct function with the same code but different types
# should create an additional cache entry
cudautils.compile_udf(lambda x: x + 1, (types.float32, ))
assert_cache_size(1)
cudautils.compile_udf(lambda x: x + 1, (types.float64, ))
assert_cache_size(2)
def test_lambda_in_loop_code_cached():
# Compiling a UDF defined in a loop should result in the code cache being
# reused for each loop iteration after the first. We check for this by
# ensuring that there is only one entry in the code cache after the loop.
for i in range(3):
cudautils.compile_udf(lambda x: x + 1, (types.float32, ))
assert_cache_size(1)
def test_code_cache_same_code_different_function_hit():
# Compilation of a distinct function with the same code and signature
# should reuse the cached entry
cudautils.compile_udf(lambda x: x + 1, (types.float32, ))
assert_cache_size(1)
cudautils.compile_udf(lambda x: x + 1, (types.float32, ))
assert_cache_size(1)
def test_code_cache_different_cvars_miss():
# Compilation of a distinct function with the same types and code as an
# existing entry but different closure variables should create an
# additional cache entry
def gen_closure(y):
return lambda x: x + y
cudautils.compile_udf(gen_closure(1), (types.float32, ))
assert_cache_size(1)
cudautils.compile_udf(gen_closure(2), (types.float32, ))
assert_cache_size(2)
def compile_masked_udf(func, dtypes):
"""
Generate an inlineable PTX function that will be injected into
a variadic kernel inside libcudf
assume all input types are `MaskedType(input_col.dtype)` and then
compile the requestied PTX function as a function over those types
"""
to_compiler_sig = tuple(
MaskedType(arg)
for arg in (numpy_support.from_dtype(np_type) for np_type in dtypes))
# Get the inlineable PTX function
ptx, numba_output_type = cudautils.compile_udf(func, to_compiler_sig)
numpy_output_type = numpy_support.as_dtype(numba_output_type.value_type)
return numpy_output_type, ptx
def _get_udf_return_type(argty, func: Callable, args=()):
"""
Get the return type of a masked UDF for a given set of argument dtypes. It
is assumed that the function consumes a dictionary whose keys are strings
and whose values are of MaskedType. Initially assume that the UDF may be
written to utilize any field in the row - including those containing an
unsupported dtype. If an unsupported dtype is actually used in the function
the compilation should fail at `compile_udf`. If compilation succeeds, one
can infer that the function does not use any of the columns of unsupported
dtype - meaning we can drop them going forward and the UDF will still end
up getting fed rows containing all the fields it actually needs to use to
compute the answer for that row.
"""
# present a row containing all fields to the UDF and try and compile
compile_sig = (argty, *(typeof(arg) for arg in args))
# Get the return type. The PTX is also returned by compile_udf, but is not
# needed here.
ptx, output_type = cudautils.compile_udf(func, compile_sig)
if not isinstance(output_type, MaskedType):
numba_output_type = numpy_support.from_dtype(np.dtype(output_type))
else:
numba_output_type = output_type
result = (numba_output_type
if not isinstance(numba_output_type, MaskedType) else
numba_output_type.value_type)
# _get_udf_return_type will throw a TypingError if the user tries to use
# a field in the row containing an unsupported dtype, except in the
# edge case where all the function does is return that element:
# def f(row):
# return row[<bad dtype key>]
# In this case numba is happy to return MaskedType(<bad dtype key>)
# because it relies on not finding overloaded operators for types to raise
# the exception, so we have to explicitly check for that case.
if isinstance(result, Poison):
raise TypingError(str(result))
return result
def test_first_compile_sets_cache_entry():
# The first compilation should put an entry in the cache
cudautils.compile_udf(lambda x: x + 1, (types.float32, ))
assert_cache_size(1)