880 lines
34 KiB
Python
880 lines
34 KiB
Python
import functools
|
|
import operator
|
|
import warnings
|
|
|
|
import numpy as np
|
|
|
|
from numba import jit, typeof
|
|
from numba.core import cgutils, types, serialize, sigutils, errors
|
|
from numba.core.extending import (is_jitted, overload_attribute,
|
|
overload_method, register_jitable,
|
|
intrinsic)
|
|
from numba.core.typing import npydecl
|
|
from numba.core.typing.templates import AbstractTemplate, signature
|
|
from numba.cpython.unsafe.tuple import tuple_setitem
|
|
from numba.np.ufunc import _internal
|
|
from numba.np.ufunc.ufunc_base import UfuncBase, UfuncLowererBase
|
|
from numba.parfors import array_analysis
|
|
from numba.np.ufunc import ufuncbuilder
|
|
from numba.np import numpy_support
|
|
from typing import Callable
|
|
from llvmlite import ir
|
|
from numba.core.compiler_lock import global_compiler_lock
|
|
|
|
|
|
class UfuncAtIterator:
|
|
|
|
def __init__(self, ufunc, a, a_ty, indices, indices_ty, b=None, b_ty=None):
|
|
self.ufunc = ufunc
|
|
self.a = a
|
|
self.a_ty = a_ty
|
|
self.indices = indices
|
|
self.indices_ty = indices_ty
|
|
self.b = b
|
|
self.b_ty = b_ty
|
|
|
|
def run(self, context, builder):
|
|
self._prepare(context, builder)
|
|
loop_indices, _ = self.indexer.begin_loops()
|
|
self._call_ufunc(context, builder, loop_indices)
|
|
self.indexer.end_loops()
|
|
|
|
def need_advanced_indexing(self):
|
|
return isinstance(self.indices_ty, types.BaseTuple)
|
|
|
|
def _prepare(self, context, builder):
|
|
from numba.np.arrayobj import normalize_indices, FancyIndexer
|
|
|
|
a, indices = self.a, self.indices
|
|
a_ty, indices_ty = self.a_ty, self.indices_ty
|
|
|
|
zero = context.get_value_type(types.intp)(0)
|
|
|
|
if self.b is not None:
|
|
self.b_indice = cgutils.alloca_once_value(builder, zero)
|
|
|
|
if self.need_advanced_indexing():
|
|
indices = cgutils.unpack_tuple(builder, indices,
|
|
count=len(indices_ty))
|
|
index_types = indices_ty.types
|
|
index_types, indices = normalize_indices(context, builder,
|
|
index_types, indices)
|
|
else:
|
|
indices = (indices,)
|
|
index_types = (indices_ty,)
|
|
index_types, indices = normalize_indices(context, builder,
|
|
index_types, indices)
|
|
|
|
self.indexer = FancyIndexer(context, builder, a_ty, a,
|
|
index_types, indices)
|
|
self.indexer.prepare()
|
|
self.cres = self._compile_ufunc(context, builder)
|
|
|
|
def _load_val(self, context, builder, loop_indices, array, array_ty):
|
|
from numba.np.arrayobj import load_item
|
|
shapes = cgutils.unpack_tuple(builder, array.shape)
|
|
strides = cgutils.unpack_tuple(builder, array.strides)
|
|
data = array.data
|
|
|
|
ptr = cgutils.get_item_pointer2(context, builder, data, shapes, strides,
|
|
array_ty.layout, loop_indices)
|
|
val = load_item(context, builder, array_ty, ptr)
|
|
return ptr, val
|
|
|
|
def _load_flat(self, context, builder, indices, array, array_ty):
|
|
idx = builder.load(indices)
|
|
sig = array_ty.dtype(array_ty, types.intp)
|
|
impl = context.get_function(operator.getitem, sig)
|
|
val = impl(builder, (array, idx))
|
|
|
|
# increment indices
|
|
one = context.get_value_type(types.intp)(1)
|
|
idx = builder.add(idx, one)
|
|
builder.store(idx, indices)
|
|
|
|
return None, val
|
|
|
|
def _store_val(self, context, builder, array, array_ty, ptr, val):
|
|
from numba.np.arrayobj import store_item
|
|
fromty = self.cres.signature.return_type
|
|
toty = array_ty.dtype
|
|
val = context.cast(builder, val, fromty, toty)
|
|
store_item(context, builder, array_ty, val, ptr)
|
|
|
|
def _compile_ufunc(self, context, builder):
|
|
ufunc = self.ufunc.key[0]
|
|
|
|
if self.b is None:
|
|
sig = (self.a_ty.dtype,)
|
|
else:
|
|
sig = (self.a_ty.dtype, self.b_ty.dtype)
|
|
|
|
cres = ufunc.add(sig)
|
|
context.add_linking_libs((cres.library,))
|
|
return cres
|
|
|
|
def _call_ufunc(self, context, builder, loop_indices):
|
|
cres = self.cres
|
|
a, a_ty = self.a, self.a_ty
|
|
|
|
ptr, val = self._load_val(context, builder, loop_indices, a, a_ty)
|
|
|
|
if self.b is None:
|
|
args = (val,)
|
|
else:
|
|
b, b_ty, b_idx = self.b, self.b_ty, self.b_indice
|
|
_, val_b = self._load_flat(context, builder, b_idx, b, b_ty)
|
|
args = (val, val_b)
|
|
|
|
res = context.call_internal(builder, cres.fndesc, cres.signature,
|
|
args)
|
|
self._store_val(context, builder, a, a_ty, ptr, res)
|
|
|
|
|
|
def make_dufunc_kernel(_dufunc):
|
|
from numba.np import npyimpl
|
|
|
|
class DUFuncKernel(npyimpl._Kernel):
|
|
"""
|
|
npyimpl._Kernel subclass responsible for lowering a DUFunc kernel
|
|
(element-wise function) inside a broadcast loop (which is
|
|
generated by npyimpl.numpy_ufunc_kernel()).
|
|
"""
|
|
dufunc = _dufunc
|
|
|
|
def __init__(self, context, builder, outer_sig):
|
|
super().__init__(context, builder, outer_sig)
|
|
self.inner_sig, self.cres = self.dufunc.find_ewise_function(
|
|
outer_sig.args)
|
|
|
|
DUFuncKernel.__name__ += _dufunc.ufunc.__name__
|
|
return DUFuncKernel
|
|
|
|
|
|
class DUFuncLowerer(UfuncLowererBase):
|
|
'''Callable class responsible for lowering calls to a specific DUFunc.
|
|
'''
|
|
def __init__(self, dufunc):
|
|
from numba.np import npyimpl
|
|
super().__init__(dufunc,
|
|
make_dufunc_kernel,
|
|
npyimpl.numpy_ufunc_kernel)
|
|
|
|
|
|
class DUFunc(serialize.ReduceMixin, _internal._DUFunc, UfuncBase):
|
|
"""
|
|
Dynamic universal function (DUFunc) intended to act like a normal
|
|
Numpy ufunc, but capable of call-time (just-in-time) compilation
|
|
of fast loops specialized to inputs.
|
|
"""
|
|
# NOTE: __base_kwargs must be kept in synch with the kwlist in
|
|
# _internal.c:dufunc_init()
|
|
__base_kwargs = set(('identity', '_keepalive', 'nin', 'nout'))
|
|
|
|
def __init__(self, py_func, identity=None, cache=False, targetoptions={}):
|
|
if is_jitted(py_func):
|
|
py_func = py_func.py_func
|
|
with ufuncbuilder._suppress_deprecation_warning_nopython_not_supplied():
|
|
dispatcher = jit(_target='npyufunc',
|
|
cache=cache,
|
|
**targetoptions)(py_func)
|
|
self._initialize(dispatcher, identity)
|
|
functools.update_wrapper(self, py_func)
|
|
|
|
def _initialize(self, dispatcher, identity):
|
|
identity = ufuncbuilder.parse_identity(identity)
|
|
super(DUFunc, self).__init__(dispatcher, identity=identity)
|
|
# Loop over a copy of the keys instead of the keys themselves,
|
|
# since we're changing the dictionary while looping.
|
|
self.reorderable = (identity != _internal.PyUFunc_None)
|
|
self.__name__ = dispatcher.py_func.__name__
|
|
self.__doc__ = dispatcher.py_func.__doc__
|
|
self._lower_me = DUFuncLowerer(self)
|
|
self._install_cg()
|
|
self._install_type()
|
|
|
|
def _reduce_states(self):
|
|
"""
|
|
NOTE: part of ReduceMixin protocol
|
|
"""
|
|
siglist = list(self._dispatcher.overloads.keys())
|
|
return dict(
|
|
dispatcher=self._dispatcher,
|
|
identity=self.identity,
|
|
frozen=self._frozen,
|
|
siglist=siglist,
|
|
)
|
|
|
|
@classmethod
|
|
def _rebuild(cls, dispatcher, identity, frozen, siglist):
|
|
"""
|
|
NOTE: part of ReduceMixin protocol
|
|
"""
|
|
self = _internal._DUFunc.__new__(cls)
|
|
self._initialize(dispatcher, identity)
|
|
# Re-add signatures
|
|
for sig in siglist:
|
|
self.add(sig)
|
|
if frozen:
|
|
self.disable_compile()
|
|
return self
|
|
|
|
def build_ufunc(self):
|
|
"""
|
|
For compatibility with the various *UFuncBuilder classes.
|
|
"""
|
|
return self
|
|
|
|
@property
|
|
def targetoptions(self):
|
|
return self._dispatcher.targetoptions
|
|
|
|
@property
|
|
def nin(self):
|
|
return self.ufunc.nin
|
|
|
|
@property
|
|
def nout(self):
|
|
return self.ufunc.nout
|
|
|
|
@property
|
|
def nargs(self):
|
|
return self.ufunc.nargs
|
|
|
|
@property
|
|
def ntypes(self):
|
|
return self.ufunc.ntypes
|
|
|
|
@property
|
|
def types(self):
|
|
return self.ufunc.types
|
|
|
|
@property
|
|
def identity(self):
|
|
return self.ufunc.identity
|
|
|
|
@property
|
|
def signature(self):
|
|
return self.ufunc.signature
|
|
|
|
def disable_compile(self):
|
|
"""
|
|
Disable the compilation of new signatures at call time.
|
|
"""
|
|
# If disabling compilation then there must be at least one signature
|
|
assert len(self._dispatcher.overloads) > 0
|
|
self._frozen = True
|
|
|
|
def add(self, sig):
|
|
"""
|
|
Compile the DUFunc for the given signature.
|
|
"""
|
|
args, return_type = sigutils.normalize_signature(sig)
|
|
return self._compile_for_argtys(args, return_type)
|
|
|
|
def __call__(self, *args, **kws):
|
|
"""
|
|
Allow any argument that has overridden __array_ufunc__ (NEP-18)
|
|
to take control of DUFunc.__call__.
|
|
"""
|
|
default = numpy_support.np.ndarray.__array_ufunc__
|
|
|
|
for arg in args + tuple(kws.values()):
|
|
if getattr(type(arg), "__array_ufunc__", default) is not default:
|
|
output = arg.__array_ufunc__(self, "__call__", *args, **kws)
|
|
if output is not NotImplemented:
|
|
return output
|
|
else:
|
|
return super().__call__(*args, **kws)
|
|
|
|
def _compile_for_args(self, *args, **kws):
|
|
nin = self.ufunc.nin
|
|
if kws:
|
|
if 'out' in kws:
|
|
out = kws.pop('out')
|
|
args += (out,)
|
|
if kws:
|
|
raise TypeError("unexpected keyword arguments to ufunc: %s"
|
|
% ", ".join(repr(k) for k in sorted(kws)))
|
|
|
|
args_len = len(args)
|
|
assert (args_len == nin) or (args_len == nin + self.ufunc.nout)
|
|
assert not kws
|
|
argtys = []
|
|
for arg in args[:nin]:
|
|
argty = typeof(arg)
|
|
if isinstance(argty, types.Array):
|
|
argty = argty.dtype
|
|
else:
|
|
# To avoid a mismatch in how Numba types scalar values as
|
|
# opposed to Numpy, we need special logic for scalars.
|
|
# For example, on 64-bit systems, numba.typeof(3) => int32, but
|
|
# np.array(3).dtype => int64.
|
|
|
|
# Note: this will not handle numpy "duckarrays" correctly,
|
|
# including but not limited to those implementing `__array__`
|
|
# and `__array_ufunc__`.
|
|
argty = numpy_support.map_arrayscalar_type(arg)
|
|
argtys.append(argty)
|
|
return self._compile_for_argtys(tuple(argtys))
|
|
|
|
@global_compiler_lock
|
|
def _compile_for_argtys(self, argtys, return_type=None):
|
|
"""
|
|
Given a tuple of argument types (these should be the array
|
|
dtypes, and not the array types themselves), compile the
|
|
element-wise function for those inputs, generate a UFunc loop
|
|
wrapper, and register the loop with the Numpy ufunc object for
|
|
this DUFunc.
|
|
"""
|
|
if self._frozen:
|
|
raise RuntimeError("compilation disabled for %s" % (self,))
|
|
assert isinstance(argtys, tuple)
|
|
if return_type is None:
|
|
sig = argtys
|
|
else:
|
|
sig = return_type(*argtys)
|
|
|
|
for k, cres in self._dispatcher.overloads.items():
|
|
if argtys == k.args:
|
|
msg = ("Compilation requested for previously compiled argument"
|
|
f" types ({argtys}). This has no effect and perhaps "
|
|
"indicates a bug in the calling code (compiling a "
|
|
"ufunc more than once for the same signature")
|
|
warnings.warn(msg, errors.NumbaWarning)
|
|
return cres
|
|
|
|
cres, argtys, return_type = ufuncbuilder._compile_element_wise_function(
|
|
self._dispatcher, self.targetoptions, sig)
|
|
actual_sig = ufuncbuilder._finalize_ufunc_signature(
|
|
cres, argtys, return_type)
|
|
dtypenums, ptr, env = ufuncbuilder._build_element_wise_ufunc_wrapper(
|
|
cres, actual_sig)
|
|
self._add_loop(int(ptr), dtypenums)
|
|
self._keepalive.append((ptr, cres.library, env))
|
|
self._lower_me.libs.append(cres.library)
|
|
return cres
|
|
|
|
def match_signature(self, ewise_types, sig):
|
|
return sig.args == ewise_types
|
|
|
|
def _install_ufunc_attributes(self, template) -> None:
|
|
|
|
def get_attr_fn(attr: str) -> Callable:
|
|
|
|
def impl(ufunc):
|
|
val = getattr(ufunc.key[0], attr)
|
|
return lambda ufunc: val
|
|
return impl
|
|
|
|
# ntypes/types needs "at" to be a BoundFunction rather than a Function
|
|
# But this fails as it cannot a weak reference to an ufunc due to NumPy
|
|
# not setting the "tp_weaklistoffset" field. See:
|
|
# https://github.com/numpy/numpy/blob/7fc72776b972bfbfdb909e4b15feb0308cf8adba/numpy/core/src/umath/ufunc_object.c#L6968-L6983 # noqa: E501
|
|
|
|
at = types.Function(template)
|
|
attributes = ('nin', 'nout', 'nargs', # 'ntypes', # 'types',
|
|
'identity', 'signature')
|
|
for attr in attributes:
|
|
attr_fn = get_attr_fn(attr)
|
|
overload_attribute(at, attr)(attr_fn)
|
|
|
|
def _install_ufunc_methods(self, template) -> None:
|
|
self._install_ufunc_reduce(template)
|
|
self._install_ufunc_at(template)
|
|
|
|
def _install_ufunc_at(self, template) -> None:
|
|
at = types.Function(template)
|
|
|
|
@overload_method(at, 'at')
|
|
def ol_at(ufunc, a, indices, b=None):
|
|
warnings.warn("ufunc.at feature is experimental",
|
|
category=errors.NumbaExperimentalFeatureWarning)
|
|
|
|
if not isinstance(a, types.Array):
|
|
msg = 'The first argument "a" must be array-like'
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
indices_arr = isinstance(indices, types.Array)
|
|
indices_list = isinstance(indices, types.List)
|
|
indices_tuple = isinstance(indices, types.Tuple)
|
|
indices_slice = isinstance(indices, types.SliceType)
|
|
indices_scalar = not (indices_arr or indices_slice or indices_tuple)
|
|
indices_empty_tuple = indices_tuple and len(indices) == 0
|
|
b_array = isinstance(b, (types.Array, types.Sequence, types.List,
|
|
types.Tuple))
|
|
b_none = cgutils.is_nonelike(b)
|
|
b_scalar = not (b_array or b_none)
|
|
need_cast = any([indices_list])
|
|
|
|
nin = self.ufunc.nin
|
|
|
|
# missing second argument?
|
|
if nin == 2 and cgutils.is_nonelike(b):
|
|
raise errors.TypingError('second operand needed for ufunc')
|
|
|
|
# extra second argument
|
|
if nin == 1 and not cgutils.is_nonelike(b):
|
|
msg = 'second operand provided when ufunc is unary'
|
|
raise errors.TypingError(msg)
|
|
|
|
if cgutils.is_nonelike(b):
|
|
self.add((a.dtype,))
|
|
elif b_scalar:
|
|
self.add((a.dtype, b))
|
|
else:
|
|
self.add((a.dtype, b.dtype))
|
|
|
|
def apply_ufunc_codegen(context, builder, sig, args):
|
|
from numba.np.arrayobj import make_array
|
|
|
|
if len(args) == 4:
|
|
_, aty, idxty, bty = sig.args
|
|
_, a, indices, b = args
|
|
else:
|
|
_, aty, idxty, bty = sig.args + (None,)
|
|
_, a, indices, b = args + (None,)
|
|
|
|
a = make_array(aty)(context, builder, a)
|
|
at_iter = UfuncAtIterator(ufunc, a, aty, indices, idxty, b, bty)
|
|
at_iter.run(context, builder)
|
|
|
|
@intrinsic
|
|
def apply_a_b_ufunc(typingctx, ufunc, a, indices, b):
|
|
sig = types.none(ufunc, a, indices, b)
|
|
return sig, apply_ufunc_codegen
|
|
|
|
@intrinsic
|
|
def apply_a_ufunc(typingctx, ufunc, a, indices):
|
|
sig = types.none(ufunc, a, indices)
|
|
return sig, apply_ufunc_codegen
|
|
|
|
def impl_cast(ufunc, a, indices, b=None):
|
|
if b_none:
|
|
return ufunc.at(a, np.asarray(indices))
|
|
else:
|
|
return ufunc.at(a,
|
|
np.asarray(indices),
|
|
np.asarray(b))
|
|
|
|
def impl_generic(ufunc, a, indices, b=None):
|
|
if b_none:
|
|
apply_a_ufunc(ufunc, a, indices,)
|
|
else:
|
|
b_ = np.asarray(b)
|
|
a_ = a[indices]
|
|
b_ = np.broadcast_to(b_, a_.shape)
|
|
apply_a_b_ufunc(ufunc, a, indices, b_.flat)
|
|
|
|
def impl_indices_empty_b_scalar(ufunc, a, indices, b=None):
|
|
a[()] = ufunc(a[()], b)
|
|
|
|
def impl_scalar_scalar(ufunc, a, indices, b=None):
|
|
if b_none:
|
|
a[indices] = ufunc(a[indices])
|
|
else:
|
|
a[indices] = ufunc(a[indices], b)
|
|
|
|
if need_cast:
|
|
return impl_cast
|
|
elif indices_empty_tuple and b_scalar:
|
|
return impl_indices_empty_b_scalar
|
|
elif indices_scalar and b_scalar:
|
|
return impl_scalar_scalar
|
|
else:
|
|
return impl_generic
|
|
|
|
def _install_ufunc_reduce(self, template) -> None:
|
|
at = types.Function(template)
|
|
|
|
@overload_method(at, 'reduce')
|
|
def ol_reduce(ufunc, array, axis=0, dtype=None, initial=None):
|
|
|
|
warnings.warn("ufunc.reduce feature is experimental",
|
|
category=errors.NumbaExperimentalFeatureWarning)
|
|
|
|
if not isinstance(array, types.Array):
|
|
msg = 'The first argument "array" must be array-like'
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
axis_int_tuple = isinstance(axis, types.UniTuple) and \
|
|
isinstance(axis.dtype, types.Integer)
|
|
axis_empty_tuple = isinstance(axis, types.Tuple) and len(axis) == 0
|
|
axis_none = cgutils.is_nonelike(axis)
|
|
|
|
identity_none = self.ufunc.identity is None
|
|
ufunc_name = self.ufunc.__name__
|
|
|
|
# In NumPy, a ufunc is reorderable if its identity type is **not**
|
|
# PyUfunc_None.
|
|
if not self.reorderable and axis_int_tuple and len(axis) > 1:
|
|
msg = (f"reduction operation '{ufunc_name}' is not "
|
|
"reorderable, so at most one axis may be specified")
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
tup_init = (0,) * (array.ndim)
|
|
tup_init_m1 = (0,) * (array.ndim - 1)
|
|
nb_dtype = array.dtype if cgutils.is_nonelike(dtype) else dtype
|
|
identity = self.identity
|
|
|
|
id_none = cgutils.is_nonelike(identity)
|
|
init_none = cgutils.is_nonelike(initial)
|
|
|
|
@register_jitable
|
|
def tuple_slice(tup, pos):
|
|
# Same as
|
|
# tup = tup[0 : pos] + tup[pos + 1:]
|
|
s = tup_init_m1
|
|
i = 0
|
|
for j, e in enumerate(tup):
|
|
if j == pos:
|
|
continue
|
|
s = tuple_setitem(s, i, e)
|
|
i += 1
|
|
return s
|
|
|
|
@register_jitable
|
|
def tuple_slice_append(tup, pos, val):
|
|
# Same as
|
|
# tup = tup[0 : pos] + val + tup[pos + 1:]
|
|
s = tup_init
|
|
i, j, sz = 0, 0, len(s)
|
|
while j < sz:
|
|
if j == pos:
|
|
s = tuple_setitem(s, j, val)
|
|
else:
|
|
e = tup[i]
|
|
s = tuple_setitem(s, j, e)
|
|
i += 1
|
|
j += 1
|
|
return s
|
|
|
|
@intrinsic
|
|
def compute_flat_idx(typingctx, strides, itemsize, idx, axis):
|
|
sig = types.intp(strides, itemsize, idx, axis)
|
|
len_idx = len(idx)
|
|
|
|
def gen_block(builder, block_pos, block_name, bb_end, args):
|
|
strides, _, idx, _ = args
|
|
bb = builder.append_basic_block(name=block_name)
|
|
|
|
with builder.goto_block(bb):
|
|
zero = ir.IntType(64)(0)
|
|
flat_idx = zero
|
|
|
|
if block_pos == 0:
|
|
for i in range(1, len_idx):
|
|
stride = builder.extract_value(strides, i - 1)
|
|
idx_i = builder.extract_value(idx, i)
|
|
m = builder.mul(stride, idx_i)
|
|
flat_idx = builder.add(flat_idx, m)
|
|
elif 0 < block_pos < len_idx - 1:
|
|
for i in range(0, block_pos):
|
|
stride = builder.extract_value(strides, i)
|
|
idx_i = builder.extract_value(idx, i)
|
|
m = builder.mul(stride, idx_i)
|
|
flat_idx = builder.add(flat_idx, m)
|
|
|
|
for i in range(block_pos + 1, len_idx):
|
|
stride = builder.extract_value(strides, i - 1)
|
|
idx_i = builder.extract_value(idx, i)
|
|
m = builder.mul(stride, idx_i)
|
|
flat_idx = builder.add(flat_idx, m)
|
|
else:
|
|
for i in range(0, len_idx - 1):
|
|
stride = builder.extract_value(strides, i)
|
|
idx_i = builder.extract_value(idx, i)
|
|
m = builder.mul(stride, idx_i)
|
|
flat_idx = builder.add(flat_idx, m)
|
|
|
|
builder.branch(bb_end)
|
|
|
|
return bb, flat_idx
|
|
|
|
def codegen(context, builder, sig, args):
|
|
strides, itemsize, idx, axis = args
|
|
|
|
bb = builder.basic_block
|
|
switch_end = builder.append_basic_block(name='axis_end')
|
|
l = []
|
|
for i in range(len_idx):
|
|
block, flat_idx = gen_block(builder, i, f"axis_{i}",
|
|
switch_end, args)
|
|
l.append((block, flat_idx))
|
|
|
|
with builder.goto_block(bb):
|
|
switch = builder.switch(axis, l[-1][0])
|
|
for i in range(len_idx):
|
|
switch.add_case(i, l[i][0])
|
|
|
|
builder.position_at_end(switch_end)
|
|
phi = builder.phi(l[0][1].type)
|
|
for block, value in l:
|
|
phi.add_incoming(value, block)
|
|
return builder.sdiv(phi, itemsize)
|
|
|
|
return sig, codegen
|
|
|
|
@register_jitable
|
|
def fixup_axis(axis, ndim):
|
|
ax = axis
|
|
for i in range(len(axis)):
|
|
val = axis[i] + ndim if axis[i] < 0 else axis[i]
|
|
ax = tuple_setitem(ax, i, val)
|
|
return ax
|
|
|
|
@register_jitable
|
|
def find_min(tup):
|
|
idx, e = 0, tup[0]
|
|
for i in range(len(tup)):
|
|
if tup[i] < e:
|
|
idx, e = i, tup[i]
|
|
return idx, e
|
|
|
|
def impl_1d(ufunc, array, axis=0, dtype=None, initial=None):
|
|
if identity_none and initial is None and len(array) == 0:
|
|
msg = ('zero-size array to reduction operation '
|
|
f'{ufunc_name} which has no identity')
|
|
raise ValueError(msg)
|
|
|
|
start = 0
|
|
if init_none and id_none:
|
|
start = 1
|
|
r = array[0]
|
|
elif init_none:
|
|
r = identity
|
|
else:
|
|
r = initial
|
|
|
|
sz = array.shape[0]
|
|
for i in range(start, sz):
|
|
r = ufunc(r, array[i])
|
|
return r
|
|
|
|
def impl_nd_axis_int(ufunc,
|
|
array,
|
|
axis=0,
|
|
dtype=None,
|
|
initial=None):
|
|
if axis is None:
|
|
raise ValueError("'axis' must be specified")
|
|
|
|
if axis < 0:
|
|
axis += array.ndim
|
|
|
|
if axis < 0 or axis >= array.ndim:
|
|
raise ValueError("Invalid axis")
|
|
|
|
if identity_none and initial is None and array.shape[axis] == 0:
|
|
msg = ('zero-size array to reduction operation '
|
|
f'{ufunc_name} which has no identity')
|
|
raise ValueError(msg)
|
|
|
|
# create result array
|
|
shape = tuple_slice(array.shape, axis)
|
|
|
|
if initial is None and identity is None:
|
|
r = np.empty(shape, dtype=nb_dtype)
|
|
for idx, _ in np.ndenumerate(r):
|
|
# shape[0:axis] + 0 + shape[axis:]
|
|
result_idx = tuple_slice_append(idx, axis, 0)
|
|
r[idx] = array[result_idx]
|
|
elif initial is None and identity is not None:
|
|
# Checking if identity is not none is redundant but required
|
|
# compile this block
|
|
r = np.full(shape, fill_value=identity, dtype=nb_dtype)
|
|
else:
|
|
r = np.full(shape, fill_value=initial, dtype=nb_dtype)
|
|
|
|
# One approach to implement reduce is to remove the axis index
|
|
# from the indexing tuple returned by "np.ndenumerate". For
|
|
# instance, if idx = (X, Y, Z) and axis=1, the result index
|
|
# is (X, Y).
|
|
# Another way is to compute the result index using strides,
|
|
# which is faster than manipulating tuples.
|
|
view = r.ravel()
|
|
if initial is None and identity is None:
|
|
for idx, val in np.ndenumerate(array):
|
|
if idx[axis] == 0:
|
|
continue
|
|
else:
|
|
flat_pos = compute_flat_idx(r.strides, r.itemsize,
|
|
idx, axis)
|
|
lhs, rhs = view[flat_pos], val
|
|
view[flat_pos] = ufunc(lhs, rhs)
|
|
else:
|
|
for idx, val in np.ndenumerate(array):
|
|
if initial is None and identity is None and \
|
|
idx[axis] == 0:
|
|
continue
|
|
flat_pos = compute_flat_idx(r.strides, r.itemsize,
|
|
idx, axis)
|
|
lhs, rhs = view[flat_pos], val
|
|
view[flat_pos] = ufunc(lhs, rhs)
|
|
return r
|
|
|
|
def impl_nd_axis_tuple(ufunc,
|
|
array,
|
|
axis=0,
|
|
dtype=None,
|
|
initial=None):
|
|
axis_ = fixup_axis(axis, array.ndim)
|
|
for i in range(0, len(axis_)):
|
|
if axis_[i] < 0 or axis_[i] >= array.ndim:
|
|
raise ValueError("Invalid axis")
|
|
|
|
for j in range(i + 1, len(axis_)):
|
|
if axis_[i] == axis_[j]:
|
|
raise ValueError("duplicate value in 'axis'")
|
|
|
|
min_idx, min_elem = find_min(axis_)
|
|
r = ufunc.reduce(array,
|
|
axis=min_elem,
|
|
dtype=dtype,
|
|
initial=initial)
|
|
if len(axis) == 1:
|
|
return r
|
|
elif len(axis) == 2:
|
|
return ufunc.reduce(r, axis=axis_[(min_idx + 1) % 2] - 1)
|
|
else:
|
|
ax = axis_tup
|
|
for i in range(len(ax)):
|
|
if i != min_idx:
|
|
ax = tuple_setitem(ax, i, axis_[i])
|
|
return ufunc.reduce(r, axis=ax)
|
|
|
|
def impl_axis_empty_tuple(ufunc,
|
|
array,
|
|
axis=0,
|
|
dtype=None,
|
|
initial=None):
|
|
return array
|
|
|
|
def impl_axis_none(ufunc,
|
|
array,
|
|
axis=0,
|
|
dtype=None,
|
|
initial=None):
|
|
return ufunc.reduce(array, axis_tup, dtype, initial)
|
|
|
|
if array.ndim == 1 and not axis_empty_tuple:
|
|
return impl_1d
|
|
elif axis_empty_tuple:
|
|
# ufunc(array, axis=())
|
|
return impl_axis_empty_tuple
|
|
elif axis_none:
|
|
# ufunc(array, axis=None)
|
|
axis_tup = tuple(range(array.ndim))
|
|
return impl_axis_none
|
|
elif axis_int_tuple:
|
|
# axis is tuple of integers
|
|
# ufunc(array, axis=(1, 2, ...))
|
|
axis_tup = (0,) * (len(axis) - 1)
|
|
return impl_nd_axis_tuple
|
|
elif axis == 0 or isinstance(axis, (types.Integer,
|
|
types.Omitted,
|
|
types.IntegerLiteral)):
|
|
# axis is default value (0) or an integer
|
|
# ufunc(array, axis=0)
|
|
return impl_nd_axis_int
|
|
|
|
def at(self, a, indices, b=None):
|
|
# dynamic compile ufunc.at
|
|
args = (a,) if cgutils.is_nonelike(b) else (a, b)
|
|
argtys = (typeof(arg) for arg in args)
|
|
ewise_types = tuple(arg.dtype if isinstance(arg, types.Array) else arg
|
|
for arg in argtys)
|
|
|
|
if self.find_ewise_function(ewise_types) == (None, None):
|
|
# cannot find a matching function and compilation is disabled
|
|
if self._frozen:
|
|
msg = "compilation disabled for %s.at(...)" % (self,)
|
|
raise RuntimeError(msg)
|
|
|
|
self._compile_for_args(*args)
|
|
|
|
# all good, just dispatch to the function
|
|
if cgutils.is_nonelike(b):
|
|
return super().at(a, indices)
|
|
else:
|
|
return super().at(*(a, indices, b))
|
|
|
|
def _install_type(self, typingctx=None):
|
|
"""Constructs and installs a typing class for a DUFunc object in the
|
|
input typing context. If no typing context is given, then
|
|
_install_type() installs into the typing context of the
|
|
dispatcher object (should be same default context used by
|
|
jit() and njit()).
|
|
"""
|
|
if typingctx is None:
|
|
typingctx = self._dispatcher.targetdescr.typing_context
|
|
_ty_cls = type('DUFuncTyping_' + self.ufunc.__name__,
|
|
(AbstractTemplate,),
|
|
dict(key=self, generic=self._type_me))
|
|
typingctx.insert_user_function(self, _ty_cls)
|
|
self._install_ufunc_attributes(_ty_cls)
|
|
self._install_ufunc_methods(_ty_cls)
|
|
|
|
def find_ewise_function(self, ewise_types):
|
|
"""
|
|
Given a tuple of element-wise argument types, find a matching
|
|
signature in the dispatcher.
|
|
|
|
Return a 2-tuple containing the matching signature, and
|
|
compilation result. Will return two None's if no matching
|
|
signature was found.
|
|
"""
|
|
if self._frozen:
|
|
# If we cannot compile, coerce to the best matching loop
|
|
loop = numpy_support.ufunc_find_matching_loop(self, ewise_types)
|
|
if loop is None:
|
|
return None, None
|
|
ewise_types = tuple(loop.inputs + loop.outputs)[:len(ewise_types)]
|
|
for sig, cres in self._dispatcher.overloads.items():
|
|
if sig.args == ewise_types:
|
|
return sig, cres
|
|
return None, None
|
|
|
|
def _type_me(self, argtys, kwtys):
|
|
"""
|
|
Implement AbstractTemplate.generic() for the typing class
|
|
built by DUFunc._install_type().
|
|
|
|
Return the call-site signature after either validating the
|
|
element-wise signature or compiling for it.
|
|
"""
|
|
assert not kwtys
|
|
ufunc = self.ufunc
|
|
_handle_inputs_result = npydecl.Numpy_rules_ufunc._handle_inputs(
|
|
ufunc, argtys, kwtys)
|
|
base_types, explicit_outputs, ndims, layout = _handle_inputs_result
|
|
explicit_output_count = len(explicit_outputs)
|
|
if explicit_output_count > 0:
|
|
ewise_types = tuple(base_types[:-len(explicit_outputs)])
|
|
else:
|
|
ewise_types = tuple(base_types)
|
|
sig, cres = self.find_ewise_function(ewise_types)
|
|
if sig is None:
|
|
# Matching element-wise signature was not found; must
|
|
# compile.
|
|
if self._frozen:
|
|
raise errors.NumbaTypeError("cannot call %s with types %s"
|
|
% (self, argtys))
|
|
self._compile_for_argtys(ewise_types)
|
|
sig, cres = self.find_ewise_function(ewise_types)
|
|
assert sig is not None
|
|
if explicit_output_count > 0:
|
|
outtys = list(explicit_outputs)
|
|
elif ufunc.nout == 1:
|
|
if ndims > 0:
|
|
outtys = [types.Array(sig.return_type, ndims, layout)]
|
|
else:
|
|
outtys = [sig.return_type]
|
|
else:
|
|
raise errors.NumbaNotImplementedError("typing gufuncs (nout > 1)")
|
|
outtys.extend(argtys)
|
|
return signature(*outtys)
|
|
|
|
|
|
array_analysis.MAP_TYPES.append(DUFunc)
|