def test_allocation_functions():
host = TargetInfo.host()
dhost = TargetInfo.host(use_tracing_allocator=True)
assert host is not dhost # check that use_tracing_allocator is part of the cache key
assert host.info[
'fn_allocate_varlen_buffer'] == 'rbclib_allocate_varlen_buffer'
assert host.info['fn_free_buffer'] == 'rbclib_free_buffer'
#
assert dhost.info[
'fn_allocate_varlen_buffer'] == 'rbclib_tracing_allocate_varlen_buffer'
assert dhost.info['fn_free_buffer'] == 'rbclib_tracing_free_buffer'
def omnisci_buffer_idx_set_null(typingctx, arr, row_idx):
T = arr.eltype
sig = types.none(arr, row_idx)
target_info = TargetInfo()
null_value = target_info.null_values[f'{T}']
# The server sends numbers as unsigned values rather than signed ones.
# Thus, 129 should be read as -127 (overflow). See rbc issue #254
bitwidth = T.bitwidth
null_value = np.dtype(f'uint{bitwidth}').type(null_value).view(
f'int{bitwidth}')
def codegen(context, builder, signature, args):
# get the operator.setitem intrinsic
fnop = context.typing_context.resolve_value_type(
omnisci_buffer_ptr_setitem_)
setitem_sig = types.none(arr, row_idx, T)
# register the intrinsic in the typing ctx
fnop.get_call_type(context.typing_context, setitem_sig.args, {})
intrinsic = context.get_function(fnop, setitem_sig)
data, index = args
# data = {T*, i64, i8}*
ty = data.type.pointee.elements[0].pointee
nv = ir.Constant(ir.IntType(T.bitwidth), null_value)
if isinstance(T, types.Float):
nv = builder.bitcast(nv, ty)
intrinsic(builder, (
data,
index,
nv,
))
return sig, codegen
def omnisci_udtfmanager_error_message_(typingctx, mgr, msg):
sig = types.int32(mgr, msg)
target_info = TargetInfo()
if target_info.software[1][:3] < (5, 9, 0):
raise UnsupportedError(error_msg % (".".join(map(str, target_info.software[1]))))
if not isinstance(msg, types.StringLiteral):
raise TypeError(f"expected StringLiteral but got {type(msg).__name__}")
def codegen(context, builder, signature, args):
mgr_ptr = args[0]
mgr_i8ptr = builder.bitcast(mgr_ptr, i8p)
msg_bytes = msg.literal_value.encode('utf-8')
msg_const = make_bytearray(msg_bytes + b'\0')
msg_global_var = global_constant(
builder.module,
"table_function_manager_error_message",
msg_const)
msg_ptr = builder.bitcast(msg_global_var, i8p)
fnty = ir.FunctionType(i32, [i8p, i8p])
fn = irutils.get_or_insert_function(
builder.module, fnty, "TableFunctionManager_error_message")
return builder.call(fn, [mgr_i8ptr, msg_ptr])
return sig, codegen
def omnisci_column_set_null_(typingctx, col_var, row_idx):
# Float values are serialized as integers by OmniSciDB
# For reference, here is the conversion table for float and double
# FLOAT: 1.1754944e-38 -> 8388608
# DOUBLE: 2.2250738585072014e-308 -> 4503599627370496
# ^ ^
# fp value serialized
T = col_var.eltype
sig = types.void(col_var, row_idx)
target_info = TargetInfo()
null_value = target_info.null_values[str(T)]
def codegen(context, builder, signature, args):
zero = int32_t(0)
data, index = args
assert data.opname == 'load'
buf = data.operands[0]
ptr = builder.load(builder.gep(buf, [zero, zero]))
ty = ptr.type.pointee
nv = ir.Constant(ir.IntType(T.bitwidth), null_value)
if isinstance(T, types.Float):
nv = builder.bitcast(nv, ty)
builder.store(nv, builder.gep(ptr, [index]))
return sig, codegen
def codegen(context, builder, signature, args):
buffers = builder_buffers[builder]
# TODO: using stdlib `free` that works only for CPU. For CUDA
# devices, we need to use omniscidb provided deallocator.
target_info = TargetInfo()
try:
free_buffer_fn_name = target_info.info['fn_free_buffer']
except KeyError as msg:
raise UnsupportedError(f'{target_info} does not provide {msg}')
free_buffer_fnty = llvm_ir.FunctionType(void_t, [int8_t.as_pointer()])
free_buffer_fn = irutils.get_or_insert_function(
builder.module, free_buffer_fnty, free_buffer_fn_name)
if isinstance(value_to_keep_alive, BufferPointer):
# free all the buffers apart value_to_keep_alive
[keep_alive] = args
keep_alive_ptr = builder.load(
builder.gep(keep_alive, [int32_t(0), int32_t(0)]))
keep_alive_ptr = builder.bitcast(keep_alive_ptr,
int8_t.as_pointer())
for ptr8 in buffers:
with builder.if_then(
builder.icmp_signed('!=', keep_alive_ptr, ptr8)):
builder.call(free_buffer_fn, [ptr8])
else:
# free all the buffers unconditionally
for ptr8 in buffers:
builder.call(free_buffer_fn, [ptr8])
del builder_buffers[builder]
def generic(self, args, kws):
# get the correct signature and function name for the current device
atypes = tuple(map(Type.fromobject, args))
t = self.obj.match_signature(atypes)
TargetInfo().add_external(t.name)
codegen = self.obj.get_codegen()
extending.lower_builtin(self.key, *t.tonumba().args)(codegen)
return t.tonumba()
def codegen(context, builder, sig, args):
target_info = TargetInfo()
if target_info.software[1][:3] < (5, 7, 0):
msg = 'set_output_row_size is only available in OmniSciDB 5.7 or newer'
raise UnsupportedError(msg)
fnty = ir.FunctionType(ir.VoidType(), [ir.IntType(64)])
fn = irutils.get_or_insert_function(builder.module, fnty, name="set_output_row_size")
assert fn.is_declaration
builder.call(fn, args) # don't return anything
def codegen(context, builder, signature, args):
# TODO: using stdlib `free` that works only for CPU. For CUDA
# devices, we need to use omniscidb provided deallocator.
target_info = TargetInfo()
free_buffer_fn_name = target_info.info['fn_free_buffer']
free_buffer_fnty = llvm_ir.FunctionType(void_t, [int8_t.as_pointer()])
free_buffer_fn = irutils.get_or_insert_function(
builder.module, free_buffer_fnty, free_buffer_fn_name)
[buf] = args
buf_ptr = builder.load(builder.gep(
buf, [int32_t(0), int32_t(0)])) # buf.ptr
buf_ptr = builder.bitcast(buf_ptr, int8_t.as_pointer())
builder.call(free_buffer_fn, [buf_ptr])
def omnisci_buffer_constructor(context, builder, sig, args):
"""
Usage:
extending.lower_builtin(MyBuffer, numba.types.Integer, ...)(omnisci_buffer_constructor)
will enable creating MyBuffer instance from a Omnisci UDF/UDTF definition:
b = MyBuffer(<size>, ...)
"""
target_info = TargetInfo()
try:
alloc_fn_name = target_info.info['fn_allocate_varlen_buffer']
except KeyError as msg:
raise UnsupportedError(f'{target_info} does not provide {msg}')
ptr_type, sz_type = sig.return_type.dtype.members[:2]
if len(sig.return_type.dtype.members) > 2:
assert len(sig.return_type.dtype.members) == 3
null_type = sig.return_type.dtype.members[2]
else:
null_type = None
assert isinstance(args[0].type, ir.IntType), (args[0].type)
element_count = builder.zext(args[0], int64_t)
element_size = int64_t(ptr_type.dtype.bitwidth // 8)
alloc_fnty = ir.FunctionType(int8_t.as_pointer(), [int64_t, int64_t])
alloc_fn = irutils.get_or_insert_function(builder.module, alloc_fnty,
alloc_fn_name)
ptr8 = builder.call(alloc_fn, [element_count, element_size])
# remember possible temporary allocations so that when leaving a
# UDF/UDTF, these will be deallocated, see pipeline.py.
builder_buffers[builder].append(ptr8)
ptr = builder.bitcast(ptr8, context.get_value_type(ptr_type))
fa = cgutils.create_struct_proxy(sig.return_type.dtype)(context, builder)
fa.ptr = ptr # T*
fa.sz = element_count # size_t
if null_type is not None:
is_zero = builder.icmp_signed('==', element_count, int64_t(0))
with builder.if_else(is_zero) as (then, orelse):
with then:
is_null = context.get_value_type(null_type)(1)
with orelse:
is_null = context.get_value_type(null_type)(0)
fa.is_null = is_null # int8_t
return fa._getpointer()
def omnisci_array_is_null_(typingctx, T, elem):
sig = types.boolean(T, elem)
target_info = TargetInfo()
null_value = target_info.null_values[f'{T.dtype}']
# The server sends numbers as unsigned values rather than signed ones.
# Thus, 129 should be read as -127 (overflow). See rbc issue #254
bitwidth = T.dtype.bitwidth
null_value = np.dtype(f'uint{bitwidth}').type(null_value).view(f'int{bitwidth}')
nv = ir.Constant(ir.IntType(bitwidth), null_value)
def codegen(context, builder, signature, args):
_, elem = args
if isinstance(T.dtype, types.Float):
elem = builder.bitcast(elem, nv.type)
return builder.icmp_signed('==', elem, nv)
return sig, codegen
def omnisci_udtfmanager_set_output_row_size_(typingctx, mgr, num_rows):
sig = types.void(mgr, num_rows)
target_info = TargetInfo()
if target_info.software[1][:3] < (5, 9, 0):
raise UnsupportedError(error_msg % (".".join(map(str, target_info.software[1]))))
def codegen(context, builder, sig, args):
mgr_ptr, num_rows_arg = args
mgr_i8ptr = builder.bitcast(mgr_ptr, i8p)
fnty = ir.FunctionType(ir.VoidType(), [i8p, i64])
fn = irutils.get_or_insert_function(
builder.module, fnty, "TableFunctionManager_set_output_row_size")
builder.call(fn, [mgr_i8ptr, num_rows_arg])
return sig, codegen
def omnisci_column_is_null_(typingctx, col_var, row_idx):
T = col_var.eltype
sig = types.boolean(col_var, row_idx)
target_info = TargetInfo()
null_value = target_info.null_values[str(T)]
nv = ir.Constant(ir.IntType(T.bitwidth), null_value)
def codegen(context, builder, signature, args):
data, index = args
ptr = irutils.get_member_value(builder, data, 0)
res = builder.load(builder.gep(ptr, [index]))
if isinstance(T, types.Float):
res = builder.bitcast(res, nv.type)
return builder.icmp_signed('==', res, nv)
return sig, codegen
def match_signature(self, atypes):
# Code here is the same found in remotejit.py::Signature::best_match
device = "CPU" if TargetInfo().is_cpu else "GPU"
available_types = tuple(map(Type.fromobject, self._signatures[device]))
ftype = None
match_penalty = None
for typ in available_types:
penalty = typ.match(atypes)
if penalty is not None:
if ftype is None or penalty < match_penalty:
ftype = typ
match_penalty = penalty
if ftype is None:
satypes = ", ".join(map(str, atypes))
available = "; ".join(map(str, available_types))
raise TypeError(
f"found no matching function type to given argument types"
f" `{satypes}` and device `{device}`. Available function types: {available}"
)
return ftype
def test_dummy():
with pytest.raises(
RuntimeError,
match=r'Target not specified.'):
Type.fromobject('int')
with TargetInfo.dummy():
t = Type.fromobject('int')
assert str(t) == 'int'
t = Type.fromobject('int(int)')
assert t.is_function
t = Type.fromobject('foo(bar)')
assert t.is_function
with pytest.raises(
RuntimeError,
match=r'Target not specified.'):
Type.fromobject('foo(bar)')
def omnisci_buffer_idx_is_null_(typingctx, col_var, row_idx):
T = col_var.eltype
sig = types.boolean(col_var, row_idx)
target_info = TargetInfo()
null_value = target_info.null_values[str(T)]
# The server sends numbers as unsigned values rather than signed ones.
# Thus, 129 should be read as -127 (overflow). See rbc issue #254
nv = ir.Constant(ir.IntType(T.bitwidth), null_value)
def codegen(context, builder, signature, args):
ptr, index = args
data = builder.extract_value(builder.load(ptr), [0])
res = builder.load(builder.gep(data, [index]))
if isinstance(T, types.Float):
res = builder.bitcast(res, nv.type)
return builder.icmp_signed('==', res, nv)
return sig, codegen
def external(cls, *args):
"""
Parameters
----------
signature : object (str, ctypes function, python callable, numba function)
Any object convertible to a Numba function via Type.fromobject(...).tonumba()
"""
ts = defaultdict(list)
key = None
for signature in args:
with TargetInfo.dummy():
t = Type.fromobject(signature)
if not t.is_function:
raise ValueError("signature must represent a function type")
if not t.name:
raise ValueError(
f"external function name not specified for signature {signature}"
)
if key is None:
key = t.name
if not key:
raise ValueError(
f"external function name not specified for signature {signature}"
)
for device in [
a for a in t.annotation() or [] if a in ["CPU", "GPU"]
] or [
"CPU",
"GPU",
]:
ts[device].append(signature)
obj = cls(key, ts)
obj.register()
return obj
def omnisci_column_is_null_(typingctx, col_var, row_idx):
T = col_var.eltype
sig = types.boolean(col_var, row_idx)
target_info = TargetInfo()
null_value = target_info.null_values[str(T)]
nv = ir.Constant(ir.IntType(T.bitwidth), null_value)
def codegen(context, builder, signature, args):
zero = int32_t(0)
data, index = args
assert data.opname == 'load'
buf = data.operands[0]
ptr = builder.load(builder.gep(buf, [zero, zero]))
res = builder.load(builder.gep(ptr, [index]))
if isinstance(T, types.Float):
res = builder.bitcast(res, nv.type)
return builder.icmp_signed('==', res, nv)
return sig, codegen
def fromobject(cls, signature, name: str = None):
"""
Parameters
----------
signature : object (str, ctypes function, python callable, numba function)
Any object convertible to a Numba function via Type.fromobject(...).tonumba()
name : str
The name of the external function
"""
# Make inner function for the actual work
target_info = TargetInfo.dummy()
with target_info:
t = Type.fromobject(signature)
if not t.is_function:
raise ValueError("signature must represent a function type")
if name is None:
name = t.name
if not name:
raise ValueError(
f"external function name not specified for signature {signature}"
)
return cls(name, t)
np_NA_message = None
except ImportError as msg:
np = None
np_NA_message = str(msg)
try:
import numpy as np
except ImportError:
np = None
import pytest
from rbc.typesystem import Type, get_signature
from rbc.utils import get_datamodel
from rbc.targetinfo import TargetInfo
target_info = TargetInfo.host()
def Type_fromstring(s):
return Type.fromstring(s, target_info)
def Type_fromobject(s):
return Type.fromobject(s, target_info)
def Type_fromcallable(s):
return Type.fromcallable(s, target_info)
def Type_fromvalue(s):
def inner(context, builder, sig, args):
impl = cpu if TargetInfo().is_cpu else gpu
return impl(context, builder, sig, args)
def codegen(context, builder, signature, args):
target_info = TargetInfo()
if target_info.is_cpu:
cgutils.printf(builder, format_type.literal_value, *args[1:])
_cg_fflush(builder)
def codegen(context, builder, signature, args):
target_info = TargetInfo()
if target_info.is_cpu:
_cg_fflush(builder)
def target_info():
target_info = TargetInfo.host()
with target_info:
yield target_info
def test_to_from_dict():
host = TargetInfo.host()
d = host.todict()
host2 = TargetInfo.fromdict(d)
assert host.__dict__ == host2.__dict__
def test_host_cache():
host = TargetInfo.host()
host2 = TargetInfo.host()
assert host is host2
def test_basic():
host = TargetInfo.host()
assert host.name == 'host_cpu'
with host:
ti = TargetInfo()
assert ti is host
def pass_by_value(self):
omnisci_version = TargetInfo().software[1][:3]
return omnisci_version <= (5, 7, 0)