1690 lines
61 KiB
Python
1690 lines
61 KiB
Python
|
"""Codegen for functions used as kernels in NumPy functions
|
||
|
|
|
||
|
|
Typically, the kernels of several ufuncs that can't map directly to
|
||
|
|
Python builtins
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
import math
|
||
|
|
|
||
|
|
import llvmlite.ir
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from numba.core.extending import overload
|
||
|
|
from numba.core.imputils import impl_ret_untracked
|
||
|
|
from numba.core import typing, types, errors, lowering, cgutils
|
||
|
|
from numba.core.extending import register_jitable
|
||
|
|
from numba.np import npdatetime
|
||
|
|
from numba.cpython import cmathimpl, mathimpl, numbers
|
||
|
|
|
||
|
|
# some NumPy constants. Note that we could generate some of them using
|
||
|
|
# the math library, but having the values copied from npy_math seems to
|
||
|
|
# yield more accurate results
|
||
|
|
_NPY_LOG2E = 1.442695040888963407359924681001892137 # math.log(math.e, 2)
|
||
|
|
_NPY_LOG10E = 0.434294481903251827651128918916605082 # math.log(math.e, 10)
|
||
|
|
_NPY_LOGE2 = 0.693147180559945309417232121458176568 # math.log(2)
|
||
|
|
|
||
|
|
|
||
|
|
def _check_arity_and_homogeneity(sig, args, arity, return_type = None):
|
||
|
|
"""checks that the following are true:
|
||
|
|
- args and sig.args have arg_count elements
|
||
|
|
- all input types are homogeneous
|
||
|
|
- return type is 'return_type' if provided, otherwise it must be
|
||
|
|
homogeneous with the input types.
|
||
|
|
"""
|
||
|
|
assert len(args) == arity
|
||
|
|
assert len(sig.args) == arity
|
||
|
|
ty = sig.args[0]
|
||
|
|
if return_type is None:
|
||
|
|
return_type = ty
|
||
|
|
# must have homogeneous args
|
||
|
|
if not( all(arg==ty for arg in sig.args) and sig.return_type == return_type):
|
||
|
|
import inspect
|
||
|
|
fname = inspect.currentframe().f_back.f_code.co_name
|
||
|
|
msg = '{0} called with invalid types: {1}'.format(fname, sig)
|
||
|
|
assert False, msg
|
||
|
|
|
||
|
|
|
||
|
|
def _call_func_by_name_with_cast(context, builder, sig, args,
|
||
|
|
func_name, ty=types.float64):
|
||
|
|
# it is quite common in NumPy to have loops implemented as a call
|
||
|
|
# to the double version of the function, wrapped in casts. This
|
||
|
|
# helper function facilitates that.
|
||
|
|
mod = builder.module
|
||
|
|
lty = context.get_argument_type(ty)
|
||
|
|
fnty = llvmlite.ir.FunctionType(lty, [lty]*len(sig.args))
|
||
|
|
fn = cgutils.insert_pure_function(mod, fnty, name=func_name)
|
||
|
|
cast_args = [context.cast(builder, arg, argty, ty)
|
||
|
|
for arg, argty in zip(args, sig.args) ]
|
||
|
|
|
||
|
|
result = builder.call(fn, cast_args)
|
||
|
|
return context.cast(builder, result, types.float64, sig.return_type)
|
||
|
|
|
||
|
|
|
||
|
|
def _dispatch_func_by_name_type(context, builder, sig, args, table, user_name):
|
||
|
|
# for most cases the functions are homogeneous on all their types.
|
||
|
|
# this code dispatches on the first argument type as it is the most useful
|
||
|
|
# for our uses (all cases but ldexp are homogeneous in all types, and
|
||
|
|
# dispatching on the first argument type works of ldexp as well)
|
||
|
|
#
|
||
|
|
# assumes that the function pointed by func_name has the type
|
||
|
|
# signature sig (but needs translation to llvm types).
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
try:
|
||
|
|
func_name = table[ty]
|
||
|
|
except KeyError as e:
|
||
|
|
msg = "No {0} function for real type {1}".format(user_name, str(e))
|
||
|
|
raise errors.LoweringError(msg)
|
||
|
|
|
||
|
|
mod = builder.module
|
||
|
|
if ty in types.complex_domain:
|
||
|
|
# In numba struct types are always passed by pointer. So the call has to
|
||
|
|
# be transformed from "result = func(ops...)" to "func(&result, ops...).
|
||
|
|
# note that the result value pointer as first argument is the convention
|
||
|
|
# used by numba.
|
||
|
|
|
||
|
|
# First, prepare the return value
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
ptrargs = [cgutils.alloca_once_value(builder, arg)
|
||
|
|
for arg in args]
|
||
|
|
call_args = [out._getpointer()] + ptrargs
|
||
|
|
# get_value_as_argument for struct types like complex allocate stack space
|
||
|
|
# and initialize with the value, the return value is the pointer to that
|
||
|
|
# allocated space (ie: pointer to a copy of the value in the stack).
|
||
|
|
# get_argument_type returns a pointer to the struct type in consonance.
|
||
|
|
call_argtys = [ty] + list(sig.args)
|
||
|
|
call_argltys = [context.get_value_type(ty).as_pointer()
|
||
|
|
for ty in call_argtys]
|
||
|
|
fnty = llvmlite.ir.FunctionType(llvmlite.ir.VoidType(), call_argltys)
|
||
|
|
# Note: the function isn't pure here (it writes to its pointer args)
|
||
|
|
fn = cgutils.get_or_insert_function(mod, fnty, func_name)
|
||
|
|
builder.call(fn, call_args)
|
||
|
|
retval = builder.load(call_args[0])
|
||
|
|
else:
|
||
|
|
argtypes = [context.get_argument_type(aty) for aty in sig.args]
|
||
|
|
restype = context.get_argument_type(sig.return_type)
|
||
|
|
fnty = llvmlite.ir.FunctionType(restype, argtypes)
|
||
|
|
fn = cgutils.insert_pure_function(mod, fnty, name=func_name)
|
||
|
|
retval = context.call_external_function(builder, fn, sig.args, args)
|
||
|
|
return retval
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# Division kernels inspired by NumPy loops.c.src code
|
||
|
|
#
|
||
|
|
# The builtins are not applicable as they rely on a test for zero in the
|
||
|
|
# denominator. If it is zero the appropriate exception is raised.
|
||
|
|
# In NumPy, a division by zero does not raise an exception, but instead
|
||
|
|
# generated a known value. Note that a division by zero in any of the
|
||
|
|
# operations of a vector may raise an exception or issue a warning
|
||
|
|
# depending on the np.seterr configuration. This is not supported
|
||
|
|
# right now (and in any case, it won't be handled by these functions
|
||
|
|
# either)
|
||
|
|
|
||
|
|
def np_int_sdiv_impl(context, builder, sig, args):
|
||
|
|
# based on the actual code in NumPy loops.c.src for signed integer types
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
num, den = args
|
||
|
|
ty = sig.args[0] # any arg type will do, homogeneous
|
||
|
|
|
||
|
|
ZERO = context.get_constant(ty, 0)
|
||
|
|
MINUS_ONE = context.get_constant(ty, -1)
|
||
|
|
MIN_INT = context.get_constant(ty, 1 << (den.type.width-1))
|
||
|
|
den_is_zero = builder.icmp_unsigned('==', ZERO, den)
|
||
|
|
den_is_minus_one = builder.icmp_unsigned('==', MINUS_ONE, den)
|
||
|
|
num_is_min_int = builder.icmp_unsigned('==', MIN_INT, num)
|
||
|
|
could_cause_sigfpe = builder.and_(den_is_minus_one, num_is_min_int)
|
||
|
|
force_zero = builder.or_(den_is_zero, could_cause_sigfpe)
|
||
|
|
with builder.if_else(force_zero, likely=False) as (then, otherwise):
|
||
|
|
with then:
|
||
|
|
bb_then = builder.basic_block
|
||
|
|
with otherwise:
|
||
|
|
bb_otherwise = builder.basic_block
|
||
|
|
div = builder.sdiv(num, den)
|
||
|
|
mod = builder.srem(num, den)
|
||
|
|
num_gt_zero = builder.icmp_signed('>', num, ZERO)
|
||
|
|
den_gt_zero = builder.icmp_signed('>', den, ZERO)
|
||
|
|
not_same_sign = builder.xor(num_gt_zero, den_gt_zero)
|
||
|
|
mod_not_zero = builder.icmp_unsigned('!=', mod, ZERO)
|
||
|
|
needs_fixing = builder.and_(not_same_sign, mod_not_zero)
|
||
|
|
fix_value = builder.select(needs_fixing, MINUS_ONE, ZERO)
|
||
|
|
result_otherwise = builder.add(div, fix_value)
|
||
|
|
|
||
|
|
result = builder.phi(ZERO.type)
|
||
|
|
result.add_incoming(ZERO, bb_then)
|
||
|
|
result.add_incoming(result_otherwise, bb_otherwise)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_srem_impl(context, builder, sig, args):
|
||
|
|
# based on the actual code in NumPy loops.c.src for signed integers
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
num, den = args
|
||
|
|
ty = sig.args[0] # any arg type will do, homogeneous
|
||
|
|
|
||
|
|
ZERO = context.get_constant(ty, 0)
|
||
|
|
den_not_zero = builder.icmp_unsigned('!=', ZERO, den)
|
||
|
|
bb_no_if = builder.basic_block
|
||
|
|
with cgutils.if_unlikely(builder, den_not_zero):
|
||
|
|
bb_if = builder.basic_block
|
||
|
|
mod = builder.srem(num,den)
|
||
|
|
num_gt_zero = builder.icmp_signed('>', num, ZERO)
|
||
|
|
den_gt_zero = builder.icmp_signed('>', den, ZERO)
|
||
|
|
not_same_sign = builder.xor(num_gt_zero, den_gt_zero)
|
||
|
|
mod_not_zero = builder.icmp_unsigned('!=', mod, ZERO)
|
||
|
|
needs_fixing = builder.and_(not_same_sign, mod_not_zero)
|
||
|
|
fix_value = builder.select(needs_fixing, den, ZERO)
|
||
|
|
final_mod = builder.add(fix_value, mod)
|
||
|
|
|
||
|
|
result = builder.phi(ZERO.type)
|
||
|
|
result.add_incoming(ZERO, bb_no_if)
|
||
|
|
result.add_incoming(final_mod, bb_if)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_sdivrem_impl(context, builder, sig, args):
|
||
|
|
div = np_int_sdiv_impl(context, builder, sig.return_type[0](*sig.args), args)
|
||
|
|
rem = np_int_srem_impl(context, builder, sig.return_type[1](*sig.args), args)
|
||
|
|
return context.make_tuple(builder, sig.return_type, [div, rem])
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_udiv_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
num, den = args
|
||
|
|
ty = sig.args[0] # any arg type will do, homogeneous
|
||
|
|
|
||
|
|
ZERO = context.get_constant(ty, 0)
|
||
|
|
div_by_zero = builder.icmp_unsigned('==', ZERO, den)
|
||
|
|
with builder.if_else(div_by_zero, likely=False) as (then, otherwise):
|
||
|
|
with then:
|
||
|
|
# division by zero
|
||
|
|
bb_then = builder.basic_block
|
||
|
|
with otherwise:
|
||
|
|
# divide!
|
||
|
|
div = builder.udiv(num, den)
|
||
|
|
bb_otherwise = builder.basic_block
|
||
|
|
|
||
|
|
result = builder.phi(ZERO.type)
|
||
|
|
result.add_incoming(ZERO, bb_then)
|
||
|
|
result.add_incoming(div, bb_otherwise)
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_urem_impl(context, builder, sig, args):
|
||
|
|
# based on the actual code in NumPy loops.c.src for signed integers
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
num, den = args
|
||
|
|
ty = sig.args[0] # any arg type will do, homogeneous
|
||
|
|
|
||
|
|
ZERO = context.get_constant(ty, 0)
|
||
|
|
den_not_zero = builder.icmp_unsigned('!=', ZERO, den)
|
||
|
|
bb_no_if = builder.basic_block
|
||
|
|
with cgutils.if_unlikely(builder, den_not_zero):
|
||
|
|
bb_if = builder.basic_block
|
||
|
|
mod = builder.urem(num,den)
|
||
|
|
|
||
|
|
result = builder.phi(ZERO.type)
|
||
|
|
result.add_incoming(ZERO, bb_no_if)
|
||
|
|
result.add_incoming(mod, bb_if)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_udivrem_impl(context, builder, sig, args):
|
||
|
|
div = np_int_udiv_impl(context, builder, sig.return_type[0](*sig.args), args)
|
||
|
|
rem = np_int_urem_impl(context, builder, sig.return_type[1](*sig.args), args)
|
||
|
|
return context.make_tuple(builder, sig.return_type, [div, rem])
|
||
|
|
|
||
|
|
|
||
|
|
# implementation of int_fmod is in fact the same as the unsigned remainder,
|
||
|
|
# that is: srem with a special case returning 0 when the denominator is 0.
|
||
|
|
np_int_fmod_impl = np_int_urem_impl
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_div_impl(context, builder, sig, args):
|
||
|
|
# in NumPy real div has the same semantics as an fdiv for generating
|
||
|
|
# NANs, INF and NINF
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
return builder.fdiv(*args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_mod_impl(context, builder, sig, args):
|
||
|
|
# note: this maps to NumPy remainder, which has the same semantics as Python
|
||
|
|
# based on code in loops.c.src
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
in1, in2 = args
|
||
|
|
ty = sig.args[0]
|
||
|
|
|
||
|
|
ZERO = context.get_constant(ty, 0.0)
|
||
|
|
res = builder.frem(in1, in2)
|
||
|
|
res_ne_zero = builder.fcmp_ordered('!=', res, ZERO)
|
||
|
|
den_lt_zero = builder.fcmp_ordered('<', in2, ZERO)
|
||
|
|
res_lt_zero = builder.fcmp_ordered('<', res, ZERO)
|
||
|
|
needs_fixing = builder.and_(res_ne_zero,
|
||
|
|
builder.xor(den_lt_zero, res_lt_zero))
|
||
|
|
fix_value = builder.select(needs_fixing, in2, ZERO)
|
||
|
|
|
||
|
|
return builder.fadd(res, fix_value)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_fmod_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
return builder.frem(*args)
|
||
|
|
|
||
|
|
|
||
|
|
def _fabs(context, builder, arg):
|
||
|
|
ZERO = llvmlite.ir.Constant(arg.type, 0.0)
|
||
|
|
arg_negated = builder.fsub(ZERO, arg)
|
||
|
|
arg_is_negative = builder.fcmp_ordered('<', arg, ZERO)
|
||
|
|
return builder.select(arg_is_negative, arg_negated, arg)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_div_impl(context, builder, sig, args):
|
||
|
|
# Extracted from numpy/core/src/umath/loops.c.src,
|
||
|
|
# inspired by complex_div_impl
|
||
|
|
# variables named coherent with loops.c.src
|
||
|
|
# This is implemented using the approach described in
|
||
|
|
# R.L. Smith. Algorithm 116: Complex division.
|
||
|
|
# Communications of the ACM, 5(8):435, 1962
|
||
|
|
|
||
|
|
in1, in2 = [context.make_complex(builder, sig.args[0], value=arg)
|
||
|
|
for arg in args]
|
||
|
|
|
||
|
|
in1r = in1.real # numerator.real
|
||
|
|
in1i = in1.imag # numerator.imag
|
||
|
|
in2r = in2.real # denominator.real
|
||
|
|
in2i = in2.imag # denominator.imag
|
||
|
|
ftype = in1r.type
|
||
|
|
assert all([i.type==ftype for i in [in1r, in1i, in2r, in2i]]), "mismatched types"
|
||
|
|
out = context.make_helper(builder, sig.return_type)
|
||
|
|
|
||
|
|
ZERO = llvmlite.ir.Constant(ftype, 0.0)
|
||
|
|
ONE = llvmlite.ir.Constant(ftype, 1.0)
|
||
|
|
|
||
|
|
# if abs(denominator.real) >= abs(denominator.imag)
|
||
|
|
in2r_abs = _fabs(context, builder, in2r)
|
||
|
|
in2i_abs = _fabs(context, builder, in2i)
|
||
|
|
in2r_abs_ge_in2i_abs = builder.fcmp_ordered('>=', in2r_abs, in2i_abs)
|
||
|
|
with builder.if_else(in2r_abs_ge_in2i_abs) as (then, otherwise):
|
||
|
|
with then:
|
||
|
|
# if abs(denominator.real) == 0 and abs(denominator.imag) == 0
|
||
|
|
in2r_is_zero = builder.fcmp_ordered('==', in2r_abs, ZERO)
|
||
|
|
in2i_is_zero = builder.fcmp_ordered('==', in2i_abs, ZERO)
|
||
|
|
in2_is_zero = builder.and_(in2r_is_zero, in2i_is_zero)
|
||
|
|
with builder.if_else(in2_is_zero) as (inn_then, inn_otherwise):
|
||
|
|
with inn_then:
|
||
|
|
# division by 0.
|
||
|
|
# fdiv generates the appropriate NAN/INF/NINF
|
||
|
|
out.real = builder.fdiv(in1r, in2r_abs)
|
||
|
|
out.imag = builder.fdiv(in1i, in2i_abs)
|
||
|
|
with inn_otherwise:
|
||
|
|
# general case for:
|
||
|
|
# abs(denominator.real) > abs(denominator.imag)
|
||
|
|
rat = builder.fdiv(in2i, in2r)
|
||
|
|
# scl = 1.0/(in2r + in2i*rat)
|
||
|
|
tmp1 = builder.fmul(in2i, rat)
|
||
|
|
tmp2 = builder.fadd(in2r, tmp1)
|
||
|
|
scl = builder.fdiv(ONE, tmp2)
|
||
|
|
# out.real = (in1r + in1i*rat)*scl
|
||
|
|
# out.imag = (in1i - in1r*rat)*scl
|
||
|
|
tmp3 = builder.fmul(in1i, rat)
|
||
|
|
tmp4 = builder.fmul(in1r, rat)
|
||
|
|
tmp5 = builder.fadd(in1r, tmp3)
|
||
|
|
tmp6 = builder.fsub(in1i, tmp4)
|
||
|
|
out.real = builder.fmul(tmp5, scl)
|
||
|
|
out.imag = builder.fmul(tmp6, scl)
|
||
|
|
with otherwise:
|
||
|
|
# general case for:
|
||
|
|
# abs(denominator.imag) > abs(denominator.real)
|
||
|
|
rat = builder.fdiv(in2r, in2i)
|
||
|
|
# scl = 1.0/(in2i + in2r*rat)
|
||
|
|
tmp1 = builder.fmul(in2r, rat)
|
||
|
|
tmp2 = builder.fadd(in2i, tmp1)
|
||
|
|
scl = builder.fdiv(ONE, tmp2)
|
||
|
|
# out.real = (in1r*rat + in1i)*scl
|
||
|
|
# out.imag = (in1i*rat - in1r)*scl
|
||
|
|
tmp3 = builder.fmul(in1r, rat)
|
||
|
|
tmp4 = builder.fmul(in1i, rat)
|
||
|
|
tmp5 = builder.fadd(tmp3, in1i)
|
||
|
|
tmp6 = builder.fsub(tmp4, in1r)
|
||
|
|
out.real = builder.fmul(tmp5, scl)
|
||
|
|
out.imag = builder.fmul(tmp6, scl)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy logaddexp
|
||
|
|
|
||
|
|
def _npy_logaddexp(x1, x2):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def _generate_logaddexp(fnoverload, const, log1pfn, expfn):
|
||
|
|
# Code generation for logaddexp and logaddexp2 is based on:
|
||
|
|
# https://github.com/numpy/numpy/blob/12c2b7dd62fc0c14b81c8892ed5f4f59cc94d09c/numpy/core/src/npymath/npy_math_internal.h.src#L467-L507
|
||
|
|
|
||
|
|
@overload(fnoverload, target='generic')
|
||
|
|
def ol_npy_logaddexp(x1, x2):
|
||
|
|
if x1 != x2:
|
||
|
|
return
|
||
|
|
shift = x1(const)
|
||
|
|
def impl(x1, x2):
|
||
|
|
x, y = x1, x2
|
||
|
|
if (x == y):
|
||
|
|
# Handles infinities of the same sign without warnings
|
||
|
|
return x + shift
|
||
|
|
else:
|
||
|
|
tmp = x - y
|
||
|
|
if (tmp > 0):
|
||
|
|
return x + log1pfn(expfn(-tmp))
|
||
|
|
elif (tmp <= 0):
|
||
|
|
return y + log1pfn(expfn(tmp))
|
||
|
|
else:
|
||
|
|
# NaN
|
||
|
|
return tmp
|
||
|
|
return impl
|
||
|
|
|
||
|
|
def _npy_logaddexp(x1, x2):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
_generate_logaddexp(_npy_logaddexp, _NPY_LOGE2, np.log1p, np.exp)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_logaddexp_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
fnty = context.typing_context.resolve_value_type(_npy_logaddexp)
|
||
|
|
sig = fnty.get_call_type(context.typing_context, (*sig.args,), {})
|
||
|
|
impl = context.get_function(fnty, sig)
|
||
|
|
return impl(builder, args)
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy logaddexp2
|
||
|
|
def _npy_logaddexp2(x1, x2):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def npy_log2_1p(x):
|
||
|
|
pass
|
||
|
|
|
||
|
|
# The following npy_log2_1p function is a translation of:
|
||
|
|
# https://github.com/numpy/numpy/blob/12c2b7dd62fc0c14b81c8892ed5f4f59cc94d09c/numpy/core/src/npymath/npy_math_internal.h.src#L457-L460
|
||
|
|
|
||
|
|
@overload(npy_log2_1p, target='generic')
|
||
|
|
def ol_npy_log2_1p(x):
|
||
|
|
LOG2E = x(_NPY_LOG2E)
|
||
|
|
def impl(x):
|
||
|
|
return LOG2E * np.log1p(x)
|
||
|
|
return impl
|
||
|
|
|
||
|
|
|
||
|
|
_generate_logaddexp(_npy_logaddexp2, 1.0, npy_log2_1p, np.exp2)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_logaddexp2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
fnty = context.typing_context.resolve_value_type(_npy_logaddexp2)
|
||
|
|
sig = fnty.get_call_type(context.typing_context, (*sig.args,), {})
|
||
|
|
impl = context.get_function(fnty, sig)
|
||
|
|
return impl(builder, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# true div kernels
|
||
|
|
|
||
|
|
def np_int_truediv_impl(context, builder, sig, args):
|
||
|
|
# in NumPy we don't check for 0 denominator... fdiv handles div by
|
||
|
|
# 0 in the way NumPy expects..
|
||
|
|
# integer truediv always yields double
|
||
|
|
num, den = args
|
||
|
|
lltype = num.type
|
||
|
|
assert all(i.type==lltype for i in args), "must have homogeneous types"
|
||
|
|
numty, denty = sig.args
|
||
|
|
|
||
|
|
num = context.cast(builder, num, numty, types.float64)
|
||
|
|
den = context.cast(builder, den, denty, types.float64)
|
||
|
|
|
||
|
|
return builder.fdiv(num,den)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# floor div kernels
|
||
|
|
|
||
|
|
def np_real_floor_div_impl(context, builder, sig, args):
|
||
|
|
res = np_real_div_impl(context, builder, sig, args)
|
||
|
|
s = typing.signature(sig.return_type, sig.return_type)
|
||
|
|
return np_real_floor_impl(context, builder, s, (res,))
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_divmod_impl(context, builder, sig, args):
|
||
|
|
div = np_real_floor_div_impl(context, builder, sig.return_type[0](*sig.args), args)
|
||
|
|
rem = np_real_mod_impl(context, builder, sig.return_type[1](*sig.args), args)
|
||
|
|
return context.make_tuple(builder, sig.return_type, [div, rem])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_floor_div_impl(context, builder, sig, args):
|
||
|
|
# this is based on the complex floor divide in Numpy's loops.c.src
|
||
|
|
# This is basically a full complex division with a complex floor
|
||
|
|
# applied.
|
||
|
|
# The complex floor seems to be defined as the real floor applied
|
||
|
|
# with the real part and zero in the imaginary part. Fully developed
|
||
|
|
# so it avoids computing anything related to the imaginary result.
|
||
|
|
float_kind = sig.args[0].underlying_float
|
||
|
|
floor_sig = typing.signature(float_kind, float_kind)
|
||
|
|
|
||
|
|
in1, in2 = [context.make_complex(builder, sig.args[0], value=arg)
|
||
|
|
for arg in args]
|
||
|
|
|
||
|
|
in1r = in1.real
|
||
|
|
in1i = in1.imag
|
||
|
|
in2r = in2.real
|
||
|
|
in2i = in2.imag
|
||
|
|
ftype = in1r.type
|
||
|
|
assert all([i.type==ftype for i in [in1r, in1i, in2r, in2i]]), "mismatched types"
|
||
|
|
|
||
|
|
ZERO = llvmlite.ir.Constant(ftype, 0.0)
|
||
|
|
|
||
|
|
out = context.make_helper(builder, sig.return_type)
|
||
|
|
out.imag = ZERO
|
||
|
|
|
||
|
|
in2r_abs = _fabs(context, builder, in2r)
|
||
|
|
in2i_abs = _fabs(context, builder, in2i)
|
||
|
|
in2r_abs_ge_in2i_abs = builder.fcmp_ordered('>=', in2r_abs, in2i_abs)
|
||
|
|
|
||
|
|
with builder.if_else(in2r_abs_ge_in2i_abs) as (then, otherwise):
|
||
|
|
with then:
|
||
|
|
rat = builder.fdiv(in2i, in2r)
|
||
|
|
# out.real = floor((in1r+in1i*rat)/(in2r + in2i*rat))
|
||
|
|
tmp1 = builder.fmul(in1i, rat)
|
||
|
|
tmp2 = builder.fmul(in2i, rat)
|
||
|
|
tmp3 = builder.fadd(in1r, tmp1)
|
||
|
|
tmp4 = builder.fadd(in2r, tmp2)
|
||
|
|
tmp5 = builder.fdiv(tmp3, tmp4)
|
||
|
|
out.real = np_real_floor_impl(context, builder, floor_sig, (tmp5,))
|
||
|
|
with otherwise:
|
||
|
|
rat = builder.fdiv(in2r, in2i)
|
||
|
|
# out.real = floor((in1i + in1r*rat)/(in2i + in2r*rat))
|
||
|
|
tmp1 = builder.fmul(in1r, rat)
|
||
|
|
tmp2 = builder.fmul(in2r, rat)
|
||
|
|
tmp3 = builder.fadd(in1i, tmp1)
|
||
|
|
tmp4 = builder.fadd(in2i, tmp2)
|
||
|
|
tmp5 = builder.fdiv(tmp3, tmp4)
|
||
|
|
out.real = np_real_floor_impl(context, builder, floor_sig, (tmp5,))
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# numpy power funcs
|
||
|
|
|
||
|
|
def np_complex_power_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
return numbers.complex_power_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# numpy float power funcs
|
||
|
|
|
||
|
|
def real_float_power_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
return numbers.real_power_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_float_power_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
return numbers.complex_power_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# numpy greatest common denominator
|
||
|
|
|
||
|
|
def np_gcd_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
return mathimpl.gcd_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# numpy lowest common multiple
|
||
|
|
|
||
|
|
def np_lcm_impl(context, builder, sig, args):
|
||
|
|
|
||
|
|
xty, yty = sig.args
|
||
|
|
assert xty == yty == sig.return_type
|
||
|
|
x, y = args
|
||
|
|
|
||
|
|
def lcm(a, b):
|
||
|
|
"""
|
||
|
|
Like gcd, heavily cribbed from Julia.
|
||
|
|
"""
|
||
|
|
return 0 if a == 0 else abs(a * (b // np.gcd(b, a)))
|
||
|
|
|
||
|
|
res = context.compile_internal(builder, lcm, sig, args)
|
||
|
|
return impl_ret_untracked(context, builder, sig.return_type, res)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# Numpy style complex sign
|
||
|
|
|
||
|
|
def np_complex_sign_impl(context, builder, sig, args):
|
||
|
|
# equivalent to complex sign in NumPy's sign
|
||
|
|
# but implemented via selects, balancing the 4 cases.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
op = args[0]
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
|
||
|
|
ZERO = context.get_constant(float_ty, 0.0)
|
||
|
|
ONE = context.get_constant(float_ty, 1.0)
|
||
|
|
MINUS_ONE = context.get_constant(float_ty, -1.0)
|
||
|
|
NAN = context.get_constant(float_ty, float('nan'))
|
||
|
|
result = context.make_complex(builder, ty)
|
||
|
|
result.real = ZERO
|
||
|
|
result.imag = ZERO
|
||
|
|
|
||
|
|
cmp_sig = typing.signature(types.boolean, *[ty] * 2)
|
||
|
|
cmp_args = [op, result._getvalue()]
|
||
|
|
arg1_ge_arg2 = np_complex_ge_impl(context, builder, cmp_sig, cmp_args)
|
||
|
|
arg1_eq_arg2 = np_complex_eq_impl(context, builder, cmp_sig, cmp_args)
|
||
|
|
arg1_lt_arg2 = np_complex_lt_impl(context, builder, cmp_sig, cmp_args)
|
||
|
|
|
||
|
|
real_when_ge = builder.select(arg1_eq_arg2, ZERO, ONE)
|
||
|
|
real_when_nge = builder.select(arg1_lt_arg2, MINUS_ONE, NAN)
|
||
|
|
result.real = builder.select(arg1_ge_arg2, real_when_ge, real_when_nge)
|
||
|
|
|
||
|
|
return result._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# Numpy rint
|
||
|
|
|
||
|
|
def np_real_rint_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
return mathimpl.call_fp_intrinsic(builder, 'llvm.rint', args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_rint_impl(context, builder, sig, args):
|
||
|
|
# based on code in NumPy's funcs.inc.src
|
||
|
|
# rint of a complex number defined as rint of its real and imag
|
||
|
|
# parts
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
in1 = context.make_complex(builder, ty, value=args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
|
||
|
|
inner_sig = typing.signature(*[float_ty]*2)
|
||
|
|
out.real = np_real_rint_impl(context, builder, inner_sig, [in1.real])
|
||
|
|
out.imag = np_real_rint_impl(context, builder, inner_sig, [in1.imag])
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy exp
|
||
|
|
|
||
|
|
def np_real_exp_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.exp_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_exp_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return cmathimpl.exp_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy exp2
|
||
|
|
|
||
|
|
def np_real_exp2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ll_ty = args[0].type
|
||
|
|
fnty = llvmlite.ir.FunctionType(ll_ty, [ll_ty,])
|
||
|
|
fn = cgutils.insert_pure_function(builder.module, fnty,
|
||
|
|
name='llvm.exp2')
|
||
|
|
return builder.call(fn, [args[0]])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_exp2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
in1 = context.make_complex(builder, ty, value=args[0])
|
||
|
|
tmp = context.make_complex(builder, ty)
|
||
|
|
loge2 = context.get_constant(float_ty, _NPY_LOGE2)
|
||
|
|
tmp.real = builder.fmul(loge2, in1.real)
|
||
|
|
tmp.imag = builder.fmul(loge2, in1.imag)
|
||
|
|
return np_complex_exp_impl(context, builder, sig, [tmp._getvalue()])
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy log
|
||
|
|
|
||
|
|
def np_real_log_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.log_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_log_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return cmathimpl.log_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy log2
|
||
|
|
|
||
|
|
def np_real_log2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ll_ty = args[0].type
|
||
|
|
fnty = llvmlite.ir.FunctionType(ll_ty, [ll_ty,])
|
||
|
|
fn = cgutils.insert_pure_function(builder.module, fnty,
|
||
|
|
name='llvm.log2')
|
||
|
|
return builder.call(fn, [args[0]])
|
||
|
|
|
||
|
|
def np_complex_log2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
tmp = np_complex_log_impl(context, builder, sig, args)
|
||
|
|
tmp = context.make_complex(builder, ty, value=tmp)
|
||
|
|
log2e = context.get_constant(float_ty, _NPY_LOG2E)
|
||
|
|
tmp.real = builder.fmul(log2e, tmp.real)
|
||
|
|
tmp.imag = builder.fmul(log2e, tmp.imag)
|
||
|
|
return tmp._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy log10
|
||
|
|
|
||
|
|
def np_real_log10_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.log10_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_log10_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
tmp = np_complex_log_impl(context, builder, sig, args)
|
||
|
|
tmp = context.make_complex(builder, ty, value=tmp)
|
||
|
|
log10e = context.get_constant(float_ty, _NPY_LOG10E)
|
||
|
|
tmp.real = builder.fmul(log10e, tmp.real)
|
||
|
|
tmp.imag = builder.fmul(log10e, tmp.imag)
|
||
|
|
return tmp._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy expm1
|
||
|
|
|
||
|
|
def np_real_expm1_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.expm1_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
def np_complex_expm1_impl(context, builder, sig, args):
|
||
|
|
# this is based on nc_expm1 in funcs.inc.src
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
float_unary_sig = typing.signature(*[float_ty]*2)
|
||
|
|
|
||
|
|
MINUS_ONE = context.get_constant(float_ty, -1.0)
|
||
|
|
in1 = context.make_complex(builder, ty, value=args[0])
|
||
|
|
a = np_real_exp_impl(context, builder, float_unary_sig, [in1.real])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
cos_imag = np_real_cos_impl(context, builder, float_unary_sig, [in1.imag])
|
||
|
|
sin_imag = np_real_sin_impl(context, builder, float_unary_sig, [in1.imag])
|
||
|
|
tmp = builder.fmul(a, cos_imag)
|
||
|
|
out.imag = builder.fmul(a, sin_imag)
|
||
|
|
out.real = builder.fadd(tmp, MINUS_ONE)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy log1p
|
||
|
|
|
||
|
|
def np_real_log1p_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.log1p_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
def np_complex_log1p_impl(context, builder, sig, args):
|
||
|
|
# base on NumPy's nc_log1p in funcs.inc.src
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
float_unary_sig = typing.signature(*[float_ty]*2)
|
||
|
|
float_binary_sig = typing.signature(*[float_ty]*3)
|
||
|
|
|
||
|
|
ONE = context.get_constant(float_ty, 1.0)
|
||
|
|
in1 = context.make_complex(builder, ty, value=args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
real_plus_one = builder.fadd(in1.real, ONE)
|
||
|
|
l = np_real_hypot_impl(context, builder, float_binary_sig,
|
||
|
|
[real_plus_one, in1.imag])
|
||
|
|
out.imag = np_real_atan2_impl(context, builder, float_binary_sig,
|
||
|
|
[in1.imag, real_plus_one])
|
||
|
|
out.real = np_real_log_impl(context, builder, float_unary_sig, [l])
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy sqrt
|
||
|
|
|
||
|
|
def np_real_sqrt_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.sqrt_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_sqrt_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return cmathimpl.sqrt_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy square
|
||
|
|
|
||
|
|
def np_int_square_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return builder.mul(args[0], args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_square_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return builder.fmul(args[0], args[0])
|
||
|
|
|
||
|
|
def np_complex_square_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
binary_sig = typing.signature(*[sig.return_type]*3)
|
||
|
|
return numbers.complex_mul_impl(context, builder, binary_sig,
|
||
|
|
[args[0], args[0]])
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy cbrt
|
||
|
|
|
||
|
|
def np_real_cbrt_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
# We enable fastmath here to force np.power(x, 1/3) to generate a
|
||
|
|
# call to libm cbrt function
|
||
|
|
@register_jitable(fastmath=True)
|
||
|
|
def cbrt(x):
|
||
|
|
if x < 0:
|
||
|
|
return -np.power(-x, 1.0 / 3.0)
|
||
|
|
else:
|
||
|
|
return np.power(x, 1.0 / 3.0)
|
||
|
|
|
||
|
|
def _cbrt(x):
|
||
|
|
if np.isnan(x):
|
||
|
|
return np.nan
|
||
|
|
return cbrt(x)
|
||
|
|
|
||
|
|
return context.compile_internal(builder, _cbrt, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy reciprocal
|
||
|
|
|
||
|
|
def np_int_reciprocal_impl(context, builder, sig, args):
|
||
|
|
# based on the implementation in loops.c.src
|
||
|
|
# integer versions for reciprocal are performed via promotion
|
||
|
|
# using double, and then converted back to the type
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
ty = sig.return_type
|
||
|
|
|
||
|
|
binary_sig = typing.signature(*[ty]*3)
|
||
|
|
in_as_float = context.cast(builder, args[0], ty, types.float64)
|
||
|
|
ONE = context.get_constant(types.float64, 1)
|
||
|
|
result_as_float = builder.fdiv(ONE, in_as_float)
|
||
|
|
return context.cast(builder, result_as_float, types.float64, ty)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_reciprocal_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
ONE = context.get_constant(sig.return_type, 1.0)
|
||
|
|
return builder.fdiv(ONE, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_reciprocal_impl(context, builder, sig, args):
|
||
|
|
# based on the implementation in loops.c.src
|
||
|
|
# Basically the same Smith method used for division, but with
|
||
|
|
# the numerator substituted by 1.0
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
float_ty = ty.underlying_float
|
||
|
|
|
||
|
|
ZERO = context.get_constant(float_ty, 0.0)
|
||
|
|
ONE = context.get_constant(float_ty, 1.0)
|
||
|
|
in1 = context.make_complex(builder, ty, value=args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
in1r = in1.real
|
||
|
|
in1i = in1.imag
|
||
|
|
in1r_abs = _fabs(context, builder, in1r)
|
||
|
|
in1i_abs = _fabs(context, builder, in1i)
|
||
|
|
in1i_abs_le_in1r_abs = builder.fcmp_ordered('<=', in1i_abs, in1r_abs)
|
||
|
|
|
||
|
|
with builder.if_else(in1i_abs_le_in1r_abs) as (then, otherwise):
|
||
|
|
with then:
|
||
|
|
r = builder.fdiv(in1i, in1r)
|
||
|
|
tmp0 = builder.fmul(in1i, r)
|
||
|
|
d = builder.fadd(in1r, tmp0)
|
||
|
|
inv_d = builder.fdiv(ONE, d)
|
||
|
|
minus_r = builder.fsub(ZERO, r)
|
||
|
|
out.real = inv_d
|
||
|
|
out.imag = builder.fmul(minus_r, inv_d)
|
||
|
|
with otherwise:
|
||
|
|
r = builder.fdiv(in1r, in1i)
|
||
|
|
tmp0 = builder.fmul(in1r, r)
|
||
|
|
d = builder.fadd(tmp0, in1i)
|
||
|
|
inv_d = builder.fdiv(ONE, d)
|
||
|
|
out.real = builder.fmul(r, inv_d)
|
||
|
|
out.imag = builder.fsub(ZERO, inv_d)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy sin
|
||
|
|
|
||
|
|
def np_real_sin_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.sin_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_sin_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return cmathimpl.sin_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy cos
|
||
|
|
|
||
|
|
def np_real_cos_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.cos_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_cos_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return cmathimpl.cos_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy tan
|
||
|
|
|
||
|
|
def np_real_tan_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.tan_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy asin
|
||
|
|
|
||
|
|
def np_real_asin_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.asin_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy acos
|
||
|
|
|
||
|
|
def np_real_acos_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.acos_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy atan
|
||
|
|
|
||
|
|
def np_real_atan_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.atan_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy atan2
|
||
|
|
|
||
|
|
def np_real_atan2_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
return mathimpl.atan2_float_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy hypot
|
||
|
|
|
||
|
|
def np_real_hypot_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
return mathimpl.hypot_float_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy sinh
|
||
|
|
|
||
|
|
def np_real_sinh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.sinh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_sinh_impl(context, builder, sig, args):
|
||
|
|
# npymath does not provide a complex sinh. The code in funcs.inc.src
|
||
|
|
# is translated here...
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
fty = ty.underlying_float
|
||
|
|
fsig1 = typing.signature(*[fty]*2)
|
||
|
|
x = context.make_complex(builder, ty, args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
xr = x.real
|
||
|
|
xi = x.imag
|
||
|
|
|
||
|
|
sxi = np_real_sin_impl(context, builder, fsig1, [xi])
|
||
|
|
shxr = np_real_sinh_impl(context, builder, fsig1, [xr])
|
||
|
|
cxi = np_real_cos_impl(context, builder, fsig1, [xi])
|
||
|
|
chxr = np_real_cosh_impl(context, builder, fsig1, [xr])
|
||
|
|
|
||
|
|
out.real = builder.fmul(cxi, shxr)
|
||
|
|
out.imag = builder.fmul(sxi, chxr)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy cosh
|
||
|
|
|
||
|
|
def np_real_cosh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.cosh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_cosh_impl(context, builder, sig, args):
|
||
|
|
# npymath does not provide a complex cosh. The code in funcs.inc.src
|
||
|
|
# is translated here...
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
fty = ty.underlying_float
|
||
|
|
fsig1 = typing.signature(*[fty]*2)
|
||
|
|
x = context.make_complex(builder, ty, args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
xr = x.real
|
||
|
|
xi = x.imag
|
||
|
|
|
||
|
|
cxi = np_real_cos_impl(context, builder, fsig1, [xi])
|
||
|
|
chxr = np_real_cosh_impl(context, builder, fsig1, [xr])
|
||
|
|
sxi = np_real_sin_impl(context, builder, fsig1, [xi])
|
||
|
|
shxr = np_real_sinh_impl(context, builder, fsig1, [xr])
|
||
|
|
|
||
|
|
out.real = builder.fmul(cxi, chxr)
|
||
|
|
out.imag = builder.fmul(sxi, shxr)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy tanh
|
||
|
|
|
||
|
|
def np_real_tanh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.tanh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_tanh_impl(context, builder, sig, args):
|
||
|
|
# npymath does not provide complex tan functions. The code
|
||
|
|
# in funcs.inc.src for tanh is translated here...
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
fty = ty.underlying_float
|
||
|
|
fsig1 = typing.signature(*[fty]*2)
|
||
|
|
ONE = context.get_constant(fty, 1.0)
|
||
|
|
x = context.make_complex(builder, ty, args[0])
|
||
|
|
out = context.make_complex(builder, ty)
|
||
|
|
|
||
|
|
xr = x.real
|
||
|
|
xi = x.imag
|
||
|
|
si = np_real_sin_impl(context, builder, fsig1, [xi])
|
||
|
|
ci = np_real_cos_impl(context, builder, fsig1, [xi])
|
||
|
|
shr = np_real_sinh_impl(context, builder, fsig1, [xr])
|
||
|
|
chr_ = np_real_cosh_impl(context, builder, fsig1, [xr])
|
||
|
|
rs = builder.fmul(ci, shr)
|
||
|
|
is_ = builder.fmul(si, chr_)
|
||
|
|
rc = builder.fmul(ci, chr_)
|
||
|
|
ic = builder.fmul(si, shr) # note: opposite sign from code in funcs.inc.src
|
||
|
|
sqr_rc = builder.fmul(rc, rc)
|
||
|
|
sqr_ic = builder.fmul(ic, ic)
|
||
|
|
d = builder.fadd(sqr_rc, sqr_ic)
|
||
|
|
inv_d = builder.fdiv(ONE, d)
|
||
|
|
rs_rc = builder.fmul(rs, rc)
|
||
|
|
is_ic = builder.fmul(is_, ic)
|
||
|
|
is_rc = builder.fmul(is_, rc)
|
||
|
|
rs_ic = builder.fmul(rs, ic)
|
||
|
|
numr = builder.fadd(rs_rc, is_ic)
|
||
|
|
numi = builder.fsub(is_rc, rs_ic)
|
||
|
|
out.real = builder.fmul(numr, inv_d)
|
||
|
|
out.imag = builder.fmul(numi, inv_d)
|
||
|
|
|
||
|
|
return out._getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy asinh
|
||
|
|
|
||
|
|
def np_real_asinh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.asinh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy acosh
|
||
|
|
|
||
|
|
def np_real_acosh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.acosh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_acosh_impl(context, builder, sig, args):
|
||
|
|
# npymath does not provide a complex acosh. The code in funcs.inc.src
|
||
|
|
# is translated here...
|
||
|
|
# log(x + sqrt(x+1) * sqrt(x-1))
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
csig2 = typing.signature(*[ty]*3)
|
||
|
|
|
||
|
|
ONE = context.get_constant_generic(builder, ty, 1.0 + 0.0j)
|
||
|
|
x = args[0]
|
||
|
|
|
||
|
|
x_plus_one = numbers.complex_add_impl(context, builder, csig2, [x,
|
||
|
|
ONE])
|
||
|
|
x_minus_one = numbers.complex_sub_impl(context, builder, csig2, [x,
|
||
|
|
ONE])
|
||
|
|
sqrt_x_plus_one = np_complex_sqrt_impl(context, builder, sig, [x_plus_one])
|
||
|
|
sqrt_x_minus_one = np_complex_sqrt_impl(context, builder, sig, [x_minus_one])
|
||
|
|
prod_sqrt = numbers.complex_mul_impl(context, builder, csig2,
|
||
|
|
[sqrt_x_plus_one,
|
||
|
|
sqrt_x_minus_one])
|
||
|
|
log_arg = numbers.complex_add_impl(context, builder, csig2, [x,
|
||
|
|
prod_sqrt])
|
||
|
|
|
||
|
|
return np_complex_log_impl(context, builder, sig, [log_arg])
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy atanh
|
||
|
|
|
||
|
|
def np_real_atanh_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
return mathimpl.atanh_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy floor
|
||
|
|
|
||
|
|
def np_real_floor_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
return mathimpl.call_fp_intrinsic(builder, 'llvm.floor', args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy ceil
|
||
|
|
|
||
|
|
def np_real_ceil_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
return mathimpl.call_fp_intrinsic(builder, 'llvm.ceil', args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy trunc
|
||
|
|
|
||
|
|
def np_real_trunc_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
return mathimpl.call_fp_intrinsic(builder, 'llvm.trunc', args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy fabs
|
||
|
|
|
||
|
|
def np_real_fabs_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
return mathimpl.call_fp_intrinsic(builder, 'llvm.fabs', args)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy style predicates
|
||
|
|
|
||
|
|
# For real and integer types rely on numbers... but complex ordering in
|
||
|
|
# NumPy is lexicographic (while Python does not provide ordering).
|
||
|
|
def np_complex_ge_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CGE in NumPy's loops.c.src
|
||
|
|
# ((xr > yr && !npy_isnan(xi) && !npy_isnan(yi)) || (xr == yr && xi >= yi))
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_gt_yr = builder.fcmp_ordered('>', xr, yr)
|
||
|
|
no_nan_xi_yi = builder.fcmp_ordered('ord', xi, yi)
|
||
|
|
xr_eq_yr = builder.fcmp_ordered('==', xr, yr)
|
||
|
|
xi_ge_yi = builder.fcmp_ordered('>=', xi, yi)
|
||
|
|
first_term = builder.and_(xr_gt_yr, no_nan_xi_yi)
|
||
|
|
second_term = builder.and_(xr_eq_yr, xi_ge_yi)
|
||
|
|
return builder.or_(first_term, second_term)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_le_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CLE in NumPy's loops.c.src
|
||
|
|
# ((xr < yr && !npy_isnan(xi) && !npy_isnan(yi)) || (xr == yr && xi <= yi))
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_lt_yr = builder.fcmp_ordered('<', xr, yr)
|
||
|
|
no_nan_xi_yi = builder.fcmp_ordered('ord', xi, yi)
|
||
|
|
xr_eq_yr = builder.fcmp_ordered('==', xr, yr)
|
||
|
|
xi_le_yi = builder.fcmp_ordered('<=', xi, yi)
|
||
|
|
first_term = builder.and_(xr_lt_yr, no_nan_xi_yi)
|
||
|
|
second_term = builder.and_(xr_eq_yr, xi_le_yi)
|
||
|
|
return builder.or_(first_term, second_term)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_gt_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CGT in NumPy's loops.c.src
|
||
|
|
# ((xr > yr && !npy_isnan(xi) && !npy_isnan(yi)) || (xr == yr && xi > yi))
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_gt_yr = builder.fcmp_ordered('>', xr, yr)
|
||
|
|
no_nan_xi_yi = builder.fcmp_ordered('ord', xi, yi)
|
||
|
|
xr_eq_yr = builder.fcmp_ordered('==', xr, yr)
|
||
|
|
xi_gt_yi = builder.fcmp_ordered('>', xi, yi)
|
||
|
|
first_term = builder.and_(xr_gt_yr, no_nan_xi_yi)
|
||
|
|
second_term = builder.and_(xr_eq_yr, xi_gt_yi)
|
||
|
|
return builder.or_(first_term, second_term)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_lt_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CLT in NumPy's loops.c.src
|
||
|
|
# ((xr < yr && !npy_isnan(xi) && !npy_isnan(yi)) || (xr == yr && xi < yi))
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_lt_yr = builder.fcmp_ordered('<', xr, yr)
|
||
|
|
no_nan_xi_yi = builder.fcmp_ordered('ord', xi, yi)
|
||
|
|
xr_eq_yr = builder.fcmp_ordered('==', xr, yr)
|
||
|
|
xi_lt_yi = builder.fcmp_ordered('<', xi, yi)
|
||
|
|
first_term = builder.and_(xr_lt_yr, no_nan_xi_yi)
|
||
|
|
second_term = builder.and_(xr_eq_yr, xi_lt_yi)
|
||
|
|
return builder.or_(first_term, second_term)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_eq_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CEQ in NumPy's loops.c.src
|
||
|
|
# (xr == yr && xi == yi)
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_eq_yr = builder.fcmp_ordered('==', xr, yr)
|
||
|
|
xi_eq_yi = builder.fcmp_ordered('==', xi, yi)
|
||
|
|
return builder.and_(xr_eq_yr, xi_eq_yi)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_ne_impl(context, builder, sig, args):
|
||
|
|
# equivalent to macro CNE in NumPy's loops.c.src
|
||
|
|
# (xr != yr || xi != yi)
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
|
||
|
|
ty = sig.args[0]
|
||
|
|
in1, in2 = [context.make_complex(builder, ty, value=arg) for arg in args]
|
||
|
|
xr = in1.real
|
||
|
|
xi = in1.imag
|
||
|
|
yr = in2.real
|
||
|
|
yi = in2.imag
|
||
|
|
|
||
|
|
xr_ne_yr = builder.fcmp_unordered('!=', xr, yr)
|
||
|
|
xi_ne_yi = builder.fcmp_unordered('!=', xi, yi)
|
||
|
|
return builder.or_(xr_ne_yr, xi_ne_yi)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy logical algebra
|
||
|
|
|
||
|
|
# these are made generic for all types for now, assuming that
|
||
|
|
# cgutils.is_true works in the underlying types.
|
||
|
|
|
||
|
|
def _complex_is_true(context, builder, ty, val):
|
||
|
|
complex_val = context.make_complex(builder, ty, value=val)
|
||
|
|
re_true = cgutils.is_true(builder, complex_val.real)
|
||
|
|
im_true = cgutils.is_true(builder, complex_val.imag)
|
||
|
|
return builder.or_(re_true, im_true)
|
||
|
|
|
||
|
|
|
||
|
|
def np_logical_and_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = cgutils.is_true(builder, args[0])
|
||
|
|
b = cgutils.is_true(builder, args[1])
|
||
|
|
return builder.and_(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_logical_and_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = _complex_is_true(context, builder, sig.args[0], args[0])
|
||
|
|
b = _complex_is_true(context, builder, sig.args[1], args[1])
|
||
|
|
return builder.and_(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_logical_or_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = cgutils.is_true(builder, args[0])
|
||
|
|
b = cgutils.is_true(builder, args[1])
|
||
|
|
return builder.or_(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_logical_or_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = _complex_is_true(context, builder, sig.args[0], args[0])
|
||
|
|
b = _complex_is_true(context, builder, sig.args[1], args[1])
|
||
|
|
return builder.or_(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_logical_xor_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = cgutils.is_true(builder, args[0])
|
||
|
|
b = cgutils.is_true(builder, args[1])
|
||
|
|
return builder.xor(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_logical_xor_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2, return_type=types.boolean)
|
||
|
|
a = _complex_is_true(context, builder, sig.args[0], args[0])
|
||
|
|
b = _complex_is_true(context, builder, sig.args[1], args[1])
|
||
|
|
return builder.xor(a, b)
|
||
|
|
|
||
|
|
|
||
|
|
def np_logical_not_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return cgutils.is_false(builder, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_logical_not_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
a = _complex_is_true(context, builder, sig.args[0], args[0])
|
||
|
|
return builder.not_(a)
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy style max/min
|
||
|
|
#
|
||
|
|
# There are 2 different sets of functions to perform max and min in
|
||
|
|
# NumPy: maximum/minimum and fmax/fmin.
|
||
|
|
# Both differ in the way NaNs are handled, so the actual differences
|
||
|
|
# come in action only on float/complex numbers. The functions used for
|
||
|
|
# integers is shared. For booleans maximum is equivalent to or, and
|
||
|
|
# minimum is equivalent to and. Datetime support will go elsewhere.
|
||
|
|
|
||
|
|
def np_int_smax_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_sge_arg2 = builder.icmp_signed('>=', arg1, arg2)
|
||
|
|
return builder.select(arg1_sge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_umax_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_uge_arg2 = builder.icmp_unsigned('>=', arg1, arg2)
|
||
|
|
return builder.select(arg1_uge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_maximum_impl(context, builder, sig, args):
|
||
|
|
# maximum prefers nan (tries to return a nan).
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = builder.fcmp_unordered('uno', arg1, arg1)
|
||
|
|
any_nan = builder.fcmp_unordered('uno', arg1, arg2)
|
||
|
|
nan_result = builder.select(arg1_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_ge_arg2 = builder.fcmp_ordered('>=', arg1, arg2)
|
||
|
|
non_nan_result = builder.select(arg1_ge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_fmax_impl(context, builder, sig, args):
|
||
|
|
# fmax prefers non-nan (tries to return a non-nan).
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg2_nan = builder.fcmp_unordered('uno', arg2, arg2)
|
||
|
|
any_nan = builder.fcmp_unordered('uno', arg1, arg2)
|
||
|
|
nan_result = builder.select(arg2_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_ge_arg2 = builder.fcmp_ordered('>=', arg1, arg2)
|
||
|
|
non_nan_result = builder.select(arg1_ge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_maximum_impl(context, builder, sig, args):
|
||
|
|
# maximum prefers nan (tries to return a nan).
|
||
|
|
# There is an extra caveat with complex numbers, as there is more
|
||
|
|
# than one type of nan. NumPy's docs state that the nan in the
|
||
|
|
# first argument is returned when both arguments are nans.
|
||
|
|
# If only one nan is found, that nan is returned.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
ty = sig.args[0]
|
||
|
|
bc_sig = typing.signature(types.boolean, ty)
|
||
|
|
bcc_sig = typing.signature(types.boolean, *[ty]*2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg1])
|
||
|
|
arg2_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg2])
|
||
|
|
any_nan = builder.or_(arg1_nan, arg2_nan)
|
||
|
|
nan_result = builder.select(arg1_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_ge_arg2 = np_complex_ge_impl(context, builder, bcc_sig, args)
|
||
|
|
non_nan_result = builder.select(arg1_ge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_fmax_impl(context, builder, sig, args):
|
||
|
|
# fmax prefers non-nan (tries to return a non-nan).
|
||
|
|
# There is an extra caveat with complex numbers, as there is more
|
||
|
|
# than one type of nan. NumPy's docs state that the nan in the
|
||
|
|
# first argument is returned when both arguments are nans.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
ty = sig.args[0]
|
||
|
|
bc_sig = typing.signature(types.boolean, ty)
|
||
|
|
bcc_sig = typing.signature(types.boolean, *[ty]*2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg1])
|
||
|
|
arg2_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg2])
|
||
|
|
any_nan = builder.or_(arg1_nan, arg2_nan)
|
||
|
|
nan_result = builder.select(arg2_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_ge_arg2 = np_complex_ge_impl(context, builder, bcc_sig, args)
|
||
|
|
non_nan_result = builder.select(arg1_ge_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_smin_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_sle_arg2 = builder.icmp_signed('<=', arg1, arg2)
|
||
|
|
return builder.select(arg1_sle_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_umin_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_ule_arg2 = builder.icmp_unsigned('<=', arg1, arg2)
|
||
|
|
return builder.select(arg1_ule_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_minimum_impl(context, builder, sig, args):
|
||
|
|
# minimum prefers nan (tries to return a nan).
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = builder.fcmp_unordered('uno', arg1, arg1)
|
||
|
|
any_nan = builder.fcmp_unordered('uno', arg1, arg2)
|
||
|
|
nan_result = builder.select(arg1_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_le_arg2 = builder.fcmp_ordered('<=', arg1, arg2)
|
||
|
|
non_nan_result = builder.select(arg1_le_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_fmin_impl(context, builder, sig, args):
|
||
|
|
# fmin prefers non-nan (tries to return a non-nan).
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = builder.fcmp_unordered('uno', arg1, arg1)
|
||
|
|
any_nan = builder.fcmp_unordered('uno', arg1, arg2)
|
||
|
|
nan_result = builder.select(arg1_nan, arg2, arg1)
|
||
|
|
|
||
|
|
arg1_le_arg2 = builder.fcmp_ordered('<=', arg1, arg2)
|
||
|
|
non_nan_result = builder.select(arg1_le_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_minimum_impl(context, builder, sig, args):
|
||
|
|
# minimum prefers nan (tries to return a nan).
|
||
|
|
# There is an extra caveat with complex numbers, as there is more
|
||
|
|
# than one type of nan. NumPy's docs state that the nan in the
|
||
|
|
# first argument is returned when both arguments are nans.
|
||
|
|
# If only one nan is found, that nan is returned.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
ty = sig.args[0]
|
||
|
|
bc_sig = typing.signature(types.boolean, ty)
|
||
|
|
bcc_sig = typing.signature(types.boolean, *[ty]*2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg1])
|
||
|
|
arg2_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg2])
|
||
|
|
any_nan = builder.or_(arg1_nan, arg2_nan)
|
||
|
|
nan_result = builder.select(arg1_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_le_arg2 = np_complex_le_impl(context, builder, bcc_sig, args)
|
||
|
|
non_nan_result = builder.select(arg1_le_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_fmin_impl(context, builder, sig, args):
|
||
|
|
# fmin prefers non-nan (tries to return a non-nan).
|
||
|
|
# There is an extra caveat with complex numbers, as there is more
|
||
|
|
# than one type of nan. NumPy's docs state that the nan in the
|
||
|
|
# first argument is returned when both arguments are nans.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
ty = sig.args[0]
|
||
|
|
bc_sig = typing.signature(types.boolean, ty)
|
||
|
|
bcc_sig = typing.signature(types.boolean, *[ty]*2)
|
||
|
|
arg1, arg2 = args
|
||
|
|
arg1_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg1])
|
||
|
|
arg2_nan = np_complex_isnan_impl(context, builder, bc_sig, [arg2])
|
||
|
|
any_nan = builder.or_(arg1_nan, arg2_nan)
|
||
|
|
nan_result = builder.select(arg2_nan, arg1, arg2)
|
||
|
|
|
||
|
|
arg1_le_arg2 = np_complex_le_impl(context, builder, bcc_sig, args)
|
||
|
|
non_nan_result = builder.select(arg1_le_arg2, arg1, arg2)
|
||
|
|
|
||
|
|
return builder.select(any_nan, nan_result, non_nan_result)
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# NumPy floating point misc
|
||
|
|
|
||
|
|
def np_int_isnan_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return cgutils.false_bit
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_isnan_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return mathimpl.is_nan(builder, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_isnan_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
|
||
|
|
x, = args
|
||
|
|
ty, = sig.args
|
||
|
|
complex_val = context.make_complex(builder, ty, value=x)
|
||
|
|
return cmathimpl.is_nan(builder, complex_val)
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_isfinite_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return cgutils.true_bit
|
||
|
|
|
||
|
|
|
||
|
|
def np_datetime_isfinite_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return builder.icmp_unsigned('!=', args[0], npdatetime.NAT)
|
||
|
|
|
||
|
|
|
||
|
|
def np_datetime_isnat_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return builder.icmp_signed('==', args[0], npdatetime.NAT)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_isfinite_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return mathimpl.is_finite(builder, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_isfinite_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
x, = args
|
||
|
|
ty, = sig.args
|
||
|
|
complex_val = context.make_complex(builder, ty, value=x)
|
||
|
|
return cmathimpl.is_finite(builder, complex_val)
|
||
|
|
|
||
|
|
|
||
|
|
def np_int_isinf_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return cgutils.false_bit
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_isinf_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
return mathimpl.is_inf(builder, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_complex_isinf_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
x, = args
|
||
|
|
ty, = sig.args
|
||
|
|
complex_val = context.make_complex(builder, ty, value=x)
|
||
|
|
return cmathimpl.is_inf(builder, complex_val)
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_signbit_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1, return_type=types.boolean)
|
||
|
|
# there's no signbit intrinsic in LLVM, so just bitcast as int, mask the
|
||
|
|
# signbit and cmp against 0.
|
||
|
|
masks = {
|
||
|
|
types.float16: context.get_constant(types.uint16, 0x8000),
|
||
|
|
types.float32: context.get_constant(types.uint32, 0x80000000),
|
||
|
|
types.float64: context.get_constant(types.uint64, 0x8000000000000000),
|
||
|
|
}
|
||
|
|
arg_ty = sig.args[0]
|
||
|
|
arg_int_ty = getattr(types, f'uint{arg_ty.bitwidth}')
|
||
|
|
arg_ll_int_ty = context.get_value_type(arg_int_ty)
|
||
|
|
int_res = builder.and_(builder.bitcast(args[0], arg_ll_int_ty),
|
||
|
|
masks[arg_ty])
|
||
|
|
bool_res = builder.icmp_unsigned('!=', int_res, int_res.type(0))
|
||
|
|
return bool_res
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_copysign_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
return mathimpl.copysign_float_impl(context, builder, sig, args)
|
||
|
|
|
||
|
|
def np_real_nextafter_impl(context, builder, sig, args):
|
||
|
|
_check_arity_and_homogeneity(sig, args, 2)
|
||
|
|
|
||
|
|
dispatch_table = {
|
||
|
|
types.float32: 'numba_nextafterf',
|
||
|
|
types.float64: 'numba_nextafter',
|
||
|
|
}
|
||
|
|
|
||
|
|
return _dispatch_func_by_name_type(context, builder, sig, args,
|
||
|
|
dispatch_table, 'nextafter')
|
||
|
|
|
||
|
|
def np_real_spacing_impl(context, builder, sig, args):
|
||
|
|
# This is different to how NumPy does it, NumPy has a specialisation of
|
||
|
|
# nextafter called _next, which is used. See:
|
||
|
|
# https://github.com/numpy/numpy/blob/12c2b7dd62fc0c14b81c8892ed5f4f59cc94d09c/numpy/core/src/npymath/ieee754.c.src#L32-L38
|
||
|
|
# Numba elects to use `nextafter` for a similar behaviour to save
|
||
|
|
# translating this very involved function. Further, the NumPy comments note
|
||
|
|
# that there is a lot of redundancy present between the two.
|
||
|
|
_check_arity_and_homogeneity(sig, args, 1)
|
||
|
|
|
||
|
|
dispatch_table = {
|
||
|
|
types.float32: 'numba_nextafterf',
|
||
|
|
types.float64: 'numba_nextafter',
|
||
|
|
}
|
||
|
|
|
||
|
|
[ty] = sig.args
|
||
|
|
inner_sig = typing.signature(sig.return_type, ty, ty)
|
||
|
|
ll_ty = args[0].type
|
||
|
|
ll_inf = ll_ty(np.inf)
|
||
|
|
fnty = llvmlite.ir.FunctionType(ll_ty, [ll_ty, ll_ty])
|
||
|
|
fn = cgutils.insert_pure_function(builder.module, fnty,
|
||
|
|
name='llvm.copysign')
|
||
|
|
ll_sinf = builder.call(fn, [ll_inf, args[0]])
|
||
|
|
inner_args = args + [ll_sinf,]
|
||
|
|
nextafter = _dispatch_func_by_name_type(context, builder, inner_sig,
|
||
|
|
inner_args, dispatch_table,
|
||
|
|
'nextafter')
|
||
|
|
return builder.fsub(nextafter, args[0])
|
||
|
|
|
||
|
|
|
||
|
|
def np_real_ldexp_impl(context, builder, sig, args):
|
||
|
|
# this one is slightly different to other ufuncs.
|
||
|
|
# arguments are not homogeneous and second arg may come as
|
||
|
|
# an 'i' or an 'l'.
|
||
|
|
|
||
|
|
# the function expects the second argument to be have a C int type
|
||
|
|
x1, x2 = args
|
||
|
|
ty1, ty2 = sig.args
|
||
|
|
# note that types.intc should be equivalent to int_ that is
|
||
|
|
# 'NumPy's default int')
|
||
|
|
x2 = context.cast(builder, x2, ty2, types.intc)
|
||
|
|
f_fi_sig = typing.signature(ty1, ty1, types.intc)
|
||
|
|
return mathimpl.ldexp_impl(context, builder, f_fi_sig, (x1, x2))
|