def get_attr_impl(context, builder, typ, value, attr):
"""
Generic getattr() for @jitclass instances.
"""
if attr in typ.struct:
# It's a struct field
inst = context.make_helper(builder, typ, value=value)
data_pointer = inst.data
data = context.make_data_helper(builder, typ.get_data_type(),
ref=data_pointer)
return imputils.impl_ret_borrowed(context, builder,
typ.struct[attr],
getattr(data, _mangle_attr(attr)))
elif attr in typ.jitprops:
# It's a jitted property
getter = typ.jitprops[attr]['get']
sig = templates.signature(None, typ)
dispatcher = types.Dispatcher(getter)
sig = dispatcher.get_call_type(context.typing_context, [typ], {})
call = context.get_function(dispatcher, sig)
out = call(builder, [value])
_add_linking_libs(context, call)
return imputils.impl_ret_new_ref(context, builder, sig.return_type, out)
raise NotImplementedError('attribute {0!r} not implemented'.format(attr))
def codegen(context, builder, signature, args):
# check that the return type is now defined
arrty = signature.return_type
assert arrty.is_precise()
shapes = unpack_tuple(builder, args[0])
# redirect implementation to np.empty
res = _empty_nd_impl(context, builder, arrty, shapes)
return impl_ret_new_ref(context, builder, arrty, res._getvalue())
def array_sinc(context, builder, sig, args):
def array_sinc_impl(arr):
out = numpy.zeros_like(arr)
for index, val in numpy.ndenumerate(arr):
out[index] = numpy.sinc(val)
return out
res = context.compile_internal(builder, array_sinc_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def imp(context, builder, sig, args):
instance_type = sig.args[0]
method = instance_type.jitmethods[attr]
disp_type = types.Dispatcher(method)
call = context.get_function(disp_type, sig)
out = call(builder, args)
return imputils.impl_ret_new_ref(context, builder,
sig.return_type, out)
def array_sinc(context, builder, sig, args):
def array_sinc_impl(arr):
out = np.zeros_like(arr)
for index, val in np.ndenumerate(arr):
out[index] = np.sinc(val)
return out
res = context.compile_internal(builder, array_sinc_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def imp(context, builder, sig, args):
instance_type = sig.args[0]
method = instance_type.jitmethods[attr]
disp_type = types.Dispatcher(method)
call = context.get_function(disp_type, sig)
out = call(builder, args)
_add_linking_libs(context, call)
return imputils.impl_ret_new_ref(context, builder,
sig.return_type, out)
def array_nonzero(context, builder, sig, args):
aryty = sig.args[0]
# Return type is a N-tuple of 1D C-contiguous arrays
retty = sig.return_type
outaryty = retty.dtype
ndim = aryty.ndim
nouts = retty.count
ary = make_array(aryty)(context, builder, args[0])
shape = cgutils.unpack_tuple(builder, ary.shape)
strides = cgutils.unpack_tuple(builder, ary.strides)
data = ary.data
layout = aryty.layout
# First count the number of non-zero elements
zero = context.get_constant(types.intp, 0)
one = context.get_constant(types.intp, 1)
count = cgutils.alloca_once_value(builder, zero)
with cgutils.loop_nest(builder, shape, zero.type) as indices:
ptr = cgutils.get_item_pointer2(builder, data, shape, strides, layout,
indices)
val = load_item(context, builder, aryty, ptr)
nz = context.is_true(builder, aryty.dtype, val)
with builder.if_then(nz):
builder.store(builder.add(builder.load(count), one), count)
# Then allocate output arrays of the right size
out_shape = (builder.load(count), )
outs = [
_empty_nd_impl(context, builder, outaryty, out_shape)._getvalue()
for i in range(nouts)
]
outarys = [make_array(outaryty)(context, builder, out) for out in outs]
out_datas = [out.data for out in outarys]
# And fill them up
index = cgutils.alloca_once_value(builder, zero)
with cgutils.loop_nest(builder, shape, zero.type) as indices:
ptr = cgutils.get_item_pointer2(builder, data, shape, strides, layout,
indices)
val = load_item(context, builder, aryty, ptr)
nz = context.is_true(builder, aryty.dtype, val)
with builder.if_then(nz):
# Store element indices in output arrays
if not indices:
# For a 0-d array, store 0 in the unique output array
indices = (zero, )
cur = builder.load(index)
for i in range(nouts):
ptr = cgutils.get_item_pointer2(builder, out_datas[i],
out_shape, (), 'C', [cur])
store_item(context, builder, outaryty, indices[i], ptr)
builder.store(builder.add(cur, one), index)
tup = context.make_tuple(builder, sig.return_type, outs)
return impl_ret_new_ref(context, builder, sig.return_type, tup)
def array_round(context, builder, sig, args):
def array_round_impl(arr, decimals, out):
if arr.shape != out.shape:
raise ValueError("invalid output shape")
for index, val in numpy.ndenumerate(arr):
out[index] = numpy.round(val, decimals)
return out
res = context.compile_internal(builder, array_round_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_round(context, builder, sig, args):
def array_round_impl(arr, decimals, out):
if arr.shape != out.shape:
raise ValueError("invalid output shape")
for index, val in np.ndenumerate(arr):
out[index] = np.round(val, decimals)
return out
res = context.compile_internal(builder, array_round_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_nonzero(context, builder, sig, args):
aryty = sig.args[0]
# Return type is a N-tuple of 1D C-contiguous arrays
retty = sig.return_type
outaryty = retty.dtype
ndim = aryty.ndim
nouts = retty.count
ary = make_array(aryty)(context, builder, args[0])
shape = cgutils.unpack_tuple(builder, ary.shape)
strides = cgutils.unpack_tuple(builder, ary.strides)
data = ary.data
layout = aryty.layout
# First count the number of non-zero elements
zero = context.get_constant(types.intp, 0)
one = context.get_constant(types.intp, 1)
count = cgutils.alloca_once_value(builder, zero)
with cgutils.loop_nest(builder, shape, zero.type) as indices:
ptr = cgutils.get_item_pointer2(builder, data, shape, strides,
layout, indices)
val = load_item(context, builder, aryty, ptr)
nz = context.is_true(builder, aryty.dtype, val)
with builder.if_then(nz):
builder.store(builder.add(builder.load(count), one), count)
# Then allocate output arrays of the right size
out_shape = (builder.load(count),)
outs = [_empty_nd_impl(context, builder, outaryty, out_shape)._getvalue()
for i in range(nouts)]
outarys = [make_array(outaryty)(context, builder, out) for out in outs]
out_datas = [out.data for out in outarys]
# And fill them up
index = cgutils.alloca_once_value(builder, zero)
with cgutils.loop_nest(builder, shape, zero.type) as indices:
ptr = cgutils.get_item_pointer2(builder, data, shape, strides,
layout, indices)
val = load_item(context, builder, aryty, ptr)
nz = context.is_true(builder, aryty.dtype, val)
with builder.if_then(nz):
# Store element indices in output arrays
if not indices:
# For a 0-d array, store 0 in the unique output array
indices = (zero,)
cur = builder.load(index)
for i in range(nouts):
ptr = cgutils.get_item_pointer2(builder, out_datas[i],
out_shape, (),
'C', [cur])
store_item(context, builder, outaryty, indices[i], ptr)
builder.store(builder.add(cur, one), index)
tup = context.make_tuple(builder, sig.return_type, outs)
return impl_ret_new_ref(context, builder, sig.return_type, tup)
def dot_2_vm(context, builder, sig, args):
"""
np.dot(vector, matrix)
"""
def dot_impl(a, b):
m, = a.shape
_m, n = b.shape
out = numpy.empty((n, ), a.dtype)
return numpy.dot(a, b, out)
res = context.compile_internal(builder, dot_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_pop(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
n = inst.size
cgutils.guard_zero(context, builder, n,
(IndexError, "pop from empty list"))
n = builder.sub(n, ir.Constant(n.type, 1))
res = inst.getitem(n)
inst.incref_value(res) # incref the pop'ed element
inst.clear_value(n) # clear the storage space
inst.resize(n)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def dot_2_vm(context, builder, sig, args):
"""
np.dot(vector, matrix)
"""
def dot_impl(a, b):
m, = a.shape
_m, n = b.shape
out = numpy.empty((n, ), a.dtype)
return numpy.dot(a, b, out)
res = context.compile_internal(builder, dot_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def dot_2_mv(context, builder, sig, args):
"""
np.dot(matrix, vector)
"""
def dot_impl(a, b):
m, n = a.shape
_n, = b.shape
out = np.empty((m, ), a.dtype)
return np.dot(a, b, out)
res = context.compile_internal(builder, dot_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_pop(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
n = inst.size
cgutils.guard_zero(context, builder, n,
(IndexError, "pop from empty list"))
n = builder.sub(n, ir.Constant(n.type, 1))
res = inst.getitem(n)
inst.incref_value(res) # incref the pop'ed element
inst.clear_value(n) # clear the storage space
inst.resize(n)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def build_list(context, builder, list_type, items):
"""
Build a list of the given type, containing the given items.
"""
nitems = len(items)
inst = ListInstance.allocate(context, builder, list_type, nitems)
# Populate list
inst.size = context.get_constant(types.intp, nitems)
for i, val in enumerate(items):
inst.setitem(context.get_constant(types.intp, i), val)
return impl_ret_new_ref(context, builder, list_type, inst.value)
def make_zip_object(context, builder, sig, args):
zip_type = sig.return_type
assert len(args) == len(zip_type.source_types)
zipobj = context.make_helper(builder, zip_type)
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
zipobj[i] = call_getiter(context, builder, srcty, arg)
res = zipobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
def build_list(context, builder, list_type, items):
"""
Build a list of the given type, containing the given items.
"""
nitems = len(items)
inst = ListInstance.allocate(context, builder, list_type, nitems)
# Populate list
inst.size = context.get_constant(types.intp, nitems)
for i, val in enumerate(items):
inst.setitem(context.get_constant(types.intp, i), val)
return impl_ret_new_ref(context, builder, list_type, inst.value)
def make_zip_object(context, builder, sig, args):
zip_type = sig.return_type
assert len(args) == len(zip_type.source_types)
zipobj = context.make_helper(builder, zip_type)
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
zipobj[i] = call_getiter(context, builder, srcty, arg)
res = zipobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
def set_constructor(context, builder, sig, args):
set_type = sig.return_type
items_type, = sig.args
items, = args
# If the argument has a len(), preallocate the set so as to
# avoid resizes.
n = call_len(context, builder, items_type, items)
inst = SetInstance.allocate(context, builder, set_type, n)
with for_iter(context, builder, items_type, items) as loop:
inst.add(loop.value)
return impl_ret_new_ref(context, builder, set_type, inst.value)
def array_angle_kwarg(context, builder, sig, args):
arg = sig.args[0]
if isinstance(arg.dtype, types.Complex):
retty = arg.dtype.underlying_float
else:
retty = arg.dtype
def array_angle_impl(arr, deg=False):
out = numpy.zeros_like(arr, dtype=retty)
for index, val in numpy.ndenumerate(arr):
out[index] = numpy.angle(val, deg)
return out
res = context.compile_internal(builder, array_angle_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def set_constructor(context, builder, sig, args):
set_type = sig.return_type
items_type, = sig.args
items, = args
# If the argument has a len(), preallocate the set so as to
# avoid resizes.
n = call_len(context, builder, items_type, items)
inst = SetInstance.allocate(context, builder, set_type, n)
with for_iter(context, builder, items_type, items) as loop:
inst.add(loop.value)
return impl_ret_new_ref(context, builder, set_type, inst.value)
def array_angle_kwarg(context, builder, sig, args):
arg = sig.args[0]
if isinstance(arg.dtype, types.Complex):
retty = arg.dtype.underlying_float
else:
retty = arg.dtype
def array_angle_impl(arr, deg=False):
out = numpy.zeros_like(arr, dtype=retty)
for index, val in numpy.ndenumerate(arr):
out[index] = numpy.angle(val, deg)
return out
res = context.compile_internal(builder, array_angle_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_angle_kwarg(context, builder, sig, args):
arg = sig.args[0]
ret_dtype = sig.return_type.dtype
def array_angle_impl(arr, deg):
out = numpy.zeros_like(arr, dtype=ret_dtype)
for index, val in numpy.ndenumerate(arr):
out[index] = numpy.angle(val, deg)
return out
if len(args) == 1:
args = args + (cgutils.false_bit,)
sig = signature(sig.return_type, *(sig.args + (types.boolean,)))
res = context.compile_internal(builder, array_angle_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_angle_kwarg(context, builder, sig, args):
arg = sig.args[0]
ret_dtype = sig.return_type.dtype
def array_angle_impl(arr, deg):
out = np.zeros_like(arr, dtype=ret_dtype)
for index, val in np.ndenumerate(arr):
out[index] = np.angle(val, deg)
return out
if len(args) == 1:
args = args + (cgutils.false_bit, )
sig = signature(sig.return_type, *(sig.args + (types.boolean, )))
res = context.compile_internal(builder, array_angle_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_pop(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
idx = inst.fix_index(args[1])
n = inst.size
cgutils.guard_zero(context, builder, n,
(IndexError, "pop from empty list"))
inst.guard_index(idx, "pop index out of range")
res = inst.getitem(idx)
one = ir.Constant(n.type, 1)
n = builder.sub(n, ir.Constant(n.type, 1))
inst.move(idx, builder.add(idx, one), builder.sub(n, idx))
inst.resize(n)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_pop(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
idx = inst.fix_index(args[1])
n = inst.size
cgutils.guard_zero(context, builder, n,
(IndexError, "pop from empty list"))
inst.guard_index(idx, "pop index out of range")
res = inst.getitem(idx)
one = ir.Constant(n.type, 1)
n = builder.sub(n, ir.Constant(n.type, 1))
inst.move(idx, builder.add(idx, one), builder.sub(n, idx))
inst.resize(n)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_cumprod(context, builder, sig, args):
scalar_dtype = sig.return_type.dtype
dtype = as_dtype(scalar_dtype)
def array_cumprod_impl(arr):
size = 1
for i in arr.shape:
size = size * i
out = numpy.empty(size, dtype)
c = 1
for idx, v in enumerate(arr.flat):
c *= v
out[idx] = c
return out
res = context.compile_internal(builder, array_cumprod_impl, sig, args, locals=dict(c=scalar_dtype))
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_mul(context, builder, sig, args):
src = ListInstance(context, builder, sig.args[0], args[0])
src_size = src.size
mult = args[1]
zero = ir.Constant(mult.type, 0)
mult = builder.select(cgutils.is_neg_int(builder, mult), zero, mult)
nitems = builder.mul(mult, src_size)
dest = ListInstance.allocate(context, builder, sig.return_type, nitems)
dest.size = nitems
with cgutils.for_range_slice(builder, zero, nitems, src_size, inc=True) as (dest_offset, _):
with cgutils.for_range(builder, src_size) as loop:
value = src.getitem(loop.index)
dest.setitem(builder.add(loop.index, dest_offset), value)
return impl_ret_new_ref(context, builder, sig.return_type, dest.value)
def any_where(context, builder, sig, args):
cond = sig.args[0]
if isinstance(cond, types.Array):
return array_where(context, builder, sig, args)
def scalar_where_impl(cond, x, y):
"""
np.where(scalar, scalar, scalar): return a 0-dim array
"""
scal = x if cond else y
# This is the equivalent of np.full_like(scal, scal),
# for compatibility with Numpy < 1.8
arr = np.empty_like(scal)
arr[()] = scal
return arr
res = context.compile_internal(builder, scalar_where_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def list_add(context, builder, sig, args):
a = ListInstance(context, builder, sig.args[0], args[0])
b = ListInstance(context, builder, sig.args[1], args[1])
a_size = a.size
b_size = b.size
nitems = builder.add(a_size, b_size)
dest = ListInstance.allocate(context, builder, sig.return_type, nitems)
dest.size = nitems
with cgutils.for_range(builder, a_size) as loop:
value = a.getitem(loop.index)
dest.setitem(loop.index, value)
with cgutils.for_range(builder, b_size) as loop:
value = b.getitem(loop.index)
dest.setitem(builder.add(loop.index, a_size), value)
return impl_ret_new_ref(context, builder, sig.return_type, dest.value)
def build_set(context, builder, set_type, items):
"""
Build a set of the given type, containing the given items.
"""
nitems = len(items)
inst = SetInstance.allocate(context, builder, set_type, nitems)
# Populate set. Inlining the insertion code for each item would be very
# costly, instead we create a LLVM array and iterate over it.
array = cgutils.pack_array(builder, items)
array_ptr = cgutils.alloca_once_value(builder, array)
count = context.get_constant(types.intp, nitems)
with cgutils.for_range(builder, count) as loop:
item = builder.load(cgutils.gep(builder, array_ptr, 0, loop.index))
inst.add(item)
return impl_ret_new_ref(context, builder, set_type, inst.value)
def string_split_impl(context, builder, sig, args):
nitems = cgutils.alloca_once(builder, lir.IntType(64))
# input str, sep, size pointer
fnty = lir.FunctionType(lir.IntType(8).as_pointer().as_pointer(),
[lir.IntType(8).as_pointer(), lir.IntType(8).as_pointer(),
lir.IntType(64).as_pointer()])
fn = builder.module.get_or_insert_function(fnty, name="str_split")
ptr = builder.call(fn, args+[nitems])
size = builder.load(nitems)
# TODO: use ptr instead of allocating and copying, use NRT_MemInfo_new
# TODO: deallocate ptr
_list = numba.targets.listobj.ListInstance.allocate(context, builder,
sig.return_type, size)
_list.size = size
with cgutils.for_range(builder, size) as loop:
value = builder.load(cgutils.gep_inbounds(builder, ptr, loop.index))
_list.setitem(loop.index, value)
return impl_ret_new_ref(context, builder, sig.return_type, _list.value)
def any_where(context, builder, sig, args):
cond = sig.args[0]
if isinstance(cond, types.Array):
return array_where(context, builder, sig, args)
def scalar_where_impl(cond, x, y):
"""
np.where(scalar, scalar, scalar): return a 0-dim array
"""
scal = x if cond else y
# This is the equivalent of numpy.full_like(scal, scal),
# for compatibility with Numpy < 1.8
arr = numpy.empty_like(scal)
arr[()] = scal
return arr
res = context.compile_internal(builder, scalar_where_impl, sig, args)
return impl_ret_new_ref(context, builder, sig.return_type, res)
def build_set(context, builder, set_type, items):
"""
Build a set of the given type, containing the given items.
"""
nitems = len(items)
inst = SetInstance.allocate(context, builder, set_type, nitems)
# Populate set. Inlining the insertion code for each item would be very
# costly, instead we create a LLVM array and iterate over it.
array = cgutils.pack_array(builder, items)
array_ptr = cgutils.alloca_once_value(builder, array)
count = context.get_constant(types.intp, nitems)
with cgutils.for_range(builder, count) as loop:
item = builder.load(cgutils.gep(builder, array_ptr, 0, loop.index))
inst.add(item)
return impl_ret_new_ref(context, builder, set_type, inst.value)
def list_add(context, builder, sig, args):
a = ListInstance(context, builder, sig.args[0], args[0])
b = ListInstance(context, builder, sig.args[1], args[1])
a_size = a.size
b_size = b.size
nitems = builder.add(a_size, b_size)
dest = ListInstance.allocate(context, builder, sig.return_type, nitems)
dest.size = nitems
with cgutils.for_range(builder, a_size) as src_index:
value = a.getitem(src_index)
dest.setitem(src_index, value)
with cgutils.for_range(builder, b_size) as src_index:
value = b.getitem(src_index)
dest.setitem(builder.add(src_index, a_size), value)
return impl_ret_new_ref(context, builder, sig.return_type, dest.value)
def _image_to_array(context, builder, shapes_array, arrtype, data, img):
# allocate array
shapes = cgutils.unpack_tuple(builder, builder.load(shapes_array))
ary = _empty_nd_impl(context, builder, arrtype, shapes)
cgutils.raw_memcpy(builder,
ary.data,
builder.load(data),
ary.nitems,
ary.itemsize,
align=1)
# clean up cv::Mat image
fnty = lir.FunctionType(lir.VoidType(), [lir.IntType(8).as_pointer()])
fn_release = builder.module.get_or_insert_function(fnty,
name="cv_mat_release")
builder.call(fn_release, [img])
return impl_ret_new_ref(context, builder, arrtype, ary._getvalue())
def list_mul(context, builder, sig, args):
src = ListInstance(context, builder, sig.args[0], args[0])
src_size = src.size
mult = args[1]
zero = ir.Constant(mult.type, 0)
mult = builder.select(cgutils.is_neg_int(builder, mult), zero, mult)
nitems = builder.mul(mult, src_size)
dest = ListInstance.allocate(context, builder, sig.return_type, nitems)
dest.size = nitems
with cgutils.for_range_slice(builder, zero, nitems, src_size, inc=True) as (dest_offset, _):
with cgutils.for_range(builder, src_size) as loop:
value = src.getitem(loop.index)
dest.setitem(builder.add(loop.index, dest_offset), value)
return impl_ret_new_ref(context, builder, sig.return_type, dest.value)
def codegen(context, builder, sig, args):
str_arr, _ = args
meminfo, meminfo_data_ptr = construct_str_arr_split_view(
context, builder)
in_str_arr = context.make_helper(builder, string_array_type, str_arr)
# (str_arr_split_view_payload* out_view, int64_t n_strs,
# uint32_t* offsets, char* data, char sep)
fnty = lir.FunctionType(lir.VoidType(), [
meminfo_data_ptr.type,
lir.IntType(64),
lir.IntType(32).as_pointer(),
lir.IntType(8).as_pointer(),
lir.IntType(8)
])
fn_impl = builder.module.get_or_insert_function(
fnty, name="str_arr_split_view_impl")
sep_val = context.get_constant(types.int8, ord(sep_typ.literal_value))
builder.call(fn_impl, [
meminfo_data_ptr, in_str_arr.num_items, in_str_arr.offsets,
in_str_arr.data, sep_val
])
view_payload = cgutils.create_struct_proxy(
str_arr_split_view_payload_type)(
context, builder, value=builder.load(meminfo_data_ptr))
out_view = context.make_helper(builder, string_array_split_view_type)
out_view.num_items = in_str_arr.num_items
out_view.index_offsets = view_payload.index_offsets
out_view.data_offsets = view_payload.data_offsets
# TODO: incref?
out_view.data = context.compile_internal(
builder, lambda S: get_data_ptr(S),
data_ctypes_type(string_array_type), [str_arr])
# out_view.null_bitmap = view_payload.null_bitmap
out_view.meminfo = meminfo
ret = out_view._getvalue()
#context.nrt.decref(builder, ty, ret)
return impl_ret_new_ref(context, builder, string_array_split_view_type,
ret)
def getslice_list(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
slice = slicing.Slice(context, builder, value=args[1])
cgutils.guard_invalid_slice(context, builder, slice)
inst.fix_slice(slice)
# Allocate result and populate it
result_size = slicing.get_slice_length(builder, slice)
result = ListInstance.allocate(context, builder, sig.return_type, result_size)
result.size = result_size
with cgutils.for_range_slice_generic(builder, slice.start, slice.stop, slice.step) as (pos_range, neg_range):
with pos_range as (idx, count):
value = inst.getitem(idx)
result.inititem(count, value)
with neg_range as (idx, count):
value = inst.getitem(idx)
result.inititem(count, value)
return impl_ret_new_ref(context, builder, sig.return_type, result.value)
def ctor_impl(context, builder, sig, args):
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt_meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt_meminfo_data(builder, meminfo)
data_pointer = builder.bitcast(data_pointer,
alloc_type.as_pointer())
# Nullify all data
builder.store(cgutils.get_null_value(alloc_type),
data_pointer)
inst_struct_typ = cgutils.create_struct_proxy(inst_typ)
inst_struct = inst_struct_typ(context, builder)
inst_struct.meminfo = meminfo
inst_struct.data = data_pointer
# Call the __init__
# TODO: extract the following into a common util
init_sig = (sig.return_type,) + sig.args
init = inst_typ.jitmethods['__init__']
init.compile(init_sig)
cres = init._compileinfos[init_sig]
realargs = [inst_struct._getvalue()] + list(args)
context.call_internal(builder, cres.fndesc, types.void(*init_sig),
realargs)
# Prepare reutrn value
ret = inst_struct._getvalue()
# Add function to link
codegen = context.codegen()
codegen.add_linking_library(cres.library)
return imputils.impl_ret_new_ref(context, builder, inst_typ, ret)
def random_arr(context, builder, sig, args, typing_key=typing_key):
from . import arrayobj
arrty = sig.return_type
dtype = arrty.dtype
scalar_sig = signature(dtype, *sig.args[:-1])
scalar_args = args[:-1]
# Allocate array...
shapes = arrayobj._parse_shape(context, builder, sig.args[-1], args[-1])
arr = arrayobj._empty_nd_impl(context, builder, arrty, shapes)
# ... and populate it in natural order
scalar_impl = context.get_function(typing_key, scalar_sig)
with cgutils.for_range(builder, arr.nitems) as loop:
val = scalar_impl(builder, scalar_args)
ptr = cgutils.gep(builder, arr.data, loop.index)
arrayobj.store_item(context, builder, arrty, val, ptr)
return impl_ret_new_ref(context, builder, sig.return_type, arr._getvalue())
def random_arr(context, builder, sig, args, typing_key=typing_key):
from . import arrayobj
arrty = sig.return_type
dtype = arrty.dtype
scalar_sig = signature(dtype, *sig.args[:-1])
scalar_args = args[:-1]
# Allocate array...
shapes = arrayobj._parse_shape(context, builder, sig.args[-1], args[-1])
arr = arrayobj._empty_nd_impl(context, builder, arrty, shapes)
# ... and populate it in natural order
scalar_impl = context.get_function(typing_key, scalar_sig)
with cgutils.for_range(builder, arr.nitems) as loop:
val = scalar_impl(builder, scalar_args)
ptr = cgutils.gep(builder, arr.data, loop.index)
arrayobj.store_item(context, builder, arrty, val, ptr)
return impl_ret_new_ref(context, builder, sig.return_type, arr._getvalue())
def list_add(context, builder, sig, args):
a = ListInstance(context, builder, sig.args[0], args[0])
b = ListInstance(context, builder, sig.args[1], args[1])
a_size = a.size
b_size = b.size
nitems = builder.add(a_size, b_size)
dest = ListInstance.allocate(context, builder, sig.return_type, nitems)
dest.size = nitems
with cgutils.for_range(builder, a_size) as loop:
value = a.getitem(loop.index)
value = context.cast(builder, value, a.dtype, dest.dtype)
dest.setitem(loop.index, value, incref=True)
with cgutils.for_range(builder, b_size) as loop:
value = b.getitem(loop.index)
value = context.cast(builder, value, b.dtype, dest.dtype)
dest.setitem(builder.add(loop.index, a_size), value, incref=True)
return impl_ret_new_ref(context, builder, sig.return_type, dest.value)
def ctor_impl(context, builder, sig, args):
"""
Generic constructor (__new__) for jitclasses.
"""
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt.meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt.meminfo_data(builder, meminfo)
data_pointer = builder.bitcast(data_pointer,
alloc_type.as_pointer())
# Nullify all data
builder.store(cgutils.get_null_value(alloc_type),
data_pointer)
inst_struct = context.make_helper(builder, inst_typ)
inst_struct.meminfo = meminfo
inst_struct.data = data_pointer
# Call the jitted __init__
# TODO: extract the following into a common util
init_sig = (sig.return_type,) + sig.args
init = inst_typ.jitmethods['__init__']
disp_type = types.Dispatcher(init)
call = context.get_function(disp_type, types.void(*init_sig))
_add_linking_libs(context, call)
realargs = [inst_struct._getvalue()] + list(args)
call(builder, realargs)
# Prepare return value
ret = inst_struct._getvalue()
return imputils.impl_ret_new_ref(context, builder, inst_typ, ret)
def ctor_impl(context, builder, sig, args):
"""
Generic constructor (__new__) for jitclasses.
"""
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt.meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt.meminfo_data(builder, meminfo)
data_pointer = builder.bitcast(data_pointer,
alloc_type.as_pointer())
# Nullify all data
builder.store(cgutils.get_null_value(alloc_type),
data_pointer)
inst_struct = context.make_helper(builder, inst_typ)
inst_struct.meminfo = meminfo
inst_struct.data = data_pointer
# Call the jitted __init__
# TODO: extract the following into a common util
init_sig = (sig.return_type,) + sig.args
init = inst_typ.jitmethods['__init__']
disp_type = types.Dispatcher(init)
call = context.get_function(disp_type, types.void(*init_sig))
_add_linking_libs(context, call)
realargs = [inst_struct._getvalue()] + list(args)
call(builder, realargs)
# Prepare return value
ret = inst_struct._getvalue()
return imputils.impl_ret_new_ref(context, builder, inst_typ, ret)
def ctor_impl(context, builder, sig, args):
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt_meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt_meminfo_data(builder, meminfo)
data_pointer = builder.bitcast(data_pointer, alloc_type.as_pointer())
# Nullify all data
builder.store(cgutils.get_null_value(alloc_type), data_pointer)
inst_struct_typ = cgutils.create_struct_proxy(inst_typ)
inst_struct = inst_struct_typ(context, builder)
inst_struct.meminfo = meminfo
inst_struct.data = data_pointer
# Call the __init__
# TODO: extract the following into a common util
init_sig = (sig.return_type, ) + sig.args
init = inst_typ.jitmethods['__init__']
init.compile(init_sig)
cres = init._compileinfos[init_sig]
realargs = [inst_struct._getvalue()] + list(args)
context.call_internal(builder, cres.fndesc, types.void(*init_sig),
realargs)
# Prepare reutrn value
ret = inst_struct._getvalue()
# Add function to link
codegen = context.codegen()
codegen.add_linking_library(cres.library)
return imputils.impl_ret_new_ref(context, builder, inst_typ, ret)
def pq_read_string_lower(context, builder, sig, args):
typ = sig.return_type
dtype = StringArrayPayloadType()
meminfo, data_pointer = construct_string_array(context, builder)
string_array = cgutils.create_struct_proxy(dtype)(context, builder)
string_array.size = args[2]
fnty = lir.FunctionType(lir.IntType(32),
[lir.IntType(8).as_pointer(), lir.IntType(64),
lir.IntType(8).as_pointer().as_pointer(),
lir.IntType(8).as_pointer().as_pointer()])
fn = builder.module.get_or_insert_function(fnty, name="pq_read_string")
res = builder.call(fn, [args[0], args[1],
string_array._get_ptr_by_name('offsets'),
string_array._get_ptr_by_name('data')])
builder.store(string_array._getvalue(),
data_pointer)
inst_struct = context.make_helper(builder, typ)
inst_struct.meminfo = meminfo
ret = inst_struct._getvalue()
return impl_ret_new_ref(context, builder, typ, ret)
def getslice_list(context, builder, sig, args):
inst = ListInstance(context, builder, sig.args[0], args[0])
slice = context.make_helper(builder, sig.args[1], args[1])
slicing.guard_invalid_slice(context, builder, sig.args[1], slice)
inst.fix_slice(slice)
# Allocate result and populate it
result_size = slicing.get_slice_length(builder, slice)
result = ListInstance.allocate(context, builder, sig.return_type,
result_size)
result.size = result_size
with cgutils.for_range_slice_generic(builder, slice.start, slice.stop,
slice.step) as (pos_range, neg_range):
with pos_range as (idx, count):
value = inst.getitem(idx)
result.inititem(count, value)
with neg_range as (idx, count):
value = inst.getitem(idx)
result.inititem(count, value)
return impl_ret_new_ref(context, builder, sig.return_type, result.value)
def getiter_unicode(context, builder, sig, args):
[ty] = sig.args
[data] = args
iterobj = context.make_helper(builder, sig.return_type)
# set the index to zero
zero = context.get_constant(types.uintp, 0)
indexptr = cgutils.alloca_once_value(builder, zero)
iterobj.index = indexptr
# wire in the unicode type data
iterobj.data = data
# incref as needed
if context.enable_nrt:
context.nrt.incref(builder, ty, data)
res = iterobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
def array_cumprod(context, builder, sig, args):
scalar_dtype = sig.return_type.dtype
dtype = as_dtype(scalar_dtype)
def array_cumprod_impl(arr):
size = 1
for i in arr.shape:
size = size * i
out = np.empty(size, dtype)
c = 1
for idx, v in enumerate(arr.flat):
c *= v
out[idx] = c
return out
res = context.compile_internal(builder,
array_cumprod_impl,
sig,
args,
locals=dict(c=scalar_dtype))
return impl_ret_new_ref(context, builder, sig.return_type, res)
def attr_impl(context, builder, typ, value, attr):
if attr in typ.struct:
inst_struct = cgutils.create_struct_proxy(typ)
inst = inst_struct(context, builder, value=value)
data_pointer = inst.data
data_struct = cgutils.create_struct_proxy(typ.get_data_type(),
kind='data')
data = data_struct(context, builder, ref=data_pointer)
return imputils.impl_ret_borrowed(context, builder, typ.struct[attr],
getattr(data, attr))
elif attr in typ.jitprops:
getter = typ.jitprops[attr]['get']
sig = templates.signature(None, typ)
dispatcher = types.Dispatcher(getter)
sig = dispatcher.get_call_type(context.typing_context, [typ], {})
call = context.get_function(dispatcher, sig)
out = call(builder, [value])
return imputils.impl_ret_new_ref(context, builder, sig.return_type,
out)
raise NotImplementedError('attribute {0!r} not implemented'.format(attr))
def getiter_unicode(context, builder, sig, args):
[ty] = sig.args
[data] = args
iterobj = context.make_helper(builder, sig.return_type)
# set the index to zero
zero = context.get_constant(types.uintp, 0)
indexptr = cgutils.alloca_once_value(builder, zero)
iterobj.index = indexptr
# wire in the unicode type data
iterobj.data = data
# incref as needed
if context.enable_nrt:
context.nrt.incref(builder, ty, data)
res = iterobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
def attr_impl(context, builder, typ, value, attr):
if attr in typ.struct:
inst_struct = cgutils.create_struct_proxy(typ)
inst = inst_struct(context, builder, value=value)
data_pointer = inst.data
data_struct = cgutils.create_struct_proxy(typ.get_data_type(),
kind='data')
data = data_struct(context, builder, ref=data_pointer)
return imputils.impl_ret_borrowed(context, builder,
typ.struct[attr],
getattr(data, attr))
elif attr in typ.jitprops:
getter = typ.jitprops[attr]['get']
sig = templates.signature(None, typ)
dispatcher = types.Dispatcher(getter)
sig = dispatcher.get_call_type(context.typing_context, [typ], {})
call = context.get_function(dispatcher, sig)
out = call(builder, [value])
return imputils.impl_ret_new_ref(context, builder, sig.return_type, out)
raise NotImplementedError('attribute {0!r} not implemented'.format(attr))
def pq_read_string_parallel_lower(context, builder, sig, args):
typ = sig.return_type
dtype = StringArrayPayloadType()
meminfo, meminfo_data_ptr = construct_string_array(context, builder)
str_arr_payload = cgutils.create_struct_proxy(dtype)(context, builder)
string_array = context.make_helper(builder, typ)
string_array.num_items = args[3]
fnty = lir.FunctionType(lir.IntType(32), [
lir.IntType(8).as_pointer(),
lir.IntType(64),
lir.IntType(32).as_pointer().as_pointer(),
lir.IntType(8).as_pointer().as_pointer(),
lir.IntType(8).as_pointer().as_pointer(),
lir.IntType(64),
lir.IntType(64)
])
fn = builder.module.get_or_insert_function(fnty,
name="pq_read_string_parallel")
res = builder.call(fn, [
args[0], args[1],
str_arr_payload._get_ptr_by_name('offsets'),
str_arr_payload._get_ptr_by_name('data'),
str_arr_payload._get_ptr_by_name('null_bitmap'), args[2], args[3]
])
builder.store(str_arr_payload._getvalue(), meminfo_data_ptr)
string_array.meminfo = meminfo
string_array.offsets = str_arr_payload.offsets
string_array.data = str_arr_payload.data
string_array.null_bitmap = str_arr_payload.null_bitmap
string_array.num_total_chars = builder.zext(
builder.load(
builder.gep(string_array.offsets, [string_array.num_items])),
lir.IntType(64))
ret = string_array._getvalue()
return impl_ret_new_ref(context, builder, typ, ret)
def make_enumerate_object(context, builder, sig, args):
assert len(args) == 1 or len(args) == 2 # enumerate(it) or enumerate(it, start)
srcty = sig.args[0]
if len(args) == 1:
src = args[0]
start_val = context.get_constant(types.intp, 0)
elif len(args) == 2:
src = args[0]
start_val = context.cast(builder, args[1], sig.args[1], types.intp)
iterobj = call_getiter(context, builder, srcty, src)
enum = context.make_helper(builder, sig.return_type)
countptr = cgutils.alloca_once(builder, start_val.type)
builder.store(start_val, countptr)
enum.count = countptr
enum.iter = iterobj
res = enum._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
def mat_inv(context, builder, sig, args):
"""
Invert a matrix through the use of its LU decomposition.
"""
xty = sig.args[0]
dtype = xty.dtype
x = make_array(xty)(context, builder, args[0])
x_shapes = cgutils.unpack_tuple(builder, x.shape)
m, n = x_shapes
check_c_int(context, builder, m)
check_c_int(context, builder, n)
# Allocate the return array (Numpy never works in place contrary to
# Scipy for which one can specify to whether or not to overwrite the
# input).
def create_out(a):
m, n = a.shape
if m != n:
raise np.linalg.LinAlgError("Last 2 dimensions of "
"the array must be square.")
return a.copy()
out = context.compile_internal(builder, create_out,
signature(sig.return_type, *sig.args), args)
o = make_array(xty)(context, builder, out)
# Allocate the array in which the pivot indices are stored.
ipiv_t = types.Array(types.intc, 1, 'C')
i = _empty_nd_impl(context, builder, ipiv_t, (m,))
ipiv = i._getvalue()
info = cgutils.alloca_once(builder, intp_t)
# Compute the LU decomposition of the matrix.
call_xxgetrf(context, builder, xty, x_shapes, o.data, i.data,
info)
info_val = builder.load(info)
zero = info_val.type(0)
lapack_error = builder.icmp_signed('!=', info_val, zero)
invalid_arg = builder.icmp_signed('<', info_val, zero)
with builder.if_then(lapack_error, False):
context.nrt_decref(builder, ipiv_t, ipiv)
with builder.if_else(invalid_arg) as (then, otherwise):
raise_err = context.call_conv.return_user_exc
with then:
raise_err(builder, ValueError,
('One argument passed to getrf is invalid',)
)
with otherwise:
raise_err(builder, ValueError,
('Matrix is singular and cannot be inverted',)
)
# Compute the optimal lwork.
lwork = make_constant_slot(context, builder, types.intc, -1)
work = cgutils.alloca_once(builder, context.get_value_type(xty.dtype))
call_xxgetri(context, builder, xty, x_shapes, o.data, i.data, work,
lwork, info)
info_val = builder.load(info)
lapack_error = builder.icmp_signed('!=', info_val, zero)
with builder.if_then(lapack_error, False):
context.nrt_decref(builder, ipiv_t, ipiv)
raise_err = context.call_conv.return_user_exc
raise_err(builder, ValueError,
('One argument passed to getri is invalid',)
)
# Allocate a work array of the optimal size as computed by getri.
def allocate_work(x, size):
"""Allocate the work array.
"""
size = int(1.01 * size.real)
return np.empty((size,), dtype=x.dtype)
wty = types.Array(dtype, 1, 'C')
work = context.compile_internal(builder, allocate_work,
signature(wty, xty, dtype),
(args[0], builder.load(work)))
w = make_array(wty)(context, builder, work)
w_shapes = cgutils.unpack_tuple(builder, w.shape)
lw, = w_shapes
builder.store(context.cast(builder, lw, types.intp, types.intc),
lwork)
# Compute the matrix inverse.
call_xxgetri(context, builder, xty, x_shapes, o.data, i.data, w.data,
lwork, info)
info_val = builder.load(info)
lapack_error = builder.icmp_signed('!=', info_val, zero)
invalid_arg = builder.icmp_signed('<', info_val, zero)
context.nrt_decref(builder, wty, work)
context.nrt_decref(builder, ipiv_t, ipiv)
with builder.if_then(lapack_error, False):
with builder.if_else(invalid_arg) as (then, otherwise):
raise_err = context.call_conv.return_user_exc
with then:
raise_err(builder, ValueError,
('One argument passed to getri is invalid',)
)
with otherwise:
raise_err(builder, ValueError,
('Matrix is singular and cannot be inverted',)
)
return impl_ret_new_ref(context, builder, sig.return_type, out)
def set_empty_constructor(context, builder, sig, args):
set_type = sig.return_type
inst = SetInstance.allocate(context, builder, set_type)
return impl_ret_new_ref(context, builder, set_type, inst.value)
def set_copy(context, builder, sig, args):
inst = SetInstance(context, builder, sig.args[0], args[0])
other = inst.copy()
return impl_ret_new_ref(context, builder, sig.return_type, other.value)
