codon/stdlib/numpy/fusion.codon

811 lines
19 KiB
Python

# Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
import util
import zmath
from .ndarray import ndarray
from .routines import asarray, broadcast_to, empty_like
# Utility
def _cast(x, dtype: type):
return util.cast(x, dtype)
def _coerce(dtype1: type, dtype2: type):
return util.coerce(dtype1, dtype2)
def _count(shape):
return util.count(shape)
def _contig_match(arrs):
def keep_arrays(arrs):
if staticlen(arrs) == 0:
return ()
else:
if hasattr(arrs[0], "_contig"):
return (arrs[0],) + keep_arrays(arrs[1:])
else:
return keep_arrays(arrs[1:])
arrs = keep_arrays(arrs)
if staticlen(arrs) == 0:
return True
all_cc = True
all_fc = True
for i in staticrange(staticlen(arrs)):
a = arrs[i]
cc, fc = a._contig
if a.ndim != arrs[0].ndim:
return False
else:
if a.shape != arrs[0].shape:
return False
all_cc = all_cc and cc
all_fc = all_fc and fc
return all_cc or all_fc
def _ptrset(p: Ptr[T], x: T, T: type):
p[0] = x
def _loop_alloc(arrays, func, extra, dtype: type):
return ndarray._loop(arrays, func, alloc=Tuple[dtype], extra=extra)[0]
def _loop_basic(arrays, func, extra):
ndarray._loop(arrays, func, extra=extra)
def _broadcast(sh1, sh2):
def bc_one(sh1, sh2, i: Static[int]):
a = sh1[i]
b = sh2[i]
if a == 1 or b == 1 or a == b:
return max(a, b)
else:
raise ValueError(f"operands could not be broadcast together with shapes {sh1} {sh2}")
def bc_same(sh1, sh2):
return tuple(bc_one(sh1, sh2, i) for i in staticrange(staticlen(sh1)))
N1: Static[int] = staticlen(sh1)
N2: Static[int] = staticlen(sh2)
if N1 == 0:
return sh2
elif N2 == 0:
return sh1
elif N1 > N2:
return sh1[:-N2] + bc_same(sh1[-N2:], sh2)
elif N1 < N2:
return sh2[:-N1] + bc_same(sh1, sh2[-N1:])
else:
return bc_same(sh1, sh2)
def _matmul_shape(x1, x2):
x1d: Static[int] = staticlen(x1)
x2d: Static[int] = staticlen(x2)
if x1d == 0:
return x2
if x2d == 0:
return x1
if x1d == 1:
y1 = (1,) + x1
else:
y1 = x1
if x2d == 1:
y2 = x2 + (1,)
else:
y2 = x2
y1d: Static[int] = staticlen(y1)
y2d: Static[int] = staticlen(y2)
base1s = y1[:-2]
base2s = y2[:-2]
mat1s = y1[-2:]
mat2s = y2[-2:]
m = mat1s[0]
k = mat1s[1]
n = mat2s[1]
if k != mat2s[0]:
raise ValueError("matmul: last dimension of first argument does not "
"match second-to-last dimension of second argument")
ans_base = _broadcast(base1s, base2s)
if x1d == 1 and x2d == 1:
return ans_base
elif x1d == 1:
return ans_base + (mat2s[1],)
elif x2d == 1:
return ans_base + (mat1s[0],)
else:
return ans_base + (mat1s[0], mat2s[1])
def _create(like, shape, dtype: type):
return empty_like(like, shape=shape, dtype=dtype)
def _shape(x):
if hasattr(x, "shape"):
return x.shape
else:
return ()
def _free(x):
util.free(x.data)
def _apply_vectorized_loop_unary(arr, out, func: Static[str]):
if arr.ndim == 0 or out.ndim == 0 or arr.ndim > out.ndim:
compile_error("[internal error] bad array dims for vectorized loop")
if out.ndim == 1:
util.call_vectorized_loop(arr.data, arr.strides[0], Ptr[arr.dtype](),
0, out.data, out.strides[0], out.size, func)
return
shape = arr.shape
arr = broadcast_to(arr, shape)
if arr._contig_match(out):
s = util.sizeof(out.dtype)
util.call_vectorized_loop(arr.data, s, Ptr[arr.dtype](), 0, out.data,
s, out.size, func)
else:
# Find smallest stride to use in vectorized loop
arr_strides = arr.strides
out_strides = out.strides
n = 0
si = 0
so = 0
loop_axis = -1
for i in staticrange(arr.ndim):
if shape[i] > 1 and (loop_axis == -1 or abs(arr_strides[i]) < abs(si)):
n = shape[i]
si = arr_strides[i]
so = out_strides[i]
loop_axis = i
if loop_axis == -1:
n = shape[0]
si = arr_strides[0]
so = out_strides[0]
loop_axis = 0
for idx in util.multirange(util.tuple_delete(shape, loop_axis)):
idx1 = util.tuple_insert(idx, loop_axis, 0)
p = arr._ptr(idx1)
q = out._ptr(idx1)
util.call_vectorized_loop(p, si, Ptr[arr.dtype](), 0, q, so, n,
func)
def _apply_vectorized_loop_binary(arr1, arr2, out, func: Static[str]):
if (arr1.ndim == 0 and arr2.ndim == 0) or out.ndim == 0 or arr1.ndim > out.ndim or arr2.ndim > out.ndim:
compile_error("[internal error] bad array dims for vectorized loop")
if arr1.ndim == 0:
st1 = 0
else:
st1 = arr1.strides[0]
if arr2.ndim == 0:
st2 = 0
else:
st2 = arr2.strides[0]
if out.ndim == 1:
util.call_vectorized_loop(arr1.data, st1, arr2.data,
st2, out.data, out.strides[0],
out.size, func)
return
shape = out.shape
arr1 = broadcast_to(arr1, shape)
arr2 = broadcast_to(arr2, shape)
if arr1._contig_match(out) and arr2._contig_match(out):
s = util.sizeof(out.dtype)
util.call_vectorized_loop(arr1.data, s, arr2.data, s, out.data, s, out.size, func)
else:
# Find smallest stride to use in vectorized loop
arr1_strides = arr1.strides
arr2_strides = arr2.strides
out_strides = out.strides
n = 0
si1 = 0
si2 = 0
so = 0
loop_axis = -1
for i in staticrange(arr1.ndim):
if shape[i] > 1 and (loop_axis == -1 or abs(arr1_strides[i]) < abs(si1)):
n = shape[i]
si1 = arr1_strides[i]
si2 = arr2_strides[i]
so = out_strides[i]
loop_axis = i
if loop_axis == -1:
n = shape[0]
si1 = arr1_strides[0]
si2 = arr2_strides[0]
so = out_strides[0]
loop_axis = 0
for idx in util.multirange(util.tuple_delete(shape, loop_axis)):
idx1 = util.tuple_insert(idx, loop_axis, 0)
p1 = arr1._ptr(idx1)
p2 = arr2._ptr(idx1)
q = out._ptr(idx1)
util.call_vectorized_loop(p1, si1, p2, si2, q, so, n, func)
# Operations
@inline
def _pos(x):
return +x
@inline
def _neg(x):
return -x
@inline
def _invert(x):
return ~x
@inline
def _abs(x):
return abs(x)
@inline
def _transpose(x):
return x.T
@inline
def _add(x, y):
return x + y
@inline
def _sub(x, y):
return x - y
@inline
def _mul(x, y):
return x * y
@inline
def _matmul(x, y):
return x @ y
@inline
def _true_div(x, y):
return x / y
@inline
def _floor_div(x, y):
X = type(x)
Y = type(y)
if isinstance(X, Int) and isinstance(Y, Int):
return util.pydiv(x, y)
else:
return x // y
@inline
def _mod(x, y):
X = type(x)
Y = type(y)
if isinstance(X, Int) and isinstance(Y, Int):
return util.pymod(x, y)
elif ((X is float and Y is float) or
(X is float32 and Y is float32) or
(X is float16 and Y is float16)):
return util.pyfmod(x, y)
else:
return x % y
@inline
def _fmod(x, y):
X = type(x)
Y = type(y)
if isinstance(X, Int) and isinstance(Y, Int):
return util.cmod_int(x, y)
elif ((X is float and Y is float) or
(X is float32 and Y is float32) or
(X is float16 and Y is float16)):
return util.cmod(x, y)
else:
return x % y
@inline
def _pow(x, y):
return x ** y
@inline
def _lshift(x, y):
return x << y
@inline
def _rshift(x, y):
return x >> y
@inline
def _and(x, y):
return x & y
@inline
def _or(x, y):
return x | y
@inline
def _xor(x, y):
return x ^ y
@inline
def _eq(x, y):
return x == y
@inline
def _ne(x, y):
return x != y
@inline
def _lt(x, y):
return x < y
@inline
def _le(x, y):
return x <= y
@inline
def _gt(x, y):
return x > y
@inline
def _ge(x, y):
return x >= y
def _apply(x, f, f_complex = None):
if f_complex is not None and (isinstance(x, complex) or isinstance(x, complex64)):
return f_complex(x)
elif isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return f(x)
else:
return f(util.to_float(x))
def _apply2(x, y, f, f_complex = None):
if type(x) is not type(y):
compile_error("type mismatch in util")
if f_complex is not None and (isinstance(x, complex) or isinstance(x, complex64)):
return f_complex(x, y)
elif isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return f(x, y)
else:
return f(util.to_float(x), util.to_float(y))
def _fabs(x):
return _apply(x, util.fabs)
def _rint(x):
def rint_complex(x):
C = type(x)
return C(util.rint(x.real), util.rint(x.imag))
return _apply(x, util.rint, rint_complex)
def _exp(x):
return _apply(x, util.exp, zmath.exp)
def _exp2(x):
return _apply(x, util.exp2, zmath.exp2)
def _expm1(x):
return _apply(x, util.expm1, zmath.expm1)
def _log(x):
return _apply(x, util.log, zmath.log)
def _log2(x):
return _apply(x, util.log2, zmath.log2)
def _log10(x):
return _apply(x, util.log10, zmath.log10)
def _log1p(x):
return _apply(x, util.log1p, zmath.log1p)
def _sqrt(x):
return _apply(x, util.sqrt, zmath.sqrt)
def _cbrt(x):
return _apply(x, util.cbrt)
def _square(x):
return x * x
def _sin(x):
return _apply(x, util.sin, zmath.sin)
def _cos(x):
return _apply(x, util.cos, zmath.cos)
def _tan(x):
return _apply(x, util.tan, zmath.tan)
def _arcsin(x):
return _apply(x, util.asin, zmath.asin)
def _arccos(x):
return _apply(x, util.acos, zmath.acos)
def _arctan(x):
return _apply(x, util.atan, zmath.atan)
def _sinh(x):
return _apply(x, util.sinh, zmath.sinh)
def _cosh(x):
return _apply(x, util.cosh, zmath.cosh)
def _tanh(x):
return _apply(x, util.tanh, zmath.tanh)
def _arcsinh(x):
return _apply(x, util.asinh, zmath.asinh)
def _arccosh(x):
return _apply(x, util.acosh, zmath.acosh)
def _arctanh(x):
return _apply(x, util.atanh, zmath.atanh)
def _rad2deg(x):
r2d = 180.0 / util.PI
x = util.to_float(x)
F = type(x)
return x * F(r2d)
def _deg2rad(x):
d2r = util.PI / 180.0
x = util.to_float(x)
F = type(x)
return x * F(d2r)
def _arctan2(x, y):
return _apply2(x, y, util.atan2)
def _hypot(x, y):
return _apply2(x, y, util.hypot)
def _logaddexp(x, y):
return _apply2(x, y, util.logaddexp)
def _logaddexp2(x, y):
return _apply2(x, y, util.logaddexp2)
def _isnan(x):
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return util.isnan(x)
elif isinstance(x, complex) or isinstance(x, complex64):
return util.isnan(x.real) or util.isnan(x.imag)
else:
return False
def _isinf(x):
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return util.isinf(x)
elif isinstance(x, complex) or isinstance(x, complex64):
return util.isinf(x.real) or util.isinf(x.imag)
else:
return False
def _isfinite(x):
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return util.isfinite(x)
elif isinstance(x, complex) or isinstance(x, complex64):
return util.isfinite(x.real) and util.isfinite(x.imag)
else:
return True
def _signbit(x):
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
return util.signbit(x)
else:
T = type(x)
return x < T()
def _copysign(x, y):
return _apply2(x, y, util.copysign)
def _nextafter(x, y):
return _apply2(x, y, util.nextafter)
def _floor(x):
return _apply(x, util.floor)
def _ceil(x):
return _apply(x, util.ceil)
def _trunc(x):
return _apply(x, util.trunc)
def _sign(x):
def sign1(x):
T = type(x)
if x < T(0):
return T(-1)
elif x > T(0):
return T(1)
else:
return x
if isinstance(x, complex):
if _isnan(x):
return complex(util.nan64(), 0.0)
return complex(sign1(x.real), 0) if x.real else complex(sign1(x.imag), 0)
elif isinstance(x, complex64):
if _isnan(x):
return complex64(util.nan64(), 0.0)
return complex64(sign1(x.real), 0) if x.real else complex64(sign1(x.imag), 0)
else:
return sign1(x)
def _heaviside(x, y):
def heaviside(x, y):
if isinstance(x, float16) and isinstance(y, float16):
if x < float16(0):
return float16(0)
elif x > float16(0):
return float16(1)
elif x == float16(0):
return y
else:
return x
elif isinstance(x, float32) and isinstance(y, float32):
if x < float32(0):
return float32(0)
elif x > float32(0):
return float32(1)
elif x == float32(0):
return y
else:
return x
elif isinstance(x, float) and isinstance(y, float):
if x < 0:
return 0.0
elif x > 0:
return 1.0
elif x == 0.0:
return y
else:
return x
return _apply2(x, y, heaviside)
def _conj(x):
if isinstance(x, complex) or isinstance(x, complex64):
return x.conjugate()
else:
return x
def _gcd(x, y):
''' # fails with optionals
if not (
isinstance(x, int) or
isinstance(x, Int) or
isinstance(x, UInt) or
isinstance(x, byte)
):
compile_error("gcd/lcm can only be used on integral types")
'''
while x:
z = x
x = y % x
y = z
return y
def _lcm(x, y):
gcd = _gcd(x, y)
return x // gcd * y if gcd else 0
def _reciprocal(x: T, T: type):
if (
isinstance(x, int) or
isinstance(x, Int) or
isinstance(x, UInt) or
isinstance(x, byte)
):
return T(1) // x
else:
return T(1) / x
def _logical_and(x, y):
return bool(x) and bool(y)
def _logical_or(x, y):
return bool(x) or bool(y)
def _logical_xor(x, y):
return bool(x) ^ bool(y)
def _logical_not(x):
return not bool(x)
def _coerce_types_for_minmax(x, y):
if isinstance(x, complex):
if isinstance(y, complex64):
return x, complex(y)
elif not isinstance(y, complex):
return x, complex(util.cast(y, float))
elif isinstance(x, complex64):
if isinstance(y, complex):
return complex(x), y
elif not isinstance(y, complex64):
return complex(x), complex(util.cast(y, float))
if isinstance(y, complex):
if isinstance(x, complex64):
return complex(x), y
elif not isinstance(x, complex):
return complex(util.cast(x, float)), y
elif isinstance(y, complex64):
if isinstance(x, complex):
return x, complex(y)
elif not isinstance(x, complex64):
return complex(util.cast(x, float)), complex(y)
T = type(util.coerce(type(x), type(y)))
return util.cast(x, T), util.cast(y, T)
def _compare_le(x, y):
if isinstance(x, complex) or isinstance(x, complex64):
return (x.real, x.imag) <= (y.real, y.imag)
else:
return x <= y
def _compare_ge(x, y):
if isinstance(x, complex) or isinstance(x, complex64):
return (x.real, x.imag) >= (y.real, y.imag)
else:
return x >= y
def _maximum(x, y):
x, y = _coerce_types_for_minmax(x, y)
if _isnan(x):
return x
if _isnan(y):
return y
return x if _compare_ge(x, y) else y
def _minimum(x, y):
x, y = _coerce_types_for_minmax(x, y)
if _isnan(x):
return x
if _isnan(y):
return y
return x if _compare_le(x, y) else y
def _fmax(x, y):
x, y = _coerce_types_for_minmax(x, y)
if _isnan(y):
return x
if _isnan(x):
return y
return x if _compare_ge(x, y) else y
def _fmin(x, y):
x, y = _coerce_types_for_minmax(x, y)
if _isnan(y):
return x
if _isnan(x):
return y
return x if _compare_le(x, y) else y
def _divmod_float(x, y):
F = type(x)
mod = util.cmod(x, y)
if not y:
return util.cdiv(x, y), mod
div = util.cdiv(x - mod, y)
if mod:
if (y < F(0)) != (mod < F(0)):
mod += y
div -= F(1)
else:
mod = util.copysign(F(0), y)
floordiv = F()
if div:
floordiv = util.floor(div)
if div - floordiv > F(0.5):
floordiv += F(1)
else:
floordiv = util.copysign(F(0), util.cdiv(x, y))
return floordiv, mod
def _divmod(x, y):
if isinstance(x, float16) and isinstance(y, float16):
return _divmod_float(x, y)
if isinstance(x, float32) and isinstance(y, float32):
return _divmod_float(x, y)
if isinstance(x, float) or isinstance(y, float):
return _divmod_float(util.cast(x, float), util.cast(y, float))
return (x // y, x % y)
def _modf(x):
return _apply(x, util.modf)
def _frexp(x):
def frexp(x):
a, b = util.frexp(x)
return a, i32(b)
return _apply(x, frexp)
def _spacing16(h: float16):
h_u16 = util.bitcast(h, u16)
h_exp = h_u16 & u16(0x7c00)
h_sig = h_u16 & u16(0x03ff)
if h_exp == u16(0x7c00):
return util.nan16()
elif h_u16 == u16(0x7bff):
return util.inf16()
elif (h_u16 & u16(0x8000)) and not h_sig:
if h_exp > u16(0x2c00):
return util.bitcast(h_exp - u16(0x2c00), float16)
elif h_exp > u16(0x0400):
return util.bitcast(u16(1) << ((h_exp >> u16(10)) - u16(2)), float16)
else:
return util.bitcast(u16(0x0001), float16)
elif h_exp > u16(0x2800):
return util.bitcast(h_exp - u16(0x2800), float16)
elif h_exp > u16(0x0400):
return util.bitcast(u16(1) << ((h_exp >> u16(10)) - u16(1)), float16)
else:
return util.bitcast(u16(0x0001), float16)
def _spacing(x):
if isinstance(x, float16):
return _spacing16(x)
elif isinstance(x, float32):
if util.isinf32(x):
return util.nan32()
p = util.inf32() if x >= float32(0) else -util.inf32()
return util.nextafter32(x, util.inf32()) - x
elif isinstance(x, float):
x = util.cast(x, float)
if util.isinf64(x):
return util.nan64()
p = util.inf64() if x >= 0 else -util.inf64()
return util.nextafter64(x, p) - x
else:
return _spacing(util.to_float(x))