diff --git a/pyopencl/array.py b/pyopencl/array.py index cb66121a5..b96c2a504 100644 --- a/pyopencl/array.py +++ b/pyopencl/array.py @@ -34,6 +34,7 @@ import numpy as np import pyopencl.elementwise as elementwise + import pyopencl as cl from pyopencl.compyte.array import ( as_strided as _as_strided, @@ -2930,10 +2931,28 @@ def if_positive(criterion, then_, else_, out=None, queue=None): return out +# }}} + + +# {{{ minimum/maximum + +@elwise_kernel_runner +def _minimum_maximum_backend(out, a, b, minmax): + from pyopencl.elementwise import get_minmaximum_kernel + return get_minmaximum_kernel(out.context, minmax, + out.dtype, + a.dtype if isinstance(a, Array) else np.dtype(type(a)), + b.dtype if isinstance(b, Array) else np.dtype(type(b)), + elementwise.get_argument_kind(a), + elementwise.get_argument_kind(b)) + def maximum(a, b, out=None, queue=None): """Return the elementwise maximum of *a* and *b*.""" - if isinstance(a, SCALAR_CLASSES) and isinstance(b, SCALAR_CLASSES): + + a_is_scalar = np.isscalar(a) + b_is_scalar = np.isscalar(b) + if a_is_scalar and b_is_scalar: result = np.maximum(a, b) if out is not None: out[...] = result @@ -2941,36 +2960,25 @@ def maximum(a, b, out=None, queue=None): return result - # silly, but functional - a_is_array = isinstance(a, Array) - b_is_array = isinstance(b, Array) - if a_is_array: - criterion = a.mul_add(1, b, -1, queue=queue) - elif b_is_array: - criterion = b.mul_add(-1, a, 1, queue=queue) - else: - raise AssertionError - - # {{{ propagate NaNs - - if a_is_array: - a_with_nan = a.mul_add(1, b, 0, queue=queue) - else: - a_with_nan = b.mul_add(0, a, 1, queue=queue) + queue = queue or a.queue or b.queue - if b_is_array: - b_with_nan = b.mul_add(1, a, 0, queue=queue) - else: - b_with_nan = a.mul_add(0, b, 1, queue=queue) + if out is None: + out_dtype = _get_common_dtype(a, b, queue) + if not a_is_scalar: + out = a._new_like_me(out_dtype, queue) + elif not b_is_scalar: + out = b._new_like_me(out_dtype, queue) - # }}} + out.add_event(_minimum_maximum_backend(out, a, b, queue=queue, minmax="max")) - return if_positive(criterion, a_with_nan, b_with_nan, queue=queue, out=out) + return out def minimum(a, b, out=None, queue=None): """Return the elementwise minimum of *a* and *b*.""" - if isinstance(a, SCALAR_CLASSES) and isinstance(b, SCALAR_CLASSES): + a_is_scalar = np.isscalar(a) + b_is_scalar = np.isscalar(b) + if a_is_scalar and b_is_scalar: result = np.minimum(a, b) if out is not None: out[...] = result @@ -2978,31 +2986,18 @@ def minimum(a, b, out=None, queue=None): return result - # silly, but functional - a_is_array = isinstance(a, Array) - b_is_array = isinstance(b, Array) - if a_is_array: - criterion = a.mul_add(1, b, -1, queue=queue) - elif b_is_array: - criterion = b.mul_add(-1, a, 1, queue=queue) - else: - raise AssertionError - - # {{{ propagate NaNs + queue = queue or a.queue or b.queue - if a_is_array: - a_with_nan = a.mul_add(1, b, 0, queue=queue) - else: - a_with_nan = b.mul_add(0, a, 1, queue=queue) - - if b_is_array: - b_with_nan = b.mul_add(1, a, 0, queue=queue) - else: - b_with_nan = a.mul_add(0, b, 1, queue=queue) + if out is None: + out_dtype = _get_common_dtype(a, b, queue) + if not a_is_scalar: + out = a._new_like_me(out_dtype, queue) + elif not b_is_scalar: + out = b._new_like_me(out_dtype, queue) - # }}} + out.add_event(_minimum_maximum_backend(out, a, b, queue=queue, minmax="min")) - return if_positive(criterion, b_with_nan, a_with_nan, queue=queue, out=out) + return out # }}} diff --git a/pyopencl/elementwise.py b/pyopencl/elementwise.py index 877504038..3132bb683 100644 --- a/pyopencl/elementwise.py +++ b/pyopencl/elementwise.py @@ -27,6 +27,8 @@ """ +from typing import Any, Tuple +import enum from pyopencl.tools import context_dependent_memoize import numpy as np import pyopencl as cl @@ -356,6 +358,38 @@ def build_inner(self, context, type_aliases=(), var_values=(), # }}} +# {{{ argument kinds + +class ArgumentKind(enum.Enum): + ARRAY = enum.auto() + DEV_SCALAR = enum.auto() + SCALAR = enum.auto() + + +def get_argument_kind(v: Any) -> ArgumentKind: + from pyopencl.array import Array + if isinstance(v, Array): + if v.shape == (): + return ArgumentKind.DEV_SCALAR + else: + return ArgumentKind.ARRAY + else: + return ArgumentKind.SCALAR + + +def get_decl_and_access_for_kind(name: str, kind: ArgumentKind) -> Tuple[str, str]: + if kind == ArgumentKind.ARRAY: + return f"*{name}", f"{name}[i]" + elif kind == ArgumentKind.SCALAR: + return f"{name}", name + elif kind == ArgumentKind.DEV_SCALAR: + return f"*{name}", f"{name}[0]" + else: + raise AssertionError() + +# }}} + + # {{{ kernels supporting array functionality @context_dependent_memoize @@ -950,6 +984,32 @@ def get_ldexp_kernel(context, out_dtype=np.float32, sig_dtype=np.float32, name="ldexp_kernel") +@context_dependent_memoize +def get_minmaximum_kernel(context, minmax, dtype_z, dtype_x, dtype_y, + kind_x: ArgumentKind, kind_y: ArgumentKind): + if dtype_z.kind == "f": + reduce_func = f"f{minmax}_nanprop" + elif dtype_z.kind in "iu": + reduce_func = minmax + else: + raise TypeError("unsupported dtype specified") + + tp_x = dtype_to_ctype(dtype_x) + tp_y = dtype_to_ctype(dtype_y) + tp_z = dtype_to_ctype(dtype_z) + decl_x, acc_x = get_decl_and_access_for_kind("x", kind_x) + decl_y, acc_y = get_decl_and_access_for_kind("y", kind_y) + + return get_elwise_kernel(context, + f"{tp_z} *z, {tp_x} {decl_x}, {tp_y} {decl_y}", + f"z[i] = {reduce_func}({acc_x}, {acc_y})", + name=f"{minmax}imum", + preamble=""" + #define fmin_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmin(a, b) + #define fmax_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmax(a, b) + """) + + @context_dependent_memoize def get_bessel_kernel(context, which_func, out_dtype=np.float64, order_dtype=np.int32, x_dtype=np.float64): @@ -1070,4 +1130,4 @@ def get_if_positive_kernel( # }}} -# vim: fdm=marker:filetype=pyopencl +# vim: fdm=marker diff --git a/pyopencl/reduction.py b/pyopencl/reduction.py index 92bf51eb1..c1f2f2ef2 100644 --- a/pyopencl/reduction.py +++ b/pyopencl/reduction.py @@ -628,6 +628,13 @@ def get_subset_dot_kernel(ctx, dtype_out, dtype_subset, dtype_a=None, dtype_b=No })) +_MINMAX_PREAMBLE = """ +#define MY_INFINITY (1./0) +#define fmin_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmin(a, b) +#define fmax_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmax(a, b) +""" + + def get_minmax_neutral(what, dtype): dtype = np.dtype(dtype) if issubclass(dtype.type, np.inexact): @@ -660,11 +667,7 @@ def get_minmax_kernel(ctx, what, dtype): reduce_expr=f"{reduce_expr}", arguments="const {tp} *in".format( tp=dtype_to_ctype(dtype), - ), preamble=""" - #define MY_INFINITY (1./0) - #define fmin_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmin(a, b) - #define fmax_nanprop(a, b) (isnan(a) || isnan(b)) ? a+b : fmax(a, b) - """) + ), preamble=_MINMAX_PREAMBLE) @context_dependent_memoize @@ -686,7 +689,7 @@ def get_subset_minmax_kernel(ctx, what, dtype, dtype_subset): "tp": dtype_to_ctype(dtype), "tp_lut": dtype_to_ctype(dtype_subset), }), - preamble="#define MY_INFINITY (1./0)") + preamble=_MINMAX_PREAMBLE) # }}} diff --git a/test/test_array.py b/test/test_array.py index 1f9cdcfe4..db59902f2 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -1847,12 +1847,31 @@ def test_branch_operations_on_pure_scalars(): cl_array.maximum, cl_array.minimum, ]) -def test_branch_operations_on_nans(ctx_factory, op): +@pytest.mark.parametrize("special_a", [ + np.nan, + np.inf, + -np.inf, +]) +@pytest.mark.parametrize("special_b", [ + np.nan, + np.inf, + -np.inf, + None +]) +def test_branch_operations_on_nans(ctx_factory, op, special_a, special_b): ctx = ctx_factory() cq = cl.CommandQueue(ctx) - x_np = np.array([np.nan, 1., np.nan, 2.], dtype=np.float64) - y_np = np.array([np.nan, np.nan, 1., 3.], dtype=np.float64) + def sb_or(x): + if special_b is None: + return x + else: + return special_b + + x_np = np.array([special_a, sb_or(1.), special_a, sb_or(2.), sb_or(3.)], + dtype=np.float64) + y_np = np.array([special_a, special_a, sb_or(1.), sb_or(3.), sb_or(2.)], + dtype=np.float64) x_cl = cl_array.to_device(cq, x_np) y_cl = cl_array.to_device(cq, y_np)