Faster type checking on numpy

docs/source/conf.py

@@ -42,7 +42,10 @@ def _defining_comment(name, obj):
     from sphinx.pycode import ModuleAnalyzer
 
     for modname, module in list(sys.modules.items()):
-        if not modname.startswith("cunumpy._") or getattr(module, name, None) is not obj:
+        if (
+            not modname.startswith("cunumpy._")
+            or getattr(module, name, None) is not obj
+        ):
             continue
         try:
             docs = ModuleAnalyzer.for_module(modname).find_attr_docs()

docs/source/guides/solvers.md

@@ -77,8 +77,8 @@ least-squares variants), `BSpline`, `make_interp_spline`, `make_lsq_spline`,
 
 ```python
 interp = xp.scipy.interpolate
-bx = interp.make_interp_spline(x, values, k=3)       # fit along x (axis 0)
-by = interp.make_interp_spline(y, bx.c.T, k=3)       # then along y
+bx = interp.make_interp_spline(x, values, k=3)  # fit along x (axis 0)
+by = interp.make_interp_spline(y, bx.c.T, k=3)  # then along y
 spline = interp.NdBSpline((bx.t, by.t), by.c.T, (3, 3))
 # RectBivariateSpline clamps points to the grid; NdBSpline extrapolates
 points = xp.stack([xp.clip(R, x[0], x[-1]), xp.clip(Z, y[0], y[-1])], axis=-1)
@@ -106,6 +106,7 @@ the secant method on all of them at once, with the steps of
 def residual(r):
     return psi(R0 + r * cos_theta, Z0 + r * sin_theta) - levels
 
+
 r = xp.optimize.newton(residual, xp.full(levels.shape, 0.3))
 ```
 
@@ -124,9 +125,11 @@ host copy of that data:
 ```python
 B = equilibrium.b_field_on_grid()  # device array
 
+
 def residual(x):
     return interpolate(B, x) - b_target
 
+
 x = xp.optimize.fsolve(residual, xp.asarray([0.5, 0.0]))  # device array
 ```
 

docs/source/kernels/cuda-headers.md

@@ -185,12 +185,16 @@ same box form a contiguous range, the starting point of tree builds on the
 GPU. On the host:
 
 ```python
-keys = xp.algorithms.morton_keys(positions, lower, upper, levels)  # (n, 2|3) -> (n,) uint64
+keys = xp.algorithms.morton_keys(
+    positions, lower, upper, levels
+)  # (n, 2|3) -> (n,) uint64
 keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
 node = keys >> np.uint64(ndim * (levels - level))  # node index at `level`
 cells = xp.algorithms.morton_decode(node, ndim)  # its integer coordinates
 key = xp.algorithms.morton_encode(ix, iy)  # from integer cells
-scales = xp.algorithms.morton_scales(lower, upper, levels)  # 2**levels / (upper - lower)
+scales = xp.algorithms.morton_scales(
+    lower, upper, levels
+)  # 2**levels / (upper - lower)
 ```
 
 In a kernel:

src/cunumpy/_dispatch.py

@@ -29,6 +29,7 @@
 from typing import Any
 
 import array_api_compat
+import numpy as np
 
 from cunumpy._cuda_kernel import CudaKernel, _compile_in_threads
 from cunumpy._kernel import (
@@ -51,8 +52,15 @@
 _is_device_array = array_api_compat.is_cupy_array
 
 
+#: Argument types that are never on the device: checked first, since a kernel
+#: call tests every argument and most are host arrays or scalars.
+_HOST_TYPES = frozenset({np.ndarray, float, int, bool, np.float64, np.int64, np.bool_})
+
+
 def _on_device(arg: Any) -> bool:
     """Whether `arg` is a CuPy array or a CUDA argument object (``__cuda_args__``)."""
+    if type(arg) in _HOST_TYPES:
+        return False
     return _is_device_array(arg) or callable(getattr(type(arg), "__cuda_args__", None))
 
 

src/cunumpy/_kernel.py

@@ -19,9 +19,9 @@
 import inspect
 import os
 import warnings
-from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
-from types import ModuleType
+from types import ModuleType, TracebackType
 from typing import Any
 
 import array_api_compat
@@ -1000,6 +1000,17 @@ def as_kernel_array(
     >>> xp.kernels.as_kernel_array([1, 2], like=grid, dtype=float)
     array([1., 2.])
     """
+    # fast path: a C-contiguous NumPy array of the dtype for a host kernel, the
+    # common case, without the conversions below (subclasses take the full path)
+    if (
+        type(value) is np.ndarray
+        and (type(like) is np.ndarray or not is_gpu(like))
+        and (dtype is None or value.dtype == dtype)
+        and value.flags.c_contiguous
+        and value.ndim > 0
+        and (value.ndim == 1 or 1 not in value.shape)
+    ):
+        return value
     if is_gpu(like):
         import cupy
 
@@ -1041,10 +1052,39 @@ def _same_memory(buffer: Any, out: Any) -> bool:
     )
 
 
-@contextmanager
+class _KernelOutput:
+    """The context manager of :func:`kernel_output` (a class: cheaper than a generator)."""
+
+    __slots__ = ("_buffer", "_out")
+
+    def __init__(self, out: Any, buffer: Any) -> None:
+        self._out = out
+        self._buffer = buffer
+
+    def __enter__(self) -> Any:
+        return self._buffer
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if exc_type is not None:
+            return
+        buffer, out = self._buffer, self._out
+        if not _same_memory(buffer, out):
+            if is_gpu(out) and not is_gpu(buffer):
+                out[...] = to_cupy(buffer)
+            elif not is_gpu(out) and is_gpu(buffer):
+                out[...] = to_numpy(buffer)
+            else:
+                out[...] = buffer
+
+
 def kernel_output(
     out: Any, like: Any, dtype: Any = None, *, strided: bool = False
-) -> Generator[Any]:
+) -> _KernelOutput:
     """Provide the output buffer for the kernel chosen for `like`, copied into `out`.
 
     The buffer is `out` itself if :func:`~cunumpy.kernels.as_kernel_array`
@@ -1062,10 +1102,10 @@ def kernel_output(
     strided : bool, optional
         As for :func:`~cunumpy.kernels.as_kernel_array`. Default False.
 
-    Yields
-    ------
-    array
-        The buffer to pass to the kernel.
+    Returns
+    -------
+    context manager
+        Yields the buffer to pass to the kernel.
 
     Examples
     --------
@@ -1075,15 +1115,7 @@ def kernel_output(
     >>> out
     array([1.5, 1.5], dtype=float32)
     """
-    buffer = as_kernel_array(out, like, dtype, strided=strided)
-    yield buffer
-    if not _same_memory(buffer, out):
-        if is_gpu(out) and not is_gpu(buffer):
-            out[...] = to_cupy(buffer)
-        elif not is_gpu(out) and is_gpu(buffer):
-            out[...] = to_numpy(buffer)
-        else:
-            out[...] = buffer
+    return _KernelOutput(out, as_kernel_array(out, like, dtype, strided=strided))
 
 
 _SCALAR_ANNOTATIONS = {"int", "float", "bool", "complex", "str"}

src/cunumpy/xp.py

@@ -21,6 +21,11 @@
 
 import array_api_compat
 import array_api_compat.numpy as np
+import numpy
+
+#: The exact type of a host array; a fast path checks ``type(a) is _NDARRAY`` before
+#: the slower ``array_api_compat`` tests (subclasses take the full path).
+_NDARRAY = numpy.ndarray
 
 from cunumpy._transfers import _ACTIVE as _COUNTERS
 from cunumpy._transfers import (
@@ -493,6 +498,8 @@ def synchronize() -> None:
 
 def _to_numpy(array: Any) -> np.ndarray:
     """`to_numpy` without transfer counting, for internal use."""
+    if type(array) is _NDARRAY:
+        return array
     if get_array_backend(array) == "cupy":
         return array.get(order="A")
 
@@ -533,6 +540,8 @@ def to_numpy(array: Any) -> np.ndarray:
     >>> xp.to_numpy([1, 2])
     array([1, 2])
     """
+    if type(array) is _NDARRAY:
+        return array
     result = _to_numpy(array)
     if _COUNTERS and get_array_backend(array) == "cupy":
         _record("to_host", f"to_numpy({_describe(array)})", nbytes=_nbytes(result))
@@ -719,6 +728,8 @@ def get_array_backend(array: Any) -> BackendType:
     >>> xp.get_array_backend(np.zeros(3))
     'numpy'
     """
+    if type(array) is _NDARRAY:
+        return "numpy"
     return "cupy" if array_api_compat.is_cupy_array(array) else "numpy"