Implement dedicated kernels for {min,max}imum

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,79 +2931,73 @@ 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
             return out
 
         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
             return out
 
         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
 
 # }}}
 

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

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)
 
 # }}}
 

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)