def range_iter_len(typingctx, val):
"""
An implementation of len(range_iter) for internal use.
"""
if isinstance(val, types.RangeIteratorType):
val_type = val.yield_type
def codegen(context, builder, sig, args):
(value,) = args
iter_type = range_impl_map[val_type][1]
iterobj = cgutils.create_struct_proxy(iter_type)(context, builder, value)
int_type = iterobj.count.type
return impl_ret_untracked(context, builder, int_type, builder.load(iterobj.count))
return signature(val_type, val), codegen
elif isinstance(val, types.ListIter):
def codegen(context, builder, sig, args):
(value,) = args
intp_t = context.get_value_type(types.intp)
iterobj = ListIterInstance(context, builder, sig.args[0], value)
return impl_ret_untracked(context, builder, intp_t, iterobj.size)
return signature(types.intp, val), codegen
elif isinstance(val, types.ArrayIterator):
def codegen(context, builder, sig, args):
(iterty,) = sig.args
(value,) = args
intp_t = context.get_value_type(types.intp)
iterobj = context.make_helper(builder, iterty, value=value)
arrayty = iterty.array_type
ary = make_array(arrayty)(context, builder, value=iterobj.array)
shape = cgutils.unpack_tuple(builder, ary.shape)
# array iterates along the outer dimension
return impl_ret_untracked(context, builder, intp_t, shape[0])
return signature(types.intp, val), codegen
def add(self, sig=None, argtypes=None, restype=None):
# Handle argtypes
if argtypes is not None:
warnings.warn("Keyword argument argtypes is deprecated",
DeprecationWarning)
assert sig is None
if restype is None:
sig = tuple(argtypes)
else:
sig = restype(*argtypes)
del argtypes
del restype
# compile core as device function
args, return_type = sigutils.normalize_signature(sig)
devfnsig = signature(return_type, *args)
funcname = self.pyfunc.__name__
kernelsource = self._get_kernel_source(self._kernel_template,
devfnsig, funcname)
corefn, return_type = self._compile_core(devfnsig)
glbl = self._get_globals(corefn)
sig = signature(types.void, *([a[:] for a in args] + [return_type[:]]))
_exec(kernelsource, glbl)
stager = glbl['__vectorized_%s' % funcname]
kernel = self._compile_kernel(stager, sig)
argdtypes = tuple(to_dtype(t) for t in devfnsig.args)
resdtype = to_dtype(return_type)
self.kernelmap[tuple(argdtypes)] = resdtype, kernel
def get_attribute(self, val, typ, attr):
if isinstance(typ, types.Module):
# Implement getattr for module-level globals.
# We are treating them as constants.
# XXX We shouldn't have to retype this
attrty = self.typing_context.resolve_module_constants(typ, attr)
if attrty is not None and not isinstance(attrty, types.Dummy):
pyval = getattr(typ.pymod, attr)
llval = self.get_constant(attrty, pyval)
@impl_attribute(typ, attr, attrty)
def imp(context, builder, typ, val):
return impl_ret_borrowed(context, builder, attrty, llval)
return imp
# No implementation required for dummies (functions, modules...),
# which are dealt with later
return None
# Lookup specific attribute implementation for this type
overloads = self.attrs[attr]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
pass
# Lookup generic getattr implementation for this type
overloads = self.attrs[None]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
raise Exception("No definition for lowering %s.%s" % (typ, attr))
def impl(context, builder, sig, args):
[tyinp, tyout] = sig.args
[inp, out] = args
ndim = tyinp.ndim
iary = context.make_array(tyinp)(context, builder, inp)
oary = context.make_array(tyout)(context, builder, out)
if asfloat:
sig = typing.signature(types.float64, types.float64)
else:
sig = typing.signature(tyout.dtype, tyinp.dtype)
fnwork = context.get_function(funckey, sig)
intpty = context.get_value_type(types.intp)
# TODO handle differing shape by mimicking broadcasting
shape = cgutils.unpack_tuple(builder, iary.shape, ndim)
with cgutils.loop_nest(builder, shape, intp=intpty) as indices:
pi = cgutils.get_item_pointer(builder, tyinp, iary, indices)
po = cgutils.get_item_pointer(builder, tyout, oary, indices)
ival = builder.load(pi)
if asfloat:
dval = context.cast(builder, ival, tyinp.dtype, types.float64)
dres = fnwork(builder, [dval])
res = context.cast(builder, dres, types.float64, tyout.dtype)
elif tyinp.dtype != tyout.dtype:
tempres = fnwork(builder, [ival])
res = context.cast(builder, tempres, tyinp.dtype, tyout.dtype)
else:
res = fnwork(builder, [ival])
builder.store(res, po)
return out
def array_median(context, builder, sig, args):
def partition(A, low, high):
mid = (low + high) // 2
# median of three {low, middle, high}
LM = A[low] <= A[mid]
MH = A[mid] <= A[high]
LH = A[low] <= A[high]
if LM == MH:
median3 = mid
elif LH != LM:
median3 = low
else:
median3 = high
# choose median3 as the pivot
A[high], A[median3] = A[median3], A[high]
x = A[high]
i = low
for j in range(low, high):
if A[j] <= x:
A[i], A[j] = A[j], A[i]
i += 1
A[i], A[high] = A[high], A[i]
return i
sig_partition = typing.signature(types.intp, *(sig.args[0], types.intp, types.intp))
_partition = context.compile_subroutine(builder, partition, sig_partition)
def select(arry, k):
n = arry.shape[0]
# XXX: assuming flat array till array.flatten is implemented
# temp_arry = arry.flatten()
temp_arry = arry.copy()
high = n - 1
low = 0
# NOTE: high is inclusive
i = _partition(temp_arry, low, high)
while i != k:
if i < k:
low = i + 1
i = _partition(temp_arry, low, high)
else:
high = i - 1
i = _partition(temp_arry, low, high)
return temp_arry[k]
sig_select = typing.signature(sig.args[0].dtype, *(sig.args[0], types.intp))
_select = context.compile_subroutine(builder, select, sig_select)
def median(arry):
n = arry.shape[0]
if n % 2 == 0:
return (_select(arry, n // 2 - 1) + _select(arry, n // 2)) / 2
else:
return _select(arry, n // 2)
res = context.compile_internal(builder, median, sig, args)
return impl_ret_untracked(context, builder, sig.return_type, res)
def test_equality(self):
self.assertEqual(types.int32, types.int32)
self.assertEqual(types.uint32, types.uint32)
self.assertEqual(types.complex64, types.complex64)
self.assertEqual(types.float32, types.float32)
# Different signedness
self.assertNotEqual(types.int32, types.uint32)
# Different width
self.assertNotEqual(types.int64, types.int32)
self.assertNotEqual(types.float64, types.float32)
self.assertNotEqual(types.complex64, types.complex128)
# Different domain
self.assertNotEqual(types.int64, types.float64)
self.assertNotEqual(types.uint64, types.float64)
self.assertNotEqual(types.complex64, types.float64)
# Same arguments but different return types
get_pointer = None
sig_a = typing.signature(types.intp, types.intp)
sig_b = typing.signature(types.voidptr, types.intp)
a = types.ExternalFunctionPointer(sig=sig_a, get_pointer=get_pointer)
b = types.ExternalFunctionPointer(sig=sig_b, get_pointer=get_pointer)
self.assertNotEqual(a, b)
# Different call convention
a = types.ExternalFunctionPointer(sig=sig_a, get_pointer=get_pointer)
b = types.ExternalFunctionPointer(sig=sig_a, get_pointer=get_pointer,
cconv='stdcall')
self.assertNotEqual(a, b)
# Different get_pointer
a = types.ExternalFunctionPointer(sig=sig_a, get_pointer=get_pointer)
b = types.ExternalFunctionPointer(sig=sig_a, get_pointer=object())
self.assertNotEqual(a, b)
def inline_array(array_var, expr, stmts, list_vars, dels):
"""Check to see if the given "array_var" is created from a list
of constants, and try to inline the list definition as array
initialization.
Extra statements produced with be appended to "stmts".
"""
callname = guard(find_callname, func_ir, expr)
require(callname and callname[1] == 'numpy' and callname[0] == 'array')
require(expr.args[0].name in list_vars)
ret_type = calltypes[expr].return_type
require(isinstance(ret_type, types.ArrayCompatible) and
ret_type.ndim == 1)
loc = expr.loc
list_var = expr.args[0]
array_typ = typemap[array_var.name]
debug_print("inline array_var = ", array_var, " list_var = ", list_var)
dtype = array_typ.dtype
seq, op = find_build_sequence(func_ir, list_var)
size = len(seq)
size_var = ir.Var(scope, mk_unique_var("size"), loc)
size_tuple_var = ir.Var(scope, mk_unique_var("size_tuple"), loc)
size_typ = types.intp
size_tuple_typ = types.UniTuple(size_typ, 1)
typemap[size_var.name] = size_typ
typemap[size_tuple_var.name] = size_tuple_typ
stmts.append(_new_definition(func_ir, size_var,
ir.Const(size, loc=loc), loc))
stmts.append(_new_definition(func_ir, size_tuple_var,
ir.Expr.build_tuple(items=[size_var], loc=loc), loc))
empty_func = ir.Var(scope, mk_unique_var("empty_func"), loc)
fnty = get_np_ufunc_typ(np.empty)
sig = context.resolve_function_type(fnty, (size_typ,), {})
typemap[empty_func.name] = fnty #
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('empty', np.empty, loc=loc), loc))
empty_call = ir.Expr.call(empty_func, [size_var], {}, loc=loc)
calltypes[empty_call] = typing.signature(array_typ, size_typ)
stmts.append(_new_definition(func_ir, array_var, empty_call, loc))
for i in range(size):
index_var = ir.Var(scope, mk_unique_var("index"), loc)
index_typ = types.intp
typemap[index_var.name] = index_typ
stmts.append(_new_definition(func_ir, index_var,
ir.Const(i, loc), loc))
setitem = ir.SetItem(array_var, index_var, seq[i], loc)
calltypes[setitem] = typing.signature(types.none, array_typ,
index_typ, dtype)
stmts.append(setitem)
stmts.extend(dels)
return True
def test_call_notation(self):
# Function call signature
i = types.int32
d = types.double
self.assertEqual(i(), typing.signature(i))
self.assertEqual(i(d), typing.signature(i, d))
self.assertEqual(i(d, d), typing.signature(i, d, d))
# Value cast
self.assertPreciseEqual(i(42.5), 42)
self.assertPreciseEqual(d(-5), -5.0)
def local_array(shape, dtype):
ndim = 1
if isinstance(shape, tuple):
ndim = len(shape)
fname = "ptx.lmem.alloc"
restype = types.Array(dtype, ndim, 'C')
if ndim == 1:
sig = typing.signature(restype, types.intp, types.Any)
else:
sig = typing.signature(restype, types.UniTuple(types.intp, ndim),
types.Any)
return ir.Intrinsic(fname, sig, args=(shape, dtype))
def get_attribute(self, val, typ, attr):
if isinstance(typ, types.Record):
# Implement get attribute for records
self.sentry_record_alignment(typ, attr)
offset = typ.offset(attr)
elemty = typ.typeof(attr)
@impl_attribute(typ, attr, elemty)
def imp(context, builder, typ, val):
dptr = cgutils.get_record_member(builder, val, offset, self.get_data_type(elemty))
return self.unpack_value(builder, elemty, dptr)
return imp
if isinstance(typ, types.Module):
# Implement getattr for module-level globals.
# We are treating them as constants.
# XXX We shouldn't have to retype this
attrty = self.typing_context.resolve_module_constants(typ, attr)
if attrty is not None:
try:
pyval = getattr(typ.pymod, attr)
llval = self.get_constant(attrty, pyval)
except NotImplementedError:
# Module attribute is not a simple constant
# (e.g. it's a function), it will be handled later on.
pass
else:
@impl_attribute(typ, attr, attrty)
def imp(context, builder, typ, val):
return llval
return imp
# No implementation
return None
# Lookup specific attribute implementation for this type
overloads = self.attrs[attr]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
pass
# Lookup generic getattr implementation for this type
overloads = self.attrs[None]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
raise Exception("No definition for lowering %s.%s" % (typ, attr))
def test_cache(self):
def times2(i):
return 2*i
def times3(i):
return i*3
i32 = lc.Type.int(32)
llvm_fnty = lc.Type.function(i32, [i32])
module = lc.Module.new("test_module")
function = module.get_or_insert_function(llvm_fnty, name='test_fn')
assert function.is_declaration
entry_block = function.append_basic_block('entry')
builder = lc.Builder.new(entry_block)
sig = typing.signature(types.int32, types.int32)
typing_context = typing.Context()
context = cpu.CPUContext(typing_context).localized()
# Ensure the cache is empty to begin with
self.assertEqual(0, len(context.cached_internal_func))
# After one compile, it should contain one entry
context.compile_internal(builder, times2, sig, function.args)
self.assertEqual(1, len(context.cached_internal_func))
# After a second compilation of the same thing, it should still contain
# one entry
context.compile_internal(builder, times2, sig, function.args)
self.assertEqual(1, len(context.cached_internal_func))
# After compilation of another function, the cache should have grown by
# one more.
context.compile_internal(builder, times3, sig, function.args)
self.assertEqual(2, len(context.cached_internal_func))
sig2 = typing.signature(types.float64, types.float64)
f64 = lc.Type.double()
llvm_fnty2 = lc.Type.function(f64, [f64])
function2 = module.get_or_insert_function(llvm_fnty2, name='test_fn_2')
assert function2.is_declaration
entry_block2 = function2.append_basic_block('entry')
builder2 = lc.Builder.new(entry_block2)
# Ensure that the same function with a different signature does not
# reuse an entry from the cache in error
context.compile_internal(builder2, times3, sig2, function2.args)
self.assertEqual(3, len(context.cached_internal_func))
def impl_setitem(d, key, value):
if not isinstance(d, types.DictType):
return
keyty, valty = d.key_type, d.value_type
def impl(d, key, value):
castedkey = _cast(key, keyty)
castedval = _cast(value, valty)
status = _dict_insert(d, castedkey, hash(castedkey), castedval)
if status == Status.OK:
return
elif status == Status.OK_REPLACED:
# replaced
# XXX handle refcount
return
elif status == Status.ERR_CMP_FAILED:
raise ValueError('key comparison failed')
else:
raise RuntimeError('dict.__setitem__ failed unexpectedly')
if d.is_precise():
# Handle the precise case.
return impl
else:
# Handle the imprecise case.
d = d.refine(key, value)
# Re-bind the key type and value type to match the arguments.
keyty, valty = d.key_type, d.value_type
# Create the signature that we wanted this impl to have.
sig = typing.signature(types.void, d, keyty, valty)
return sig, impl
def grid_expand(ndim):
"""grid(ndim)
Return the absolute position of the current thread in the entire
grid of blocks. *ndim* should correspond to the number of dimensions
declared when instantiating the kernel. If *ndim* is 1, a single integer
is returned. If *ndim* is 2 or 3, a tuple of the given number of
integers is returned.
Computation of the first integer is as follows::
cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
and is similar for the other two indices, but using the ``y`` and ``z``
attributes.
"""
if ndim == 1:
fname = "ptx.grid.1d"
restype = types.int32
elif ndim == 2:
fname = "ptx.grid.2d"
restype = types.UniTuple(types.int32, 2)
elif ndim == 3:
fname = "ptx.grid.3d"
restype = types.UniTuple(types.int32, 3)
else:
raise ValueError('argument can only be 1, 2, 3')
return ir.Intrinsic(fname, typing.signature(restype, types.intp),
args=[ndim])
def gridsize_expand(ndim):
"""
Return the absolute size (or shape) in threads of the entire grid of
blocks. *ndim* should correspond to the number of dimensions declared when
instantiating the kernel.
Computation of the first integer is as follows::
cuda.blockDim.x * cuda.gridDim.x
and is similar for the other two indices, but using the ``y`` and ``z``
attributes.
"""
if ndim == 1:
fname = "ptx.gridsize.1d"
restype = types.int32
elif ndim == 2:
fname = "ptx.gridsize.2d"
restype = types.UniTuple(types.int32, 2)
elif ndim == 3:
fname = "ptx.gridsize.3d"
restype = types.UniTuple(types.int32, 3)
else:
raise ValueError('argument can only be 1, 2 or 3')
return ir.Intrinsic(fname, typing.signature(restype, types.intp),
args=[ndim])
def resolve_call(self, fnty, pos_args, kw_args):
"""
Resolve a call to a given function type. A signature is returned.
"""
if isinstance(fnty, types.RecursiveCall) and not self._skip_recursion:
# Recursive call
disp = fnty.dispatcher_type.dispatcher
pysig, args = disp.fold_argument_types(pos_args, kw_args)
frame = self.context.callstack.match(disp.py_func, args)
# If the signature is not being compiled
if frame is None:
sig = self.context.resolve_function_type(fnty.dispatcher_type,
pos_args, kw_args)
fndesc = disp.overloads[args].fndesc
fnty.overloads[args] = qualifying_prefix(fndesc.modname,
fndesc.unique_name)
return sig
fnid = frame.func_id
fnty.overloads[args] = qualifying_prefix(fnid.modname,
fnid.unique_name)
# Resume propagation in parent frame
return_type = frame.typeinfer.return_types_from_partial()
# No known return type
if return_type is None:
raise TypingError("cannot type infer runaway recursion")
sig = typing.signature(return_type, *args)
sig.pysig = pysig
return sig
else:
# Normal non-recursive call
return self.context.resolve_function_type(fnty, pos_args, kw_args)
def local_array(shape, dtype):
shape = _legalize_shape(shape)
ndim = len(shape)
fname = "ptx.lmem.alloc"
restype = types.Array(dtype, ndim, 'C')
sig = typing.signature(restype, types.UniTuple(types.intp, ndim), types.Any)
return ir.Intrinsic(fname, sig, args=(shape, dtype))
def dot_2_vv(context, builder, sig, args, conjugate=False):
"""
np.dot(vector, vector)
np.vdot(vector, vector)
"""
aty, bty = sig.args
dtype = sig.return_type
a = make_array(aty)(context, builder, args[0])
b = make_array(bty)(context, builder, args[1])
n, = cgutils.unpack_tuple(builder, a.shape)
def check_args(a, b):
m, = a.shape
n, = b.shape
if m != n:
raise ValueError("incompatible array sizes for np.dot(a, b) "
"(vector * vector)")
context.compile_internal(builder, check_args,
signature(types.none, *sig.args), args)
check_c_int(context, builder, n)
out = cgutils.alloca_once(builder, context.get_value_type(dtype))
call_xxdot(context, builder, conjugate, dtype, n, a.data, b.data, out)
return builder.load(out)
def dot_3_vm(context, builder, sig, args):
"""
np.dot(vector, matrix, out)
np.dot(matrix, vector, out)
"""
xty, yty, outty = sig.args
assert outty == sig.return_type
dtype = xty.dtype
x = make_array(xty)(context, builder, args[0])
y = make_array(yty)(context, builder, args[1])
out = make_array(outty)(context, builder, args[2])
x_shapes = cgutils.unpack_tuple(builder, x.shape)
y_shapes = cgutils.unpack_tuple(builder, y.shape)
out_shapes = cgutils.unpack_tuple(builder, out.shape)
if xty.ndim < yty.ndim:
# Vector * matrix
# Asked for x * y, we will compute y.T * x
mty = yty
m_shapes = y_shapes
do_trans = yty.layout == 'F'
m_data, v_data = y.data, x.data
def check_args(a, b, out):
m, = a.shape
_m, n = b.shape
if m != _m:
raise ValueError("incompatible array sizes for "
"np.dot(a, b) (vector * matrix)")
if out.shape != (n,):
raise ValueError("incompatible output array size for "
"np.dot(a, b, out) (vector * matrix)")
else:
# Matrix * vector
# We will compute x * y
mty = xty
m_shapes = x_shapes
do_trans = xty.layout == 'C'
m_data, v_data = x.data, y.data
def check_args(a, b, out):
m, _n = a.shape
n, = b.shape
if n != _n:
raise ValueError("incompatible array sizes for np.dot(a, b) "
"(matrix * vector)")
if out.shape != (m,):
raise ValueError("incompatible output array size for "
"np.dot(a, b, out) (matrix * vector)")
context.compile_internal(builder, check_args,
signature(types.none, *sig.args), args)
for val in m_shapes:
check_c_int(context, builder, val)
call_xxgemv(context, builder, do_trans, mty, m_shapes, m_data,
v_data, out.data)
return impl_ret_borrowed(context, builder, sig.return_type,
out._getvalue())
def _gauss_impl(context, builder, sig, args, state):
# The type for all computations (either float or double)
ty = sig.return_type
llty = context.get_data_type(ty)
state_ptr = get_state_ptr(context, builder, state)
_random = {"py": random.random,
"np": np.random.random}[state]
ret = cgutils.alloca_once(builder, llty, name="result")
gauss_ptr = get_gauss_ptr(builder, state_ptr)
has_gauss_ptr = get_has_gauss_ptr(builder, state_ptr)
has_gauss = cgutils.is_true(builder, builder.load(has_gauss_ptr))
with builder.if_else(has_gauss) as (then, otherwise):
with then:
# if has_gauss: return it
builder.store(builder.load(gauss_ptr), ret)
builder.store(const_int(0), has_gauss_ptr)
with otherwise:
# if not has_gauss: compute a pair of numbers using the Box-Muller
# transform; keep one and return the other
pair = context.compile_internal(builder,
_gauss_pair_impl(_random),
signature(types.UniTuple(ty, 2)),
())
first, second = cgutils.unpack_tuple(builder, pair, 2)
builder.store(first, gauss_ptr)
builder.store(second, ret)
builder.store(const_int(1), has_gauss_ptr)
mu, sigma = args
return builder.fadd(mu,
builder.fmul(sigma, builder.load(ret)))
def hypot_u64_impl(context, builder, sig, args):
[x, y] = args
y = builder.sitofp(y, Type.double())
x = builder.sitofp(x, Type.double())
fsig = signature(types.float64, types.float64, types.float64)
res = hypot_float_impl(context, builder, fsig, (x, y))
return impl_ret_untracked(context, builder, sig.return_type, res)
def local_array(shape, dtype):
ndim = 1
if isinstance(shape, tuple):
ndim = len(shape)
elif not isinstance(shape, int):
raise TypeError("invalid type for shape; got {0}".format(type(shape)))
fname = "ptx.lmem.alloc"
restype = types.Array(dtype, ndim, 'C')
if isinstance(shape, int):
sig = typing.signature(restype, types.intp, types.Any)
else:
sig = typing.signature(restype, types.UniTuple(types.intp, ndim),
types.Any)
return ir.Intrinsic(fname, sig, args=(shape, dtype))
def grid_expand(ndim):
"""grid(ndim)
ndim: [int] 1, 2 or 3
if ndim == 1:
return cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
elif ndim == 2:
x = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
y = cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
return x, y
elif ndim == 3:
x = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
y = cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
z = cuda.threadIdx.z + cuda.blockIdx.z * cuda.blockDim.z
return x, y, z
"""
if ndim == 1:
fname = "ptx.grid.1d"
restype = types.int32
elif ndim == 2:
fname = "ptx.grid.2d"
restype = types.UniTuple(types.int32, 2)
elif ndim == 3:
fname = "ptx.grid.3d"
restype = types.UniTuple(types.int32, 3)
else:
raise ValueError('argument can only be 1, 2, 3')
return ir.Intrinsic(fname, typing.signature(restype, types.intp),
args=[ndim])
def _backend(self, lowerfn, objectmode):
"""
Back-end: Generate LLVM IR from Numba IR, compile to machine code
"""
if self.library is None:
codegen = self.targetctx.codegen()
self.library = codegen.create_library(self.bc.func_qualname)
# Enable object caching upfront, so that the library can
# be later serialized.
self.library.enable_object_caching()
lowered = lowerfn()
signature = typing.signature(self.return_type, *self.args)
self.cr = compile_result(typing_context=self.typingctx,
target_context=self.targetctx,
entry_point=lowered.cfunc,
typing_error=self.status.fail_reason,
type_annotation=self.type_annotation,
library=self.library,
call_helper=lowered.call_helper,
signature=signature,
objectmode=objectmode,
interpmode=False,
lifted=self.lifted,
fndesc=lowered.fndesc,
environment=lowered.env,
has_dynamic_globals=lowered.has_dynamic_globals,
)
def rect_impl(context, builder, sig, args):
[r, phi] = args
# We can't call math.isfinite() inside rect() below because it
# only exists on 3.2+.
phi_is_finite = mathimpl.is_finite(builder, phi)
def rect(r, phi, phi_is_finite):
if not phi_is_finite:
if not r:
# cmath.rect(0, phi={inf, nan}) = 0
return abs(r)
if math.isinf(r):
# cmath.rect(inf, phi={inf, nan}) = inf + j phi
return complex(r, phi)
real = math.cos(phi)
imag = math.sin(phi)
if real == 0. and math.isinf(r):
# 0 * inf would return NaN, we want to keep 0 but xor the sign
real /= r
else:
real *= r
if imag == 0. and math.isinf(r):
# ditto
imag /= r
else:
imag *= r
return complex(real, imag)
inner_sig = signature(sig.return_type, *sig.args + (types.boolean,))
res = context.compile_internal(builder, rect, inner_sig,
args + [phi_is_finite])
return impl_ret_untracked(context, builder, sig, res)
def get_function(self, fn, sig):
"""
Return the implementation of function *fn* for signature *sig*.
The return value is a callable with the signature (builder, args).
"""
if isinstance(fn, types.Function):
key = fn.template.key
if isinstance(key, MethodType):
overloads = self.defns[key.im_func]
elif sig.recvr:
sig = typing.signature(sig.return_type,
*((sig.recvr,) + sig.args))
overloads = self.defns[key]
else:
overloads = self.defns[key]
elif isinstance(fn, types.Dispatcher):
key = fn.overloaded.get_overload(sig.args)
overloads = self.defns[key]
else:
key = fn
overloads = self.defns[key]
try:
return _wrap_impl(overloads.find(sig), self, sig)
except NotImplementedError:
raise Exception("No definition for lowering %s%s" % (key, sig))
def test_normalize_signature(self):
f = sigutils.normalize_signature
def check(sig, args, return_type):
self.assertEqual(f(sig), (args, return_type))
def check_error(sig, msg):
with self.assertRaises(TypeError) as raises:
f(sig)
self.assertIn(msg, str(raises.exception))
f32 = types.float32
c64 = types.complex64
i16 = types.int16
a = types.Array(f32, 1, "C")
check((c64,), (c64,), None)
check((f32, i16), (f32, i16), None)
check(a(i16), (i16,), a)
check("int16(complex64)", (c64,), i16)
check("(complex64, int16)", (c64, i16), None)
check(typing.signature(i16, c64), (c64,), i16)
check_error((types.Integer,), "invalid signature")
check_error((None,), "invalid signature")
check_error([], "invalid signature")
def resolve_call(self, fnty, pos_args, kw_args):
"""
Resolve a call to a given function type. A signature is returned.
"""
if isinstance(fnty, types.RecursiveCall):
# Self-recursive call
disp = fnty.dispatcher_type.dispatcher
pysig, args = disp.fold_argument_types(pos_args, kw_args)
# Fetch the return type as given by the user
rettypes = set()
for retvar in self._get_return_vars():
typevar = self.typevars[retvar.name]
if not typevar.defined:
raise TypeError("recursive calls need an explicit signature in jit()")
rettypes.add(typevar.getone())
return_type = self._unify_return_types(rettypes)
# Match call arguments with current inference arguments
assert len(args) == len(self.arg_names)
formal_args = [self.typevars[self.arg_names[i]].getone()
for i in range(len(args))]
for formal_ty, actual_ty in zip(formal_args, args):
if not self.context.can_convert(actual_ty, formal_ty):
raise TypeError("bad self-recursive call with argument types %s" % (args,))
sig = typing.signature(return_type, *formal_args)
sig.pysig = pysig
return sig
else:
# Normal non-recursive call
return self.context.resolve_function_type(fnty, pos_args, kw_args)
def implementer(context, builder, sig, args):
val, = args
input_type = sig.args[0]
fpval = context.cast(builder, val, input_type, types.float64)
inner_sig = signature(types.float64, types.float64)
res = wrapped_impl(context, builder, inner_sig, (fpval,))
return context.cast(builder, res, types.float64, sig.return_type)
def test_error_model(self):
"""
Caching must not mix up different error models.
"""
def inv(x):
return 1.0 / x
inv_sig = typing.signature(types.float64, types.float64)
def compile_inv(context):
return context.compile_subroutine(builder, inv, inv_sig)
context, builder, sig, args = self._context_builder_sig_args()
py_error_model = callconv.create_error_model('python', context)
np_error_model = callconv.create_error_model('numpy', context)
py_context1 = context.subtarget(error_model=py_error_model)
py_context2 = context.subtarget(error_model=py_error_model)
np_context = context.subtarget(error_model=np_error_model)
# Note the parent context's cache is shared by subtargets
self.assertEqual(0, len(context.cached_internal_func))
# Compiling with the same error model reuses the same cache slot
compile_inv(py_context1)
self.assertEqual(1, len(context.cached_internal_func))
compile_inv(py_context2)
self.assertEqual(1, len(context.cached_internal_func))
# Compiling with another error model creates a new cache slot
compile_inv(np_context)
self.assertEqual(2, len(context.cached_internal_func))
def const_array_like(ndarray):
fname = "ptx.cmem.arylike"
from .descriptor import CUDATargetDesc
aryty = CUDATargetDesc.typingctx.resolve_argument_type(ndarray)
sig = typing.signature(aryty, aryty)
return ir.Intrinsic(fname, sig, args=[ndarray])
raise NotImplementedError
[aryarg] = args
ary = aryarg.value
count = reduce(operator.mul, ary.shape)
dtype = types.from_dtype(numpy.dtype(ary.dtype))
def impl(context, args, argtys, retty):
builder = context.builder
lmod = builder.basic_block.function.module
addrspace = nvvm.ADDRSPACE_CONSTANT
data_t = dtype.llvm_as_value()
flat = ary.flatten(order='A') # preserve order
constvals = [dtype.llvm_const(flat[i]) for i in range(flat.size)]
constary = lc.Constant.array(data_t, constvals)
gv = lmod.add_global_variable(constary.type, "cmem", addrspace)
gv.linkage = lc.LINKAGE_INTERNAL
gv.global_constant = True
gv.initializer = constary
byte = lc.Type.int(8)
byte_ptr_as = lc.Type.pointer(byte, addrspace)
to_generic = nvvmutils.insert_addrspace_conv(lmod, byte, addrspace)
rawdata = builder.call(to_generic, [builder.bitcast(gv, byte_ptr_as)])
data = builder.bitcast(rawdata, lc.Type.pointer(data_t))
llintp = types.intp.llvm_as_value()
cshape = lc.Constant.array(llintp,
map(types.const_intp, ary.shape))
cstrides = lc.Constant.array(llintp,
map(types.const_intp, ary.strides))
res = lc.Constant.struct([lc.Constant.null(data.type), cshape,
cstrides])
res = builder.insert_value(res, data, 0)
return res
if ary.flags['C_CONTIGUOUS']:
contig = 'C'
elif ary.flags['F_CONTIGUOUS']:
contig = 'F'
else:
raise TypeError("array must be either C/F contiguous to be used as a "
"constant")
impl.codegen = True
impl.return_type = types.arraytype(dtype, ary.ndim, 'A')
return impl
def generic(self, args, kws):
assert not kws
assert len(args) == 6
# array_ty = types.Array(ndim=1, layout='C', dtype=args[2])
return signature(types.int32, *unliteral_all(args))
def atan2_u64_impl(context, builder, sig, args):
[y, x] = args
y = builder.uitofp(y, Type.double())
x = builder.uitofp(x, Type.double())
fsig = signature(types.float64, types.float64, types.float64)
return atan2_float_impl(context, builder, fsig, (y, x))
def generic(self, args, kws):
assert not kws
assert len(args) == 1
if isinstance(args[0].dtype, (types.Integer, types.Boolean)):
return signature(types.float64, *args)
return signature(args[0].dtype, *args)
def is_true(self, builder, typ, val):
"""
Return the truth value of a value of the given Numba type.
"""
impl = self.get_function(bool, typing.signature(types.boolean, typ))
return impl(builder, (val, ))
def generic(self, args, kws):
arr = args[0] # array or series
# hiframes_typed pass replaces series input with array after typing
from sdc.hiframes.pd_series_ext import if_series_to_array_type
ret_typ = if_series_to_array_type(arr).copy(dtype=types.float64)
return signature(ret_typ, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 0
return signature(types.float64, *unliteral_all(args))
def binop_impl(context, builder, sig, args):
if reverse_args:
args = args[::-1]
sig = typing.signature(sig.return_type, *sig.args[::-1])
impl = context.get_function(op, sig)
return impl(builder, args)
def generic(self, args, kws):
assert not kws
assert len(args) == 2 # value and reduce_op
return signature(args[0], *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 2 # array and val
return signature(types.none, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 1 # array
return signature(args[0], *args)
def generic(self, args, kws):
assert not kws
assert len(args) in [4, 5]
return signature(mpi_req_numba_type, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 1
assert (args[0] == types.Array(types.NPDatetime('ns'), 1, 'C')
or args[0] == types.Array(types.int64, 1, 'C'))
return signature(types.Array(types.int64, 1, 'C'), *args)
def generic(self, args, kws):
assert not kws
assert len(args) == 4
return signature(string_array_type, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 2
return signature(types.int32, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 2
return signature(types.intp, args[0], types.unliteral(args[1]))
def lower_expr(self, resty, expr):
if expr.op == 'binop':
return self.lower_binop(resty, expr)
elif expr.op == 'inplace_binop':
lty = self.typeof(expr.lhs.name)
if not lty.mutable:
# inplace operators on non-mutable types reuse the same
# definition as the corresponding copying operators.
return self.lower_binop(resty, expr)
elif expr.op == 'unary':
val = self.loadvar(expr.value.name)
typ = self.typeof(expr.value.name)
# Get function
signature = self.fndesc.calltypes[expr]
impl = self.context.get_function(expr.fn, signature)
# Convert argument to match
val = self.context.cast(self.builder, val, typ, signature.args[0])
res = impl(self.builder, [val])
return self.context.cast(self.builder, res, signature.return_type,
resty)
elif expr.op == 'call':
argvals = [self.loadvar(a.name) for a in expr.args]
argtyps = [self.typeof(a.name) for a in expr.args]
signature = self.fndesc.calltypes[expr]
if isinstance(expr.func, ir.Intrinsic):
fnty = expr.func.name
castvals = expr.func.args
else:
assert not expr.kws, expr.kws
fnty = self.typeof(expr.func.name)
castvals = [
self.context.cast(self.builder, av, at, ft)
for av, at, ft in zip(argvals, argtyps, signature.args)
]
if isinstance(fnty, types.Method):
# Method of objects are handled differently
fnobj = self.loadvar(expr.func.name)
res = self.context.call_class_method(self.builder, fnobj,
signature, castvals)
elif isinstance(fnty, types.FunctionPointer):
# Handle function pointer
pointer = fnty.funcptr
res = self.context.call_function_pointer(
self.builder, pointer, signature, castvals, fnty.cconv)
elif isinstance(fnty, cffi_support.ExternCFunction):
# XXX unused?
fndesc = ExternalFunctionDescriptor(fnty.symbol, fnty.restype,
fnty.argtypes)
func = self.context.declare_external_function(
cgutils.get_module(self.builder), fndesc)
res = self.context.call_external_function(
self.builder, func, fndesc.argtypes, castvals)
else:
# Normal function resolution
impl = self.context.get_function(fnty, signature)
if signature.recvr:
# The "self" object is passed as the function object
# for bounded function
the_self = self.loadvar(expr.func.name)
# Prepend the self reference
castvals = [the_self] + castvals
res = impl(self.builder, castvals)
libs = getattr(impl, "libs", ())
for lib in libs:
self.library.add_linking_library(lib)
return self.context.cast(self.builder, res, signature.return_type,
resty)
elif expr.op == 'pair_first':
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
item = self.context.pair_first(self.builder, val, ty)
return self.context.get_argument_value(self.builder, ty.first_type,
item)
elif expr.op == 'pair_second':
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
item = self.context.pair_second(self.builder, val, ty)
return self.context.get_argument_value(self.builder,
ty.second_type, item)
elif expr.op in ('getiter', 'iternext'):
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
signature = self.fndesc.calltypes[expr]
impl = self.context.get_function(expr.op, signature)
[fty] = signature.args
castval = self.context.cast(self.builder, val, ty, fty)
res = impl(self.builder, (castval, ))
return self.context.cast(self.builder, res, signature.return_type,
resty)
elif expr.op == 'exhaust_iter':
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
# If we have a heterogenous tuple, we needn't do anything,
# and we can't iterate over it anyway.
if isinstance(ty, types.Tuple):
return val
itemty = ty.iterator_type.yield_type
tup = self.context.get_constant_undef(resty)
pairty = types.Pair(itemty, types.boolean)
getiter_sig = typing.signature(ty.iterator_type, ty)
getiter_impl = self.context.get_function('getiter', getiter_sig)
iternext_sig = typing.signature(pairty, ty.iterator_type)
iternext_impl = self.context.get_function('iternext', iternext_sig)
iterobj = getiter_impl(self.builder, (val, ))
excid = self.add_exception(ValueError)
# We call iternext() as many times as desired (`expr.count`).
for i in range(expr.count):
pair = iternext_impl(self.builder, (iterobj, ))
is_valid = self.context.pair_second(self.builder, pair, pairty)
with cgutils.if_unlikely(self.builder,
self.builder.not_(is_valid)):
self.context.return_user_exc(self.builder, excid)
item = self.context.pair_first(self.builder, pair, pairty)
tup = self.builder.insert_value(tup, item, i)
# Call iternext() once more to check that the iterator
# is exhausted.
pair = iternext_impl(self.builder, (iterobj, ))
is_valid = self.context.pair_second(self.builder, pair, pairty)
with cgutils.if_unlikely(self.builder, is_valid):
self.context.return_user_exc(self.builder, excid)
return tup
elif expr.op == "getattr":
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
if isinstance(resty, types.BoundFunction):
# if we are getting out a method, assume we have typed this
# properly and just build a bound function object
res = self.context.get_bound_function(self.builder, val, ty)
else:
impl = self.context.get_attribute(val, ty, expr.attr)
if impl is None:
# ignore the attribute
res = self.context.get_dummy_value()
else:
res = impl(self.context, self.builder, ty, val, expr.attr)
return res
elif expr.op == "static_getitem":
baseval = self.loadvar(expr.value.name)
indexval = self.context.get_constant(types.intp, expr.index)
if cgutils.is_struct(baseval.type):
# Statically extract the given element from the structure
# (structures aren't dynamically indexable).
return self.builder.extract_value(baseval, expr.index)
else:
# Fall back on the generic getitem() implementation
# for this type.
signature = typing.signature(resty,
self.typeof(expr.value.name),
types.intp)
impl = self.context.get_function("getitem", signature)
argvals = (baseval, indexval)
res = impl(self.builder, argvals)
return self.context.cast(self.builder, res,
signature.return_type, resty)
elif expr.op == "getitem":
baseval = self.loadvar(expr.value.name)
indexval = self.loadvar(expr.index.name)
signature = self.fndesc.calltypes[expr]
impl = self.context.get_function("getitem", signature)
argvals = (baseval, indexval)
argtyps = (self.typeof(expr.value.name),
self.typeof(expr.index.name))
castvals = [
self.context.cast(self.builder, av, at, ft)
for av, at, ft in zip(argvals, argtyps, signature.args)
]
res = impl(self.builder, castvals)
return self.context.cast(self.builder, res, signature.return_type,
resty)
elif expr.op == "build_tuple":
itemvals = [self.loadvar(i.name) for i in expr.items]
itemtys = [self.typeof(i.name) for i in expr.items]
castvals = [
self.context.cast(self.builder, val, fromty, toty)
for val, toty, fromty in zip(itemvals, resty, itemtys)
]
tup = self.context.get_constant_undef(resty)
for i in range(len(castvals)):
tup = self.builder.insert_value(tup, castvals[i], i)
return tup
elif expr.op == "cast":
val = self.loadvar(expr.value.name)
ty = self.typeof(expr.value.name)
castval = self.context.cast(self.builder, val, ty, resty)
return castval
raise NotImplementedError(expr)
def generic(self, args, kws):
assert not kws
return signature(types.none, *args)
def generic(self, args, kws):
assert not kws
assert len(args) in [2, 3]
return signature(types.float64, *args)
def numpy_multiply(context, builder, sig, args):
dtype = sig.args[0].dtype
coresig = typing.signature(types.Any, dtype, dtype)
core = context.get_function("*", coresig)
imp = numpy_binary_ufunc(core)
return imp(context, builder, sig, args)
def generic(self, args, kws):
assert not kws
assert len(args) == 3
# args: out_arr, in_arr, value
return signature(types.none, *args)
def generic(self, args, kws):
assert not kws
assert len(args)==3
return signature(string_array_type, *args)
def generic(self, args, kws):
assert not kws
assert len(args) == 1
# arg: in_arr
return signature(_expand_integer(args[0].dtype), *args)
def generic(self, args, kws):
assert not kws
assert len(args) == 2
# args: str_arr, pat
return signature(types.Array(types.boolean, 1, 'C'), *args)
def generic(self, args, kws):
assert not kws
assert len(args) == 1
return signature(types.intp, *args)
def generic(self, args, kws):
assert not kws
return signature(types.none, *unliteral_all(args))
def generic(self, args, kws):
assert not kws
assert len(args) == 1
return signature(array_datetime_date, *args)
def get_attribute(self, val, typ, attr):
if isinstance(typ, types.Record):
# Implement get attribute for records
self.sentry_record_alignment(typ, attr)
offset = typ.offset(attr)
elemty = typ.typeof(attr)
if isinstance(elemty, types.NestedArray):
# Inside a structured type only the array data is stored, so we
# create an array structure to point to that data.
aryty = arrayobj.make_array(elemty)
@impl_attribute(typ, attr, elemty)
def imp(context, builder, typ, val):
ary = aryty(context, builder)
dtype = elemty.dtype
newshape = [
self.get_constant(types.intp, s) for s in elemty.shape
]
newstrides = [
self.get_constant(types.intp, s)
for s in elemty.strides
]
newdata = cgutils.get_record_member(
builder, val, offset, self.get_data_type(dtype))
arrayobj.populate_array(
ary,
data=newdata,
shape=cgutils.pack_array(builder, newshape),
strides=cgutils.pack_array(builder, newstrides),
itemsize=context.get_constant(types.intp, elemty.size),
meminfo=None,
parent=None,
)
res = ary._getvalue()
return impl_ret_borrowed(context, builder, typ, res)
else:
@impl_attribute(typ, attr, elemty)
def imp(context, builder, typ, val):
dptr = cgutils.get_record_member(
builder, val, offset, context.get_data_type(elemty))
align = None if typ.aligned else 1
res = self.unpack_value(builder, elemty, dptr, align)
return impl_ret_borrowed(context, builder, typ, res)
return imp
if isinstance(typ, types.Module):
# Implement getattr for module-level globals.
# We are treating them as constants.
# XXX We shouldn't have to retype this
attrty = self.typing_context.resolve_module_constants(typ, attr)
if attrty is not None and not isinstance(attrty, types.Dummy):
pyval = getattr(typ.pymod, attr)
llval = self.get_constant(attrty, pyval)
@impl_attribute(typ, attr, attrty)
def imp(context, builder, typ, val):
return impl_ret_borrowed(context, builder, attrty, llval)
return imp
# No implementation required for dummies (functions, modules...),
# which are dealt with later
return None
# Lookup specific attribute implementation for this type
overloads = self.attrs[attr]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
pass
# Lookup generic getattr implementation for this type
overloads = self.attrs[None]
try:
return overloads.find(typing.signature(types.Any, typ))
except NotImplementedError:
raise Exception("No definition for lowering %s.%s" % (typ, attr))
outside the context of CUDA-python.
'''
_description_ = '<ptx special value>'
__slots__ = () # don't allocate __dict__
def __new__(cls):
raise NotImplementedError("%s is not instantiable" % cls)
def __repr__(self):
return self._description_
#-------------------------------------------------------------------------------
# SREG
SREG_SIGNATURE = typing.signature(types.int32)
class threadIdx(Stub):
'''threadIdx.{x, y, z}
'''
_description_ = '<threadIdx.{x,y,z}>'
x = macro.Macro('tid.x', SREG_SIGNATURE)
y = macro.Macro('tid.y', SREG_SIGNATURE)
z = macro.Macro('tid.z', SREG_SIGNATURE)
class blockIdx(Stub):
'''blockIdx.{x, y, z}
'''
def generic(self, args, kws):
assert not kws
assert len(args) == 1
return signature(args[0], *unliteral_all(args))
def is_true(self, builder, typ, val):
impl = self.get_function(bool, typing.signature(types.boolean, typ))
return impl(builder, (val, ))