diff --git a/docs/source/conf.py b/docs/source/conf.py index 5f966c0..96f089f 100644 --- a/docs/source/conf.py +++ b/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() diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index ee013fe..c13277e 100644 --- a/docs/source/guides/solvers.md +++ b/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 ``` diff --git a/docs/source/kernels/cuda-headers.md b/docs/source/kernels/cuda-headers.md index 6c2115f..1d2eea3 100644 --- a/docs/source/kernels/cuda-headers.md +++ b/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: diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index d7bb198..cebb9ca 100644 --- a/src/cunumpy/_dispatch.py +++ b/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)) diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index 16de9c4..44f7339 100644 --- a/src/cunumpy/_kernel.py +++ b/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"} diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 4ae131b..5a16a20 100644 --- a/src/cunumpy/xp.py +++ b/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"