def test_ckernel_deferred_from_pyfunc(self):
# Test wrapping make_assignment_ckernel as a deferred ckernel
def instantiate_assignment(out_ckb, ckb_offset, types, meta,
kerntype, ectx):
out_ckb = _lowlevel.CKernelBuilderStruct.from_address(out_ckb)
return _lowlevel.make_assignment_ckernel(out_ckb, ckb_offset,
types[0], meta[0],
types[1], meta[1],
'expr', kerntype, ectx)
ckd = _lowlevel.ckernel_deferred_from_pyfunc(instantiate_assignment,
[ndt.string, ndt.date])
self.assertEqual(nd.as_py(ckd.types), [ndt.string, ndt.date])
out = nd.empty(ndt.string)
in0 = nd.array('2012-11-05', ndt.date)
ckd.__call__(out, in0)
self.assertEqual(nd.as_py(out), '2012-11-05')
# Also test it as a lifted kernel
ckd_lifted = _lowlevel.lift_ckernel_deferred(ckd,
['3 * var * string', '3 * var * date'])
self.assertEqual(nd.as_py(ckd_lifted.types),
[ndt.type('3 * var * string'), ndt.type('3 * var * date')])
out = nd.empty('3 * var * string')
from datetime import date
in0 = nd.array([['2013-03-11', date(2010, 10, 10)],
[date(1999, 12, 31)],
[]], type='3 * var * date')
ckd_lifted.__call__(out, in0)
self.assertEqual(nd.as_py(out),
[['2013-03-11', '2010-10-10'],
['1999-12-31'], []])
def test_ctypes_callback_deferred(self):
# Create a deferred ckernel via a closure
def instantiate_ckernel(out_ckb, ckb_offset, types, meta, kerntype,
ectx):
out_ckb = _lowlevel.CKernelBuilder(out_ckb)
def my_kernel_func_single(dst_ptr, src_ptr, kdp):
dst = ctypes.c_int32.from_address(dst_ptr)
src = ctypes.c_double.from_address(src_ptr[0])
dst.value = int(src.value * 3.5)
def my_kernel_func_strided(dst_ptr, dst_stride, src_ptr,
src_stride, count, kdp):
src_ptr0 = src_ptr[0]
src_stride0 = src_stride[0]
for i in range(count):
my_kernel_func_single(dst_ptr, [src_ptr0], kdp)
dst_ptr += dst_stride
src_ptr0 += src_stride0
if kerntype == 'single':
kfunc = _lowlevel.ExprSingleOperation(my_kernel_func_single)
else:
kfunc = _lowlevel.ExprStridedOperation(my_kernel_func_strided)
return ckernel.wrap_ckernel_func(out_ckb, ckb_offset, kfunc, kfunc)
ckd = _lowlevel.ckernel_deferred_from_pyfunc(instantiate_ckernel,
[ndt.int32, ndt.float64])
# Test calling the ckd
out = nd.empty(ndt.int32)
in0 = nd.array(4.0, type=ndt.float64)
ckd.__call__(out, in0)
self.assertEqual(nd.as_py(out), 14)
# Also call it lifted
ckd_lifted = _lowlevel.lift_ckernel_deferred(
ckd, ['2 * var * int32', '2 * var * float64'])
out = nd.empty('2 * var * int32')
in0 = nd.array([[1.0, 3.0, 2.5], [1.25, -1.5]],
type='2 * var * float64')
ckd_lifted.__call__(out, in0)
self.assertEqual(nd.as_py(out), [[3, 10, 8], [4, -5]])
def test_ctypes_callback_deferred(self):
# Create a deferred ckernel via a closure
def instantiate_ckernel(out_ckb, ckb_offset, types, meta,
kerntype, ectx):
out_ckb = _lowlevel.CKernelBuilder(out_ckb)
def my_kernel_func_single(dst_ptr, src_ptr, kdp):
dst = ctypes.c_int32.from_address(dst_ptr)
src = ctypes.c_double.from_address(src_ptr[0])
dst.value = int(src.value * 3.5)
def my_kernel_func_strided(dst_ptr, dst_stride, src_ptr, src_stride, count, kdp):
src_ptr0 = src_ptr[0]
src_stride0 = src_stride[0]
for i in range(count):
my_kernel_func_single(dst_ptr, [src_ptr0], kdp)
dst_ptr += dst_stride
src_ptr0 += src_stride0
if kerntype == 'single':
kfunc = _lowlevel.ExprSingleOperation(my_kernel_func_single)
else:
kfunc = _lowlevel.ExprStridedOperation(my_kernel_func_strided)
return ckernel.wrap_ckernel_func(out_ckb, ckb_offset,
kfunc, kfunc)
ckd = _lowlevel.ckernel_deferred_from_pyfunc(instantiate_ckernel,
[ndt.int32, ndt.float64])
# Test calling the ckd
out = nd.empty(ndt.int32)
in0 = nd.array(4.0, type=ndt.float64)
ckd.__call__(out, in0)
self.assertEqual(nd.as_py(out), 14)
# Also call it lifted
ckd_lifted = _lowlevel.lift_ckernel_deferred(ckd,
['2 * var * int32', '2 * var * float64'])
out = nd.empty('2 * var * int32')
in0 = nd.array([[1.0, 3.0, 2.5], [1.25, -1.5]], type='2 * var * float64')
ckd_lifted.__call__(out, in0)
self.assertEqual(nd.as_py(out), [[3, 10, 8], [4, -5]])
def jit_compile_ckernel_deferred(bek, out_dshape):
"""
Creates a ckernel_deferred from the blaze element kernel.
Actual JIT compilation is done at instantiation.
Parameters
----------
bek : BlazeElementKernel
The blaze kernel.
"""
# Create a deferred ckernel via a closure
def instantiate_ckernel(out_ckb, ckb_offset, types, meta, kerntype):
out_ckb = _lowlevel.CKernelBuilder(out_ckb)
strided = (kerntype == 'strided')
# TODO cache the compiled kernels
module, lfunc = create_ckernel_interface(bek, strided)
optimize(module, lfunc)
ee, func_ptr = get_pointer(module, lfunc)
# TODO: Something like the following needs to be reenabled
# to handle array types
# Build llvm and ctypes structures for the kernel data, using
# the argument types.
##kd_llvmtype, kd_ctypestype = args_to_kernel_data_struct(bek.kinds, bek.argtypes)
# Cast the extra pointer to the right llvm type
#extra_struct = builder.bitcast(extra_ptr_arg,
# Type.pointer(kd_llvmtype))
# Create a function which copies the shape from data
# descriptors to the extra data struct.
##if len(kd_ctypestype._fields_) == 1:
## # If there were no extra data fields, it's a no-op function
## def bind_func(estruct, dst_dd, src_dd_list):
## pass
##else:
## def bind_func(estruct, dst_dd, src_dd_list):
## for i, (ds, dd) in enumerate(
## izip(bek.dshapes, src_dd_list + [dst_dd])):
## shape = [operator.index(dim)
## for dim in dd.dshape[-len(ds):-1]]
## cshape = getattr(estruct, 'operand_%d' % i)
## for j, dim_size in enumerate(shape):
## cshape[j] = dim_size
if strided:
optype = _lowlevel.ExprStridedOperation
else:
optype = _lowlevel.ExprSingleOperation
return wrap_ckernel_func(out_ckb, ckb_offset, optype(func_ptr),
(ee, func_ptr))
# Wrap the function in a ckernel_deferred
in_dshapes = list(bek.dshapes)
# HACK: sometimes the return type is there, sometimes not,
# exclude it unconditionally.
in_dshapes = in_dshapes[:len(bek.kinds)-1]
out_type = out_dshape.measure
in_types = [in_dshape.measure for in_dshape in in_dshapes]
return _lowlevel.ckernel_deferred_from_pyfunc(instantiate_ckernel,
[ndt.type(str(t)) for t in [out_type] + in_types])
def jit_compile_ckernel_deferred(bek, out_dshape):
"""
Creates a ckernel_deferred from the blaze element kernel.
Actual JIT compilation is done at instantiation.
Parameters
----------
bek : BlazeElementKernel
The blaze kernel.
"""
# Create a deferred ckernel via a closure
def instantiate_ckernel(out_ckb, ckb_offset, types, meta, kerntype):
out_ckb = _lowlevel.CKernelBuilder(out_ckb)
strided = (kerntype == 'strided')
# TODO cache the compiled kernels
module, lfunc = create_ckernel_interface(bek, strided)
optimize(module, lfunc)
ee, func_ptr = get_pointer(module, lfunc)
# TODO: Something like the following needs to be reenabled
# to handle array types
# Build llvm and ctypes structures for the kernel data, using
# the argument types.
##kd_llvmtype, kd_ctypestype = args_to_kernel_data_struct(bek.kinds, bek.argtypes)
# Cast the extra pointer to the right llvm type
#extra_struct = builder.bitcast(extra_ptr_arg,
# Type.pointer(kd_llvmtype))
# Create a function which copies the shape from data
# descriptors to the extra data struct.
##if len(kd_ctypestype._fields_) == 1:
## # If there were no extra data fields, it's a no-op function
## def bind_func(estruct, dst_dd, src_dd_list):
## pass
##else:
## def bind_func(estruct, dst_dd, src_dd_list):
## for i, (ds, dd) in enumerate(
## izip(bek.dshapes, src_dd_list + [dst_dd])):
## shape = [operator.index(dim)
## for dim in dd.dshape[-len(ds):-1]]
## cshape = getattr(estruct, 'operand_%d' % i)
## for j, dim_size in enumerate(shape):
## cshape[j] = dim_size
if strided:
optype = _lowlevel.ExprStridedOperation
else:
optype = _lowlevel.ExprSingleOperation
return wrap_ckernel_func(out_ckb, ckb_offset, optype(func_ptr),
(ee, func_ptr))
# Wrap the function in a ckernel_deferred
in_dshapes = list(bek.dshapes)
# HACK: sometimes the return type is there, sometimes not,
# exclude it unconditionally.
in_dshapes = in_dshapes[:len(bek.kinds)-1]
out_type = out_dshape.measure
in_types = [in_dshape.measure for in_dshape in in_dshapes]
return _lowlevel.ckernel_deferred_from_pyfunc(instantiate_ckernel,
[ndt.type(str(t)) for t in [out_type] + in_types])
