def test_ptr(omnisci, c_name, device):
omnisci.reset()
if not omnisci.has_cuda and device == 'gpu':
pytest.skip('test requires CUDA-enabled omniscidb server')
from rbc.external import external
if omnisci.compiler is None:
pytest.skip('test requires clang C/C++ compiler')
ctype, cname = c_name.split()
c_code = f'''
#include <stdint.h>
#ifdef __cplusplus
extern "C" {{
#endif
{ctype} mysum_impl({ctype}* x, int n) {{
{ctype} r = 0;
for (int i=0; i < n; i++) {{
r += x[i];
}}
return r;
}}
{ctype} myval_impl({ctype}* x) {{
return *x;
}}
#ifdef __cplusplus
}}
#endif
'''
omnisci.user_defined_llvm_ir[device] = omnisci.compiler(c_code)
mysum_impl = external(f'{ctype} mysum_impl({ctype}*, int32_t)')
myval_impl = external(f'{ctype} myval_impl({ctype}*)')
@omnisci(f'{ctype}({ctype}[])', devices=[device])
def mysum_ptr(x):
return mysum_impl(x.ptr(), len(x))
@omnisci(f'{ctype}({ctype}[], int32_t)', devices=[device])
def myval_ptr(x, i):
return myval_impl(x.ptr(i))
desrc, result = omnisci.sql_execute(
f'select {cname}, mysum_ptr({cname}) from {omnisci.table_name}')
for a, r in result:
if cname == 'i1':
assert sum(a) % 256 == r % 256
else:
assert sum(a) == r
desrc, result = omnisci.sql_execute(
f'select {cname}, myval_ptr({cname}, 0), myval_ptr({cname}, 2) from {omnisci.table_name}'
)
for a, r0, r2 in result:
assert a[0] == r0
assert a[2] == r2
def test_unnamed_external(heavydb):
with pytest.raises(ValueError) as excinfo:
external("f64(f64)")
assert "external function name not specified for signature" in str(excinfo)
with pytest.raises(ValueError) as excinfo:
external("f64(f64)")
assert "external function name not specified for signature" in str(excinfo)
def test_replace_declaration(omnisci):
_ = external("f64(f64)", name="fma")
fma = external("f64(f64, f64, f64)", name="fma")
@omnisci("double(double, double, double)")
def test_fma(a, b, c):
return fma(a, b, c)
_, result = omnisci.sql_execute("select test_fma(1.0, 2.0, 3.0)")
assert list(result) == [(5.0, )]
def test_replace_declaration(heavydb):
_ = external("f64 fma(f64)|CPU")
fma = external("f64 fma(f64, f64, f64)|CPU")
@heavydb("double(double, double, double)", devices=["cpu"])
def test_fma(a, b, c):
return fma(a, b, c)
_, result = heavydb.sql_execute("select test_fma(1.0, 2.0, 3.0)")
assert list(result) == [(5.0, )]
def inner(fname, signature):
cmath_fn = external(signature, name=fname)
t = Type.fromstring(signature)
retty = str(t[0])
argtypes = tuple(map(str, t[1]))
arity = len(argtypes)
# define omnisci callable
if arity == 1:
def fn(a):
return cmath_fn(a)
elif arity == 2:
def fn(a, b):
return cmath_fn(a, b)
else:
def fn(a, b, c):
return cmath_fn(a, b, c)
fn.__name__ = f"{omnisci.table_name}_{fname}"
fn = omnisci(f"{retty}({', '.join(argtypes)})")(fn)
def test_scalar_pointer_access_remote(rjit, memman, T):
rjit.reset()
calloc_prototype, free_prototype = memory_managers[memman]
with Type.alias(T=T):
arr = np.array([1, 2, 3, 4], dtype=T)
arr2 = np.array([11, 12, 13, 14], dtype=T)
calloc = external(calloc_prototype)
free = external(free_prototype)
@rjit(calloc_prototype)
def rcalloc(nmemb, size):
return calloc(nmemb, size)
@rjit(free_prototype)
def rfree(ptr):
return free(ptr)
@rjit('T(T* x, int i)')
def pitem(x, i):
return x[i]
@rjit('void(T* x, int i, T)')
def sitem(x, i, v):
x[i] = v
ptr = rcalloc(arr.itemsize, len(arr))
assert ptr.value
for i in range(len(arr)):
sitem(ptr, i, arr[i])
for i in range(len(arr2)):
arr2[i] = pitem(ptr, i)
rfree(ptr)
assert (arr == arr2).all()
def test_require_target_info(omnisci):
log2 = external("double log2(double)", name="log2")
@omnisci("double(double)")
def test_log2(a):
return log2(a)
_, result = omnisci.sql_execute("select test_log2(8.0)")
assert list(result) == [(3.0, )]
def test_require_target_info(heavydb):
log2 = external("double log2(double)|CPU")
@heavydb("double(double)", devices=["cpu"])
def test_log2(a):
return log2(a)
_, result = heavydb.sql_execute("select test_log2(8.0)")
assert list(result) == [(3.0, )]
def test_boston_house_prices(heavydb, boston_house_prices):
if heavydb.compiler is None:
pytest.skip('test requires C/C++ to LLVM IR compiler')
treelite = pytest.importorskip("treelite")
xgboost = pytest.importorskip("xgboost")
device = 'cpu'
import numpy as np
import tempfile
# Get training data from server:
table_name = f'{heavydb.table_name}bhp'
descr, result = heavydb.sql_execute(
f'SELECT rowid, data, medv FROM {table_name} ORDER BY rowid LIMIT 50')
result = list(result)
medv = np.array([medv for _, data, medv in result])
data = np.array([data for _, data, medv in result])
assert len(medv) == 50
# Train model using xgboost
dtrain = xgboost.DMatrix(data, label=medv)
params = {
'max_depth': 3,
'eta': 1,
'objective': 'reg:squarederror',
'eval_metric': 'rmse'
}
bst = xgboost.train(params, dtrain, len(medv), [(dtrain, 'train')])
# Compile model to C
working_dir = tempfile.mkdtemp()
model = treelite.Model.from_xgboost(bst)
model.compile(working_dir)
model_c = open(os.path.join(working_dir, 'main.c')).read()
# The C model implements
# float predict(union Entry* data, int pred_margin)
# but we wrap it using
model_c += '''
#ifdef __cplusplus
extern "C" {
#endif
float predict_float(float* data, int pred_margin) {
return predict((union Entry*)data, pred_margin);
}
#ifdef __cplusplus
}
#endif
'''
predict_float = external('float predict_float(float*, int32)')
# to make UDF construction easier. Notice that predict_float can
# be now called from a UDF.
# Define predict function as UDF. Notice that the xgboost null
# values are different from heavydb null values, so we'll remap
# before calling predict_float:
null_value = np.int32(-1).view(np.float32)
@heavydb('float(float[], int32)', devices=[device])
def predict(data, pred_margin):
for i in range(len(data)):
if data.is_null(i):
data[i] = null_value
return predict_float(data.ptr(), pred_margin)
# Compile C model to LLVM IR. In future, we might want this
# compilation to happen in the server side as the client might not
# have clang compiler installed.
model_llvmir = heavydb.compiler(model_c, flags=['-I' + working_dir])
# RBC will link_in the LLVM IR module
heavydb.user_defined_llvm_ir[device] = model_llvmir
# Call predict on data in the server:
descr, result = heavydb.sql_execute('SELECT rowid, predict(data, 2) FROM'
f' {table_name} ORDER BY rowid')
result = list(result)
predict_medv = np.array([r[1] for r in result])
# predict on the first 50 elements should be close to training labels
error = abs(predict_medv[:len(medv)] - medv).max() / len(medv)
assert error < 1e-4, error
def test_invalid_signature(omnisci):
with pytest.raises(ValueError) as excinfo:
external(types.int64, name="test")
assert "signature must represent a function type" in str(excinfo)
def test_valid_signatures(omnisci):
assert external("f64 log2(f64)").name == "log2"
assert external("f64(f64)", name="log2").name == "log2"
def test_invalid_signature(heavydb):
with pytest.raises(ValueError) as excinfo:
external(types.int64)
assert "signature must represent a function type" in str(excinfo)
def test_valid_signatures(heavydb):
external("f64 log2(f64)")
assert True
def test_valid_signatures(omnisci):
external("f64 log2(f64)")
assert True
