def test_compile_arith_masked_ops(op, ty1, ty2, masked):
def func(x, y):
return op(x, y)
cc = (7, 5)
if masked[0]:
ty1 = MaskedType(ty1)
if masked[1]:
ty2 = MaskedType(ty2)
ptx, resty = compile_ptx(func, (ty1, ty2), cc=cc, device=True)
def test_compile_arith_masked_ops(op, left_dtype, right_dtype, masked):
def func(x, y):
return op(x, y)
cc = (7, 5)
ty1 = from_dtype(np.dtype(left_dtype))
ty2 = from_dtype(np.dtype(right_dtype))
if masked[0]:
ty1 = MaskedType(ty1)
if masked[1]:
ty2 = MaskedType(ty2)
ptx, resty = compile_ptx(func, (ty1, ty2), cc=cc, device=True)
def test_compile_arith_masked_vs_na(op, ty):
def func(x):
return op(x, NA)
cc = (7, 5)
ptx, resty = compile_ptx(func, (MaskedType(ty),), cc=cc, device=True)
assert isinstance(resty, MaskedType)
def test_compile_arith_constant_vs_masked(op, ty, constant):
def func(x):
return op(constant, x)
cc = (7, 5)
ptx, resty = compile_ptx(func, (MaskedType(ty),), cc=cc, device=True)
assert isinstance(resty, MaskedType)
def test_compile_arith_masked_vs_constant(op, ty, constant):
def func(x):
return op(x, constant)
cc = (7, 5)
ptx, resty = compile_ptx(func, (MaskedType(ty),), cc=cc, device=True)
assert isinstance(resty, MaskedType)
# Check that the masked typing matches that of the unmasked typing
um_ptx, um_resty = compile_ptx(func, (ty,), cc=cc, device=True)
assert resty.value_type == um_resty
def masked_scalar_null_op_impl(context, builder, sig, args):
"""
Implement `MaskedType` <op> `NAType`
or `NAType` <op> `MaskedType`
The answer to this is known up front so no actual operation
needs to take place
"""
return_type = sig.return_type # MaskedType(...)
result = cgutils.create_struct_proxy(MaskedType(return_type.value_type))(
context, builder)
# Invalidate the struct and leave `value` uninitialized
result.valid = context.get_constant(types.boolean, 0)
return result._getvalue()
def _get_frame_row_type(dtype):
"""
Get the numba `Record` type corresponding to a frame.
Models each column and its mask as a MaskedType and
models the row as a dictionary like data structure
containing these MaskedTypes.
Large parts of this function are copied with comments
from the Numba internals and slightly modified to
account for validity bools to be present in the final
struct.
See numba.np.numpy_support.from_struct_dtype for details.
"""
# Create the numpy structured type corresponding to the numpy dtype.
fields = []
offset = 0
sizes = [val[0].itemsize for val in dtype.fields.values()]
for i, (name, info) in enumerate(dtype.fields.items()):
# *info* consists of the element dtype, its offset from the beginning
# of the record, and an optional "title" containing metadata.
# We ignore the offset in info because its value assumes no masking;
# instead, we compute the correct offset based on the masked type.
elemdtype = info[0]
title = info[2] if len(info) == 3 else None
ty = numpy_support.from_dtype(elemdtype)
infos = {
"type": MaskedType(ty),
"offset": offset,
"title": title,
}
fields.append((name, infos))
# increment offset by itemsize plus one byte for validity
offset += elemdtype.itemsize + 1
# Align the next member of the struct to be a multiple of the
# memory access size, per PTX ISA 7.4/5.4.5
if i < len(sizes) - 1:
next_itemsize = sizes[i + 1]
offset = int(math.ceil(offset / next_itemsize) * next_itemsize)
# Numba requires that structures are aligned for the CUDA target
_is_aligned_struct = True
return Record(fields, offset, _is_aligned_struct)
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_scalar_kernel(sr, func, args):
sr_type = MaskedType(numpy_support.from_dtype(sr.dtype))
scalar_return_type = _get_udf_return_type(sr_type, func, args)
sig = _construct_signature(sr, scalar_return_type, args=args)
f_ = cuda.jit(device=True)(func)
global_exec_context = {
"f_": f_,
"cuda": cuda,
"Masked": Masked,
"_mask_get": _mask_get,
"pack_return": pack_return,
}
kernel_string = _scalar_kernel_string_from_template(sr, args=args)
kernel = _get_kernel(kernel_string, global_exec_context, sig, func)
return kernel, scalar_return_type
def masked_scalar_bool_impl(context, builder, sig, args):
indata = cgutils.create_struct_proxy(MaskedType(types.boolean))(
context, builder, value=args[0])
return indata.value
def test_compile_masked_unary(op, ty):
def func(x):
return op(x)
cc = (7, 5)
ptx, resty = compile_ptx(func, (MaskedType(ty),), cc=cc, device=True)
def test_compile_arith_na_vs_masked(op, ty):
def func(x):
return op(NA, x)
cc = (7, 5)
ptx, resty = compile_ptx(func, (MaskedType(ty),), cc=cc, device=True)
