From a9fc724395098e129005fd50faa9033e9bef2c0c Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 31 Oct 2026 13:33:49 +0100 Subject: [PATCH] Interop: copy and device semantics, explicit dtypes and a host-only Julia adapter. --- CHANGELOG.md | 26 ++ README.md | 1 + pyproject.toml | 1 + src/cunumpy/_transfers.py | 8 +- src/cunumpy/interop/__init__.py | 6 +- src/cunumpy/interop/_common.py | 374 ++++++++++++++++++++++++-- src/cunumpy/interop/dlpack.py | 18 +- src/cunumpy/interop/jax.py | 128 +++++++-- src/cunumpy/interop/julia.py | 182 +++++++++++++ src/cunumpy/interop/mlx.py | 122 +++++++-- src/cunumpy/interop/torch.py | 103 +++++-- tests/unit/test_interop.py | 214 ++++++++++++++- tests/unit/test_interop_frameworks.py | 184 +++++++++++++ tests/unit/test_interop_julia.py | 98 +++++++ 14 files changed, 1375 insertions(+), 90 deletions(-) create mode 100644 src/cunumpy/interop/julia.py create mode 100644 tests/unit/test_interop_frameworks.py create mode 100644 tests/unit/test_interop_julia.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 17305bc..78ced9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `cunumpy.interop.julia` (extra `cunumpy[julia]`, juliacall): `to_julia` / `from_julia` + between NumPy arrays and Julia `Array`s that share memory. Julia is column-major, so a + C-ordered array arrives with its dimensions reversed and an F-ordered one with the same + dimensions; `from_julia(order="C" | "F")` picks the order back. Host arrays only; CuPy to + CUDA.jl is planned. +- `copy=None | True | False` on every interop conversion (as in the array API): `False` raises + when the conversion cannot share memory, `True` always copies. +- `supports_float64(device)` in the JAX, PyTorch and MLX adapters and + `cunumpy.interop.jax.enable_x64()`. + +### Changed +- Interop conversions never change a dtype without being asked: float64 into JAX without + `jax_enable_x64`, onto an Apple GPU in PyTorch (MPS) or into MLX for the GPU raises a + `TypeError` with advice instead of becoming float32. MLX keeps float64 for + `device=mlx.core.cpu`. +- The transfer profiler tells interop conversions apart: a conversion that shares memory records + nothing, one between host and device records `to_host` / `to_device`, and one within a + memory `device_copy` (before, every conversion was a `device_copy`). +- `from_framework` copies between host and device explicitly (a recorded transfer, an error + under `copy=False`) instead of falling back to `asarray` silently, and a framework that + cannot take an array's device raises a `ValueError` naming the device. +- `from_jax` raises a clear `TypeError` for traced values (inside `jax.jit`, `grad`, `vmap`). +- `to_mlx(device=...)` names the device the array is meant for (MLX arrays have no device); + before, it passed `device` to `mlx.core.array`, which does not take it. + ## [0.7.2] - 2026-10-27 ### Fixed diff --git a/README.md b/README.md index 836fd19..574f2b4 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ never hide a NumPy name: | `xp.profiling` | `timed_region`, `nvtx_range`, `count_transfers` | | `xp.memory` | `HostStaging`, `DeviceMirror` | | `xp.petsc` | `petsc_vec` | +| `xp.interop` | `dlpack`, `jax`, `torch`, `mlx`, `julia`: explicit conversions that share memory where they can (`copy=False` to require it) and never change a dtype silently | | `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs | Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends. diff --git a/pyproject.toml b/pyproject.toml index 061995b..2234992 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ optional-dependencies.dev = [ optional-dependencies.docs = [ "griffe", "ipykernel" ] optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] optional-dependencies.jax = [ "jax" ] +optional-dependencies.julia = [ "juliacall" ] optional-dependencies.mlx = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] optional-dependencies.torch = [ "torch" ] optional-dependencies.test = [ "coverage", "pytest", "scipy" ] diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py index 1165369..d889f2b 100644 --- a/src/cunumpy/_transfers.py +++ b/src/cunumpy/_transfers.py @@ -341,12 +341,16 @@ def count_transfers( * ``to_host`` / ``to_device``: a physical copy between host and device (conversions such as :func:`~cunumpy.to_numpy` and :func:`~cunumpy.to_cupy`, mirrors, staging, serial MPI, kernel output - copy-back), with its size in bytes; + copy-back, an :mod:`~cunumpy.interop` conversion between host and + device), with its size in bytes; * ``kernel_conversion``: a :class:`~cunumpy.kernels.PyccelKernel` call that copied device arrays to the host and back, one marker per call; * ``fallback``: a :class:`~cunumpy.kernels.Kernel` without CUDA version that ran its host kernel on the CuPy backend, one marker per call; - * ``device_copy``: a device-only dtype/layout conversion; + * ``device_copy``: a copy that stays in one memory: a device-only + dtype/layout conversion, or an :mod:`~cunumpy.interop` conversion + that could not share memory (an interop conversion that shares memory + records nothing); * ``sync``: the host waited for the device (:func:`~cunumpy.synchronize`, the MPI helpers, the CUDA debug mode and, on the fake CuPy, scalar reads such as ``float(a)``). Not included in diff --git a/src/cunumpy/interop/__init__.py b/src/cunumpy/interop/__init__.py index e7deedb..58dd6e5 100644 --- a/src/cunumpy/interop/__init__.py +++ b/src/cunumpy/interop/__init__.py @@ -1,9 +1,9 @@ -"""Explicit interoperability with scientific machine-learning libraries. +"""Explicit interoperability with scientific machine-learning libraries and Julia. The adapters are optional and import their framework lazily. They do not participate in CuNumpy backend selection. """ -from cunumpy.interop import dlpack, jax, mlx, torch +from cunumpy.interop import dlpack, jax, julia, mlx, torch -__all__ = ["dlpack", "jax", "mlx", "torch"] +__all__ = ["dlpack", "jax", "julia", "mlx", "torch"] diff --git a/src/cunumpy/interop/_common.py b/src/cunumpy/interop/_common.py index b16baec..2d79afd 100644 --- a/src/cunumpy/interop/_common.py +++ b/src/cunumpy/interop/_common.py @@ -1,7 +1,17 @@ -"""Shared implementation for optional framework adapters.""" +"""Shared implementation for optional framework adapters. + +The adapters share three rules here. Memory: an array lives on the host or on +a device, as its ``__dlpack_device__`` says (Apple GPUs share the host's +memory and count as host). Copies: a conversion shares memory when it can, +``copy=False`` raises where it cannot, ``copy=True`` always copies, and the +transfer profiler records a copy as ``to_host``, ``to_device`` or +``device_copy`` (a copy that stays in one memory) and a shared conversion not +at all. Dtypes: a conversion never changes a dtype that was not asked for. +""" from __future__ import annotations +from collections.abc import Callable from importlib import import_module from typing import Any @@ -9,6 +19,21 @@ from cunumpy._transfers import _ACTIVE, _describe, _nbytes, _record +HOST = (1, 0) +"""The DLPack device of host memory, ``(kDLCPU, 0)``.""" + +# kDLCPU, kDLCUDAHost, kDLMetal (Apple's unified memory) and kDLROCMHost +_HOST_DEVICE_TYPES = {1, 3, 8, 11} +_DEVICE_NAMES = { + 2: "CUDA", + 4: "OpenCL", + 7: "Vulkan", + 10: "ROCm", + 13: "CUDA managed", + 14: "oneAPI", +} +_DLPACK_ERRORS = (AttributeError, BufferError, RuntimeError, TypeError, ValueError) + def require(module: str, *, extra: str | None = None) -> Any: """Import an optional framework and give a consistent installation error. @@ -35,23 +60,6 @@ def require(module: str, *, extra: str | None = None) -> Any: ) from error -def record(source: Any, target: Any) -> None: - """Record an adapter conversion when transfer profiling is active. - - Parameters - ---------- - source, target : object - Source and converted arrays. - """ - if _ACTIVE: - _record( - "device_copy", - f"interop {type(source).__name__} -> {type(target).__name__} " - f"({_describe(target)})", - nbytes=_nbytes(target), - ) - - def target_backend(backend: str) -> Any: """Load a NumPy or CuPy target module. @@ -94,3 +102,333 @@ def cast(array: Any, dtype: Any) -> Any: The original or cast array. """ return array if dtype is None else array.astype(dtype) + + +def check_copy(copy: bool | None) -> None: + """Reject a `copy` argument other than None, True or False. + + Parameters + ---------- + copy : bool or None + The `copy` argument of a conversion. + """ + if copy is not None and copy is not True and copy is not False: + raise ValueError(f"copy must be None, True or False, got {copy!r}") + + +def memory(array: Any) -> tuple[int, int]: + """Return the DLPack device of the memory that holds `array`. + + Parameters + ---------- + array : object + Any array; one without ``__dlpack_device__`` counts as host memory. + + Returns + ------- + tuple of int + ``(device_type, device_id)``; :data:`HOST` for host memory, also for + CUDA/ROCm pinned host memory and Apple's unified memory. + """ + get = getattr(array, "__dlpack_device__", None) + if get is None: + return HOST + try: + device_type, device_id = get() + except Exception: # noqa: BLE001 - e.g. a JAX tracer has no device + return HOST + device_type = int(device_type) + if device_type in _HOST_DEVICE_TYPES: + return HOST + return (device_type, int(device_id)) + + +def describe_memory(device: tuple[int, int]) -> str: + """Name a DLPack device for error messages. + + Parameters + ---------- + device : tuple of int + ``(device_type, device_id)`` as returned by :func:`memory`. + + Returns + ------- + str + E.g. ``"the host"`` or ``"CUDA device 0"``. + """ + if device == HOST: + return "the host" + name = _DEVICE_NAMES.get(device[0], f"DLPack device type {device[0]}") + return f"{name} device {device[1]}" + + +def pointer(array: Any) -> int | None: + """Return the address of the first element of `array`, or None if unknown. + + Parameters + ---------- + array : object + A NumPy, CuPy, JAX, PyTorch, MLX or Julia (juliacall) array. + + Returns + ------- + int or None + The address, or None when the array does not expose it. + """ + getters = ( + lambda: array.data_ptr(), # PyTorch + lambda: array.unsafe_buffer_pointer(), # JAX + lambda: array.__cuda_array_interface__["data"][0], # CuPy + lambda: array.__array_interface__["data"][0], # NumPy, juliacall + lambda: np.asarray(memoryview(array)).ctypes.data, # MLX + ) + for get in getters: + try: + return int(get()) + except Exception: # noqa: BLE001, S112 - try the next protocol + continue + return None + + +def shares_memory(source: Any, result: Any) -> bool: + """Whether a conversion returned the memory of its source. + + Parameters + ---------- + source, result : object + The array that was converted and the converted array. + + Returns + ------- + bool + True for the same object, an empty array or the same first-element + address; False otherwise, also when an address is unknown. + """ + if result is source or getattr(result, "size", None) == 0: + return True + address = pointer(source) + return address is not None and address == pointer(result) + + +def copy_error(source: Any, target: str, reason: str) -> ValueError: + """Return the error for a conversion that needs a copy under ``copy=False``. + + Parameters + ---------- + source : object + The array that was converted. + target : str + The framework or array library converted to. + reason : str + Why a copy is needed. + + Returns + ------- + ValueError + The error to raise. + """ + return ValueError( + f"converting {type(source).__name__} ({_describe(source)}) to {target} needs " + f"a copy ({reason}), but copy=False; pass copy=None to allow it" + ) + + +def device_error(source: Any, target: str, error: Exception) -> ValueError: + """Return the error for a framework that cannot take an array's device. + + Parameters + ---------- + source : object + The array that was converted. + target : str + The framework converted to. + error : Exception + The framework's own error. + + Returns + ------- + ValueError + The error to raise. + """ + return ValueError( + f"{target} cannot take a {type(source).__name__} on " + f"{describe_memory(memory(source))} ({error}); move it to a device " + f"{target} supports first, e.g. to the host with xp.to_numpy" + ) + + +def dtype_name(dtype: Any) -> str: + """Return a framework-neutral dtype name, e.g. ``"float64"``. + + Parameters + ---------- + dtype : object + A NumPy, CuPy, JAX, PyTorch or MLX dtype, or a dtype name. + + Returns + ------- + str + The name, e.g. ``"float32"`` for ``torch.float32`` and + ``mlx.core.float32``. + """ + try: + return np.dtype(dtype).name + except TypeError: + return str(dtype).rsplit(".", 1)[-1] + + +def check_dtype(source: Any, result: Any, dtype: Any, target: str, advice: str) -> None: + """Raise when a conversion changed a dtype that was not asked for. + + Parameters + ---------- + source, result : object + The array that was converted and the converted array. + dtype : object or None + The requested dtype, or None to keep the source's. + target : str + The framework converted to. + advice : str + How to keep the dtype or ask for the cast. + """ + expected = dtype if dtype is not None else getattr(source, "dtype", None) + if expected is None: + return + wanted, got = dtype_name(expected), dtype_name(result.dtype) + if wanted != got: + raise TypeError(f"{target} would turn {wanted} into {got}; {advice}") + + +def _size(array: Any) -> int | None: + """Payload bytes of an array of any framework, or None if unknown.""" + nbytes = getattr(array, "nbytes", None) + if isinstance(nbytes, int): + return nbytes + try: + return _nbytes(array) + except AttributeError: + return None + + +def record(source: Any, result: Any) -> None: + """Record an adapter copy when transfer profiling is active. + + A copy from host to device is a ``to_device`` event, from device to host a + ``to_host`` event, and one within a memory a ``device_copy`` event. + + Parameters + ---------- + source, result : object + The array that was converted and its copy. + """ + if not _ACTIVE: + return + before, after = memory(source) == HOST, memory(result) == HOST + if before == after: + kind = "device_copy" + else: + kind = "to_device" if before else "to_host" + _record( + kind, + f"interop {type(source).__name__} -> {type(result).__name__} " + f"({_describe(result)})", + nbytes=_size(result), + ) + + +def finish( + source: Any, + result: Any, + *, + copy: bool | None, + target: str, + duplicate: Callable[[Any], Any], +) -> Any: + """Apply the `copy` rule to a converted array and record the copy. + + Parameters + ---------- + source, result : object + The array that was converted and the converted array. + copy : bool or None + None: share memory when the conversion did; False: raise if it did + not; True: copy if it did. + target : str + The framework or array library converted to, for the error message. + duplicate : callable + Copies a result in its own framework. + + Returns + ------- + object + The result, copied for ``copy=True``. + """ + shared = shares_memory(source, result) + if copy is False and not shared: + raise copy_error(source, target, "the conversion could not share memory") + if copy and shared: + result = duplicate(result) + shared = False + if not shared: + record(source, result) + return result + + +def to_backend( + source: Any, + backend: str, + *, + dtype: Any, + copy: bool | None, + host_array: Callable[[Any], Any], +) -> Any: + """Convert a framework array to NumPy or CuPy (the adapters' ``from_framework``). + + DLPack is used when the source and the target share a memory; between + host and device the array is copied (``copy=False`` raises). + + Parameters + ---------- + source : object + Framework array to convert. + backend : {"numpy", "cupy"} + Array library to construct. + dtype : object or None + Dtype for the result. + copy : bool or None + See :func:`finish`. + host_array : callable + Returns `source` as a NumPy array (a view where the framework allows), + used when DLPack cannot. + + Returns + ------- + object + NumPy or CuPy array. + """ + check_copy(copy) + target = target_backend(backend) + on_host = memory(source) == HOST + if backend == "numpy": + if not on_host: + if copy is False: + raise copy_error(source, "NumPy", "device to host") + result = np.asarray(host_array(source)) + else: + try: + result = np.from_dlpack(source) + except _DLPACK_ERRORS: + result = np.asarray(host_array(source)) + elif on_host: + if copy is False: + raise copy_error(source, "CuPy", "host to device") + result = target.asarray(host_array(source)) + else: + try: + result = target.from_dlpack(source) + except _DLPACK_ERRORS as error: + raise device_error(source, "CuPy", error) from error + result = cast(result, dtype) + return finish( + source, result, copy=copy, target=backend, duplicate=lambda array: array.copy() + ) diff --git a/src/cunumpy/interop/dlpack.py b/src/cunumpy/interop/dlpack.py index 4258974..7be0660 100644 --- a/src/cunumpy/interop/dlpack.py +++ b/src/cunumpy/interop/dlpack.py @@ -4,7 +4,7 @@ from typing import Any -from cunumpy.interop._common import record, target_backend +from cunumpy.interop._common import check_copy, finish, target_backend def to_dlpack(array: Any) -> Any: @@ -28,7 +28,9 @@ def to_dlpack(array: Any) -> Any: ) from error -def from_dlpack(source: Any, *, backend: str = "numpy") -> Any: +def from_dlpack( + source: Any, *, backend: str = "numpy", copy: bool | None = None +) -> Any: """Construct a NumPy or CuPy array from a DLPack source. Parameters @@ -37,18 +39,24 @@ def from_dlpack(source: Any, *, backend: str = "numpy") -> Any: DLPack capsule or object implementing the protocol. backend : {"numpy", "cupy"}, default="numpy" Array library to construct. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Passed on to the target's ``from_dlpack``. Returns ------- object The constructed array. """ + check_copy(copy) target = target_backend(backend) if not hasattr(target, "from_dlpack"): raise TypeError(f"{backend} does not provide from_dlpack") - result = target.from_dlpack(source) - record(source, result) - return result + kwargs = {} if copy is None else {"copy": copy} + result = target.from_dlpack(source, **kwargs) + return finish( + source, result, copy=None, target=backend, duplicate=lambda array: array.copy() + ) __all__ = ["from_dlpack", "to_dlpack"] diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py index 98393f4..f48903d 100644 --- a/src/cunumpy/interop/jax.py +++ b/src/cunumpy/interop/jax.py @@ -4,14 +4,73 @@ from typing import Any -from cunumpy.interop._common import cast, record, require, target_backend +import numpy as np +from cunumpy.interop._common import ( + HOST, + cast, + check_copy, + check_dtype, + device_error, + finish, + memory, + require, + to_backend, +) -def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: +_X64_ADVICE = ( + "enable float64 with cunumpy.interop.jax.enable_x64() before creating JAX arrays, " + "or ask for the cast with dtype=..." +) + + +def enable_x64(enabled: bool = True) -> None: + """Turn JAX's 64-bit types (``jax_enable_x64``) on or off. + + Without them JAX turns float64 into float32 and int64 into int32, which + :func:`to_framework` refuses unless the cast is asked for. Call it before + creating JAX arrays. + + Parameters + ---------- + enabled : bool, default=True + Whether JAX keeps 64-bit types. + """ + jax = require("jax") + jax.config.update("jax_enable_x64", enabled) + + +def supports_float64(device: Any = None) -> bool: + """Whether JAX keeps float64 arrays on `device`. + + Parameters + ---------- + device : jax.Device or None, optional + The device to check; the default device when None. + + Returns + ------- + bool + True when ``jax_enable_x64`` is on and the device is not an Apple GPU + (jax-metal has no float64). + """ + jax = require("jax") + if not jax.config.read("jax_enable_x64"): + return False + device = jax.devices()[0] if device is None else device + return device.platform.lower() != "metal" + + +def to_framework( + array: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: """Convert a NumPy or CuPy array to a JAX array. - DLPack is used when available; `device` and `dtype` are explicit optional - conversions. + DLPack shares the memory when JAX can (a contiguous array on a device JAX + has; on the CPU, XLA also needs the buffer aligned, which small NumPy + arrays may not be); a host array it cannot share is copied. `device` and `dtype` are + explicit conversions. A float64 array without ``jax_enable_x64`` raises + instead of becoming float32 (see :func:`enable_x64`). Parameters ---------- @@ -21,6 +80,9 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: JAX device on which to place the result. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Returns ------- @@ -30,22 +92,40 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: jax = require("jax") import jax.numpy as jnp + check_copy(copy) if type(array).__module__.startswith("jax"): result = array elif hasattr(array, "__dlpack__"): - result = jax.dlpack.from_dlpack(array) + try: + result = jax.dlpack.from_dlpack(array) + except Exception as error: + if memory(array) != HOST: + raise device_error(array, "JAX", error) from error + result = jnp.asarray(array) else: result = jnp.asarray(array) result = cast(result, dtype) if device is not None: result = jax.device_put(result, device) - record(array, result) - return result - - -def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: + check_dtype(array, result, dtype, "JAX", _X64_ADVICE) + return finish( + array, + result, + copy=copy, + target="JAX", + duplicate=lambda value: jnp.array(value, copy=True), + ) + + +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: """Convert a JAX array to a NumPy or CuPy array. + DLPack shares the memory when the JAX array is in the target's memory; + between host and device the array is copied. A traced value (inside + ``jax.jit``, ``jax.grad`` or ``jax.vmap``) has no values and raises. + Parameters ---------- array : object @@ -54,23 +134,33 @@ def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Array library to construct. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Returns ------- object NumPy or CuPy array. """ - target = target_backend(backend) - if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): - result = target.from_dlpack(array) - else: - result = target.asarray(array) - result = cast(result, dtype) - record(array, result) - return result + jax = require("jax") + if isinstance(array, jax.core.Tracer): + raise TypeError( + "cannot convert a traced JAX value (inside jax.jit, jax.grad or jax.vmap): " + "it has no concrete values; convert outside the transformation, or call " + "the cunumpy code through jax.pure_callback" + ) + return to_backend(array, backend, dtype=dtype, copy=copy, host_array=np.asarray) to_jax = to_framework from_jax = from_framework -__all__ = ["from_framework", "from_jax", "to_framework", "to_jax"] +__all__ = [ + "enable_x64", + "from_framework", + "from_jax", + "supports_float64", + "to_framework", + "to_jax", +] diff --git a/src/cunumpy/interop/julia.py b/src/cunumpy/interop/julia.py new file mode 100644 index 0000000..8f54372 --- /dev/null +++ b/src/cunumpy/interop/julia.py @@ -0,0 +1,182 @@ +"""Optional interoperability with Julia arrays through juliacall. + +A NumPy array becomes a Julia ``Array`` that shares its memory, so Julia code +(e.g. a KernelAbstractions.jl kernel on the CPU backend) reads and writes the +NumPy data in place. Julia is column-major: a C-ordered array arrives with its +dimensions reversed (NumPy ``a[i, j]`` is Julia ``A[j+1, i+1]``), an +F-ordered one with the same dimensions. A cunumpy component-major array of +shape ``(3, n)`` is therefore a Julia ``n x 3`` matrix with contiguous columns. +:func:`from_framework` reverses the dimensions back (``order="C"``) or keeps +them (``order="F"``). + +Only host arrays so far; CuPy arrays to CUDA.jl ``CuArray`` through DLPack +are planned. +""" + +from __future__ import annotations + +import functools +from typing import Any + +import numpy as np + +from cunumpy.interop._common import ( + HOST, + check_copy, + copy_error, + device_error, + finish, + memory, + record, + require, + to_backend, +) + +_JULIA_TYPES = { + "bool": "Bool", + "int8": "Int8", + "int16": "Int16", + "int32": "Int32", + "int64": "Int64", + "uint8": "UInt8", + "uint16": "UInt16", + "uint32": "UInt32", + "uint64": "UInt64", + "float16": "Float16", + "float32": "Float32", + "float64": "Float64", + "complex64": "ComplexF32", + "complex128": "ComplexF64", +} + +# The Julia array keeps the NumPy array alive through its finalizer's closure. +_HELPERS = """ +module CunumpyInterop +release(owner) = nothing +function wrap(::Type{T}, pointer::Integer, dims::Tuple, owner) where {T} + array = unsafe_wrap(Array, Ptr{T}(UInt(pointer)), Tuple(Int.(dims)); own=false) + finalizer(_ -> release(owner), array) + return array +end +end +""" + + +@functools.cache +def _julia() -> Any: + """Start Julia (once) and return ``juliacall.Main`` with the helpers defined.""" + juliacall = require("juliacall", extra="julia") + juliacall.Main.seval(_HELPERS) + return juliacall.Main + + +def _is_julia(array: Any) -> bool: + return type(array).__module__.startswith("juliacall") + + +def to_framework(array: Any, *, dtype: Any = None, copy: bool | None = None) -> Any: + """Convert a NumPy array to a Julia ``Array``. + + The Julia array shares the memory of a C- or F-contiguous, writeable + array (see the module docstring for the order of the dimensions) and + keeps it alive; any other array is first copied to C order. + + Parameters + ---------- + array : object + Host array to convert. + dtype : object or None, optional + NumPy dtype for the result (a copy when it differs). + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. + + Returns + ------- + juliacall.ArrayValue + The Julia array. + """ + jl = _julia() + check_copy(copy) + if _is_julia(array): + if dtype is not None: + raise TypeError("cast Julia arrays in Julia, e.g. Float32.(A)") + return finish(array, array, copy=copy, target="Julia", duplicate=jl.copy) + if memory(array) != HOST: + raise device_error( + array, "Julia", NotImplementedError("only host arrays so far") + ) + host = np.asarray(array) + if dtype is not None: + host = host.astype(dtype, copy=False) + name = _JULIA_TYPES.get(host.dtype.name) + if name is None: + raise TypeError(f"no Julia element type for dtype {host.dtype}") + if not host.flags.writeable or not ( + host.flags.c_contiguous or host.flags.f_contiguous + ): + if copy is False: + raise copy_error(array, "Julia", "not contiguous or read-only") + host = np.array(host, order="C") + dims = host.shape[::-1] if host.flags.c_contiguous else host.shape + result = jl.CunumpyInterop.wrap(jl.seval(name), host.ctypes.data, dims, host) + return finish(array, result, copy=copy, target="Julia", duplicate=jl.copy) + + +def from_framework( + array: Any, + *, + backend: str = "numpy", + dtype: Any = None, + copy: bool | None = None, + order: str = "C", +) -> Any: + """Convert a Julia array to a NumPy or CuPy array. + + A NumPy result is a view of a Julia ``Array`` and keeps it alive; other + Julia arrays (views, ranges) are first collected into an ``Array``. + + Parameters + ---------- + array : juliacall.ArrayValue + Julia array to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. + order : {"C", "F"}, default="C" + "C": reverse the dimensions, giving a C-ordered array (the inverse of + :func:`to_framework` for a C-ordered array); "F": keep them, giving + an F-ordered array. + + Returns + ------- + object + NumPy or CuPy array. + """ + jl = _julia() + if order not in ("C", "F"): + raise ValueError(f"order must be 'C' or 'F', got {order!r}") + if not _is_julia(array) or not jl.isa(array, jl.AbstractArray): + raise TypeError(f"expected a Julia array, got {type(array).__name__}") + source = array + if not jl.isa(array, jl.Array): + if copy is False: + raise copy_error(array, backend, "not a Julia Array") + source = jl.collect(array) + record(array, source) + + def host_array(value: Any) -> np.ndarray: + view = np.asarray(value) + return view.T if order == "C" else view + + return to_backend(source, backend, dtype=dtype, copy=copy, host_array=host_array) + + +to_julia = to_framework +from_julia = from_framework + +__all__ = ["from_framework", "from_julia", "to_framework", "to_julia"] diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py index 976978c..4efef0a 100644 --- a/src/cunumpy/interop/mlx.py +++ b/src/cunumpy/interop/mlx.py @@ -4,43 +4,107 @@ from typing import Any -from cunumpy.interop._common import cast, record, require, target_backend +import numpy as np +from cunumpy.interop._common import ( + HOST, + check_copy, + check_dtype, + copy_error, + device_error, + dtype_name, + finish, + memory, + require, + to_backend, +) -def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: - """Convert a NumPy or CuPy array to an MLX array. +_GPU_ADVICE = ( + "MLX has float64 on the CPU only: pass device=mlx.core.cpu to keep float64, " + "or ask for the cast with dtype=mlx.core.float32" +) + + +def supports_float64(device: Any = None) -> bool: + """Whether MLX can compute with float64 arrays on `device`. + + Parameters + ---------- + device : mlx.core.Device or None, optional + The device to check; MLX's default device when None. + + Returns + ------- + bool + True on the CPU, False on the GPU. + """ + mx = require("mlx.core", extra="mlx") + device = mx.default_device() if device is None else device + return device == mx.cpu + + +def to_framework( + array: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert a NumPy array to an MLX array. + + MLX always copies into its own memory, so ``copy=False`` raises. MLX + arrays have no device; `device` is the device they will be used on, which + decides whether float64 is kept (CPU) or raises (GPU) instead of + becoming float32, as ``mlx.core.array`` does. Parameters ---------- array : object - Array to convert. + Host array to convert. device : object or None, optional - MLX device on which to place the result. + MLX device the array is meant for; MLX's default device when None. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None or True: copy (MLX cannot share); False: raise. Returns ------- object MLX array. """ - mlx = require("mlx.core", extra="mlx") + mx = require("mlx.core", extra="mlx") + check_copy(copy) if type(array).__module__.startswith("mlx"): - result = array - elif hasattr(mlx, "from_dlpack") and hasattr(array, "__dlpack__"): - result = mlx.from_dlpack(array) + result = array if dtype is None else array.astype(dtype) else: - result = mlx.array(array) - result = cast(result, dtype) - if device is not None: - result = mlx.array(result, device=device) - record(array, result) - return result + if memory(array) != HOST: + raise device_error(array, "MLX", TypeError("MLX takes host arrays")) + if copy is False: + raise copy_error(array, "MLX", "MLX copies into its own memory") + source_dtype = getattr(array, "dtype", None) + if ( + dtype is None + and source_dtype is not None + and dtype_name(source_dtype) == "float64" + and supports_float64(device) + ): + dtype = mx.float64 + result = mx.array(array, dtype=dtype) + check_dtype(array, result, dtype, "MLX", _GPU_ADVICE) + return finish( + array, + result, + copy=copy, + target="MLX", + duplicate=lambda value: mx.array(np.array(value)), + ) -def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: """Convert an MLX array to a NumPy or CuPy array. + A NumPy result shares the MLX array's memory (Apple's unified memory); a + CuPy result is a copy. + Parameters ---------- array : object @@ -49,23 +113,31 @@ def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Array library to construct. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Returns ------- object NumPy or CuPy array. """ - target = target_backend(backend) - if hasattr(array, "__dlpack__") and hasattr(target, "from_dlpack"): - result = target.from_dlpack(array) - else: - result = target.asarray(array) - result = cast(result, dtype) - record(array, result) - return result + return to_backend( + array, + backend, + dtype=dtype, + copy=copy, + host_array=lambda value: np.asarray(memoryview(value)), + ) to_mlx = to_framework from_mlx = from_framework -__all__ = ["from_framework", "from_mlx", "to_framework", "to_mlx"] +__all__ = [ + "from_framework", + "from_mlx", + "supports_float64", + "to_framework", + "to_mlx", +] diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py index 6b77123..51a3113 100644 --- a/src/cunumpy/interop/torch.py +++ b/src/cunumpy/interop/torch.py @@ -4,12 +4,48 @@ from typing import Any -from cunumpy.interop._common import cast, record, require, target_backend +from cunumpy.interop._common import ( + HOST, + check_copy, + check_dtype, + device_error, + finish, + memory, + require, + to_backend, +) +_MPS_ADVICE = ( + "Apple GPUs (MPS) have no float64; ask for the cast with dtype=torch.float32" +) -def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + +def supports_float64(device: Any = None) -> bool: + """Whether PyTorch keeps float64 tensors on `device`. + + Parameters + ---------- + device : torch.device, str or None, optional + The device to check; the CPU when None. + + Returns + ------- + bool + False on an Apple GPU (``"mps"``), True otherwise. + """ + torch = require("torch") + return torch.device("cpu" if device is None else device).type != "mps" + + +def to_framework( + array: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: """Convert a NumPy or CuPy array to a PyTorch tensor. + DLPack shares the memory when PyTorch can; a host array it cannot share is + copied. `device` and `dtype` are explicit conversions. A float64 array + for an Apple GPU raises instead of becoming float32. + Parameters ---------- array : object @@ -18,6 +54,9 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: PyTorch device on which to place the result. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Returns ------- @@ -25,23 +64,43 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: PyTorch tensor. """ torch = require("torch") - if type(array).__module__.startswith("torch"): + check_copy(copy) + if isinstance(array, torch.Tensor): result = array - elif hasattr(torch, "from_dlpack") and hasattr(array, "__dlpack__"): - result = torch.from_dlpack(array) + elif hasattr(array, "__dlpack__"): + try: + result = torch.from_dlpack(array) + except Exception as error: + if memory(array) != HOST: + raise device_error(array, "PyTorch", error) from error + result = torch.as_tensor(array) else: result = torch.as_tensor(array) if dtype is not None: result = result.to(dtype=dtype) if device is not None: + if result.dtype == torch.float64 and not supports_float64(device): + raise TypeError(f"PyTorch would turn float64 into float32; {_MPS_ADVICE}") result = result.to(device=device) - record(array, result) - return result + check_dtype(array, result, dtype, "PyTorch", _MPS_ADVICE) + return finish( + array, + result, + copy=copy, + target="PyTorch", + duplicate=lambda value: value.clone(), + ) -def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: """Convert a PyTorch tensor to a NumPy or CuPy array. + DLPack shares the memory when the tensor is in the target's memory; + between host and device (and from an Apple GPU) the tensor is copied. The + result is detached from autograd. + Parameters ---------- array : object @@ -50,23 +109,33 @@ def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Array library to construct. dtype : object or None, optional Dtype for the result. + copy : bool or None, optional + None: share memory when possible; False: raise when a copy is + needed; True: always copy. Returns ------- object NumPy or CuPy array. """ - target = target_backend(backend) - if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): - result = target.from_dlpack(array) - else: - result = target.asarray(array.detach().cpu().numpy()) - result = cast(result, dtype) - record(array, result) - return result + detach = getattr(array, "detach", None) + source = array if detach is None else detach() + return to_backend( + source, + backend, + dtype=dtype, + copy=copy, + host_array=lambda tensor: tensor.cpu().numpy(), + ) to_torch = to_framework from_torch = from_framework -__all__ = ["from_framework", "from_torch", "to_framework", "to_torch"] +__all__ = [ + "from_framework", + "from_torch", + "supports_float64", + "to_framework", + "to_torch", +] diff --git a/tests/unit/test_interop.py b/tests/unit/test_interop.py index 22442ff..152a5e0 100644 --- a/tests/unit/test_interop.py +++ b/tests/unit/test_interop.py @@ -1,12 +1,14 @@ """Tests for optional framework interoperability.""" +import sys import types import numpy as np import pytest -from cunumpy.interop import dlpack +from cunumpy.interop import _common, dlpack from cunumpy.interop._common import require, target_backend +from cunumpy.profiling import count_transfers def test_dlpack_helpers_use_array_protocol(monkeypatch): @@ -37,3 +39,213 @@ def test_optional_dependency_error_is_actionable(): def test_target_backend_rejects_unknown_backend(): with pytest.raises(ValueError, match="backend must be"): target_backend("jax") + + +class _Device: + """An array on a fake CUDA device, with a DLPack device and a pointer.""" + + shape, dtype, size = (3,), np.dtype("float64"), 3 + + def __init__(self, ptr=1000, device=(2, 0)): + self.ptr, self.device = ptr, device + + def __dlpack_device__(self): + return self.device + + @property + def __cuda_array_interface__(self): + return {"data": (self.ptr, False)} + + +def test_memory_reads_the_dlpack_device(): + assert _common.memory(np.ones(2)) == _common.HOST + assert _common.memory([1.0]) == _common.HOST + assert _common.memory(_Device()) == (2, 0) + assert _common.memory(_Device(device=(8, 0))) == _common.HOST # Apple GPU + + class Broken: + def __dlpack_device__(self): + raise RuntimeError("no device") + + assert _common.memory(Broken()) == _common.HOST + + +def test_describe_memory(): + assert _common.describe_memory(_common.HOST) == "the host" + assert _common.describe_memory((2, 1)) == "CUDA device 1" + assert _common.describe_memory((99, 0)) == "DLPack device type 99 device 0" + + +def test_pointer_and_shared_memory(): + a = np.arange(4.0) + assert _common.pointer(a) == a.ctypes.data + assert _common.pointer(_Device(ptr=7)) == 7 + assert _common.pointer(memoryview(bytearray(4))) is not None + assert _common.pointer(object()) is None + assert _common.shares_memory(a, a) + assert _common.shares_memory(a, a[:]) + assert not _common.shares_memory(a, a.copy()) + assert _common.shares_memory(a, np.ones(0)) + assert not _common.shares_memory(object(), a) + + +def test_check_copy_rejects_other_values(): + for copy in (None, True, False): + _common.check_copy(copy) + with pytest.raises(ValueError, match="copy must be"): + _common.check_copy("never") + + +def test_dtype_name_is_framework_neutral(): + assert _common.dtype_name(np.float32) == "float32" + assert _common.dtype_name("float64") == "float64" + + class TorchLike: + def __str__(self): + return "torch.float64" + + assert _common.dtype_name(TorchLike()) == "float64" + + +def test_check_dtype_raises_on_an_unrequested_cast(): + a = np.ones(2) + _common.check_dtype(a, a, None, "X", "advice") + _common.check_dtype(a, a.astype(np.float32), np.float32, "X", "advice") + _common.check_dtype([1.0], a, None, "X", "advice") # no source dtype + with pytest.raises(TypeError, match="X would turn float64 into float32; ask"): + _common.check_dtype(a, a.astype(np.float32), None, "X", "ask") + + +def test_finish_applies_the_copy_rule_and_records_copies(): + a = np.arange(3.0) + with count_transfers() as counter: + assert _common.finish(a, a, copy=None, target="X", duplicate=np.copy) is a + assert counter.total == 0 + copied = _common.finish(a, a, copy=True, target="X", duplicate=np.copy) + assert copied is not a and counter.count("device_copy") == 1 + _common.finish(a, a.copy(), copy=None, target="X", duplicate=np.copy) + assert counter.count("device_copy") == 2 + with pytest.raises(ValueError, match="to X needs a copy"): + _common.finish(a, a.copy(), copy=False, target="X", duplicate=np.copy) + + +def test_record_tells_host_and_device_copies_apart(): + host, device = np.ones(3), _Device() + with count_transfers() as counter: + _common.record(host, device) + _common.record(device, host) + _common.record(device, _Device(ptr=2)) + assert [event.kind for event in counter.events] == [ + "to_device", + "to_host", + "device_copy", + ] + assert counter.bytes("to_device") == 24 + _common.record(host, device) # no counter active: nothing to do + + +def test_to_backend_numpy_shares_host_memory(): + a = np.arange(3.0) + with count_transfers() as counter: + result = _common.to_backend( + a, "numpy", dtype=None, copy=False, host_array=np.asarray + ) + assert np.shares_memory(result, a) and counter.total == 0 + cast = _common.to_backend( + a, "numpy", dtype=np.float32, copy=None, host_array=np.asarray + ) + assert cast.dtype == np.float32 + with pytest.raises(ValueError, match="needs a copy"): + _common.to_backend( + a, "numpy", dtype=np.float32, copy=False, host_array=np.asarray + ) + + +def test_to_backend_numpy_falls_back_to_the_host_array(): + a = np.arange(3.0) + result = _common.to_backend( + [0.0, 1.0, 2.0], "numpy", dtype=None, copy=None, host_array=lambda _: a + ) + assert result is a + + +def test_to_backend_numpy_copies_a_device_array(): + device = _Device() + with count_transfers() as counter: + result = _common.to_backend( + device, "numpy", dtype=None, copy=None, host_array=lambda _: np.ones(3) + ) + np.testing.assert_array_equal(result, np.ones(3)) + assert counter.count("to_host") == 1 + with pytest.raises(ValueError, match="device to host"): + _common.to_backend( + device, "numpy", dtype=None, copy=False, host_array=lambda _: np.ones(3) + ) + + +def test_to_backend_cupy_paths(monkeypatch): + received = [] + + class Target: + @staticmethod + def asarray(array): + received.append(("asarray", array)) + return _Device(ptr=5) + + @staticmethod + def from_dlpack(source): + if source.device != (2, 0): + raise BufferError("wrong device") + received.append(("from_dlpack", source)) + return source + + monkeypatch.setattr(_common, "target_backend", lambda backend: Target) + a = np.ones(3) + with count_transfers() as counter: + _common.to_backend(a, "cupy", dtype=None, copy=None, host_array=np.asarray) + assert received[-1][0] == "asarray" and counter.count("to_device") == 1 + with pytest.raises(ValueError, match="host to device"): + _common.to_backend(a, "cupy", dtype=None, copy=False, host_array=np.asarray) + + device = _Device() + result = _common.to_backend(device, "cupy", dtype=None, copy=False, host_array=None) + assert result is device + with pytest.raises(ValueError, match="CuPy cannot take a _Device on ROCm device 0"): + _common.to_backend( + _Device(device=(10, 0)), "cupy", dtype=None, copy=None, host_array=None + ) + + +def test_from_dlpack_passes_copy_on(): + calls = [] + target = types.SimpleNamespace( + from_dlpack=lambda source, **kwargs: calls.append(kwargs) or source + ) + a = np.ones(2) + with pytest.MonkeyPatch.context() as patch: + patch.setattr("cunumpy.interop.dlpack.target_backend", lambda backend: target) + assert dlpack.from_dlpack(a) is a + dlpack.from_dlpack(a, copy=False) + assert calls == [{}, {"copy": False}] + with pytest.raises(ValueError, match="copy must be"): + dlpack.from_dlpack(a, copy=1) + assert dlpack.from_dlpack(a, copy=False).ctypes.data == a.ctypes.data + + +def test_from_dlpack_rejects_a_target_without_dlpack(monkeypatch): + monkeypatch.setattr( + "cunumpy.interop.dlpack.target_backend", lambda backend: types.SimpleNamespace() + ) + with pytest.raises(TypeError, match="does not provide from_dlpack"): + dlpack.from_dlpack(np.ones(2)) + + +def test_target_backend_cupy_without_cupy(monkeypatch): + monkeypatch.setitem(sys.modules, "cupy", None) + with pytest.raises(ImportError, match="CuPy interoperability needs CuPy"): + target_backend("cupy") + + +def test_device_error_names_the_device(): + error = _common.device_error(_Device(), "JAX", RuntimeError("boom")) + assert "JAX cannot take a _Device on CUDA device 0 (boom)" in str(error) diff --git a/tests/unit/test_interop_frameworks.py b/tests/unit/test_interop_frameworks.py new file mode 100644 index 0000000..8b24af7 --- /dev/null +++ b/tests/unit/test_interop_frameworks.py @@ -0,0 +1,184 @@ +"""Tests for the JAX, PyTorch and MLX adapters with the real frameworks (skipped without them).""" + +import numpy as np +import pytest + +from cunumpy.profiling import assert_no_transfers, count_transfers + + +def _aligned(n, alignment=64): + """A float64 array whose buffer is aligned for XLA, which shares only aligned memory.""" + buffer = np.empty(n * 8 + alignment, dtype=np.uint8) + start = -buffer.ctypes.data % alignment + return buffer[start : start + n * 8].view(np.float64) + + +@pytest.fixture +def jax_x64(): + jax = pytest.importorskip("jax") + from cunumpy.interop import jax as adapter + + enabled = jax.config.read("jax_enable_x64") + yield jax, adapter + jax.config.update("jax_enable_x64", enabled) + + +def test_jax_refuses_float64_without_x64(jax_x64): + _, adapter = jax_x64 + adapter.enable_x64(False) + assert not adapter.supports_float64() + with pytest.raises(TypeError, match="JAX would turn float64 into float32; enable"): + adapter.to_jax(np.arange(3.0)) + assert adapter.to_jax(np.arange(3.0), dtype="float32").dtype == np.float32 + + +def test_jax_round_trip_shares_memory(jax_x64): + jax, adapter = jax_x64 + adapter.enable_x64() + assert adapter.supports_float64() + assert adapter.supports_float64(jax.devices()[0]) + a = _aligned(4) + a[:] = np.arange(4.0) + with assert_no_transfers(): + value = adapter.to_jax(a, copy=False) + back = adapter.from_jax(value, copy=False) + assert value.dtype == np.float64 + assert back.ctypes.data == a.ctypes.data + assert adapter.to_jax(value) is value + + +def test_jax_copies_and_reports_them(jax_x64): + jax, adapter = jax_x64 + adapter.enable_x64() + strided = np.arange(6.0).reshape(2, 3)[:, ::2] + with count_transfers() as counter: + value = adapter.to_jax(strided) # JAX's DLPack takes no strides: a copy + adapter.to_jax(np.arange(3.0), copy=True) + adapter.from_jax(value, copy=True) + adapter.to_jax([1.0, 2.0]) + assert counter.count("device_copy") == 4 + np.testing.assert_array_equal(np.asarray(value), strided) + with pytest.raises(ValueError, match="to JAX needs a copy"): + adapter.to_jax(strided, copy=False) + with pytest.raises(ValueError, match="to JAX needs a copy"): + adapter.to_jax(np.arange(3.0), dtype="float32", copy=False) + placed = adapter.to_jax(np.arange(3.0), device=jax.devices()[0]) + assert placed.dtype == np.float64 + + +def test_jax_rejects_tracers(jax_x64): + jax, adapter = jax_x64 + + @jax.jit + def convert(value): + adapter.from_jax(value) + return value + + with pytest.raises(TypeError, match="traced JAX value"): + convert(jax.numpy.ones(2)) + + +def test_jax_rejects_a_device_it_does_not_have(jax_x64): + _, adapter = jax_x64 + + class OnCuda: + dtype = np.dtype("float64") + + def __dlpack__(self, **kwargs): + raise RuntimeError("not a capsule") + + def __dlpack_device__(self): + return (2, 0) + + with pytest.raises(ValueError, match="JAX cannot take a OnCuda on CUDA device 0"): + adapter.to_jax(OnCuda()) + + +def test_torch_round_trip_shares_memory(): + torch = pytest.importorskip("torch") + from cunumpy.interop import torch as adapter + + a = np.arange(4.0) + with assert_no_transfers(): + tensor = adapter.to_torch(a, copy=False) + back = adapter.from_torch(tensor, copy=False) + assert tensor.dtype == torch.float64 and back.ctypes.data == a.ctypes.data + assert adapter.to_torch(tensor) is tensor + grad = torch.ones(2, requires_grad=True) + np.testing.assert_array_equal(adapter.from_torch(grad), [1.0, 1.0]) + + +def test_torch_copies_and_dtypes(): + torch = pytest.importorskip("torch") + from cunumpy.interop import torch as adapter + + with count_transfers() as counter: + adapter.to_torch(np.arange(3.0), copy=True) + adapter.to_torch([1.0, 2.0]) + adapter.to_torch(np.arange(3.0), dtype=torch.float32) + assert counter.count("device_copy") == 3 + with pytest.raises(ValueError, match="to PyTorch needs a copy"): + adapter.to_torch(np.arange(3.0), dtype=torch.float32, copy=False) + assert adapter.supports_float64() and adapter.supports_float64("cpu") + assert not adapter.supports_float64("mps") + with pytest.raises(TypeError, match="PyTorch would turn float64 into float32"): + adapter.to_torch(np.arange(3.0), device="mps") + assert adapter.to_torch(np.arange(3.0), device="cpu").dtype == torch.float64 + + +def test_torch_on_an_apple_gpu(): + torch = pytest.importorskip("torch") + if not torch.backends.mps.is_available(): + pytest.skip("needs an Apple GPU") + from cunumpy.interop import torch as adapter + + tensor = adapter.to_torch(np.ones(3, np.float32), device="mps") + with count_transfers() as counter: + back = adapter.from_torch(tensor) + np.testing.assert_array_equal(back, np.ones(3)) + assert counter.count("device_copy") == 1 # unified memory: no host transfer + + +def test_mlx_float64_and_copies(): + mx = pytest.importorskip("mlx.core") + from cunumpy.interop import mlx as adapter + + a = np.arange(3.0) + with pytest.raises(TypeError, match="MLX would turn float64 into float32"): + adapter.to_mlx(a, device=mx.gpu) + assert adapter.to_mlx(a, device=mx.cpu).dtype == mx.float64 + assert adapter.to_mlx(a, dtype=mx.float32).dtype == mx.float32 + assert adapter.supports_float64(mx.cpu) and not adapter.supports_float64(mx.gpu) + assert adapter.supports_float64() == (mx.default_device() == mx.cpu) + with pytest.raises(ValueError, match="MLX copies into its own memory"): + adapter.to_mlx(np.ones(3, np.float32), copy=False) + with count_transfers() as counter: + value = adapter.to_mlx(np.ones(3, np.float32)) + same = adapter.to_mlx(value) + copied = adapter.to_mlx(value, copy=True) + assert same is value and copied is not value + assert counter.count("device_copy") == 2 + assert adapter.to_mlx(value, dtype=mx.float16).dtype == mx.float16 + + +def test_mlx_to_numpy_shares_memory(): + mx = pytest.importorskip("mlx.core") + from cunumpy.interop import mlx as adapter + + value = mx.array(np.ones(3, np.float32)) + with assert_no_transfers(): + view = adapter.from_mlx(value, copy=False) + view[0] = 5.0 + assert value[0].item() == 5.0 + + +def test_mlx_rejects_device_arrays(): + pytest.importorskip("mlx.core") + from cunumpy.interop import mlx as adapter + + class OnCuda: + def __dlpack_device__(self): + return (2, 0) + + with pytest.raises(ValueError, match="MLX cannot take a OnCuda"): + adapter.to_mlx(OnCuda()) diff --git a/tests/unit/test_interop_julia.py b/tests/unit/test_interop_julia.py new file mode 100644 index 0000000..52226f5 --- /dev/null +++ b/tests/unit/test_interop_julia.py @@ -0,0 +1,98 @@ +"""Tests for the Julia adapter (skipped without juliacall).""" + +import gc + +import numpy as np +import pytest + +from cunumpy.profiling import assert_no_transfers, count_transfers + +juliacall = pytest.importorskip("juliacall") + +from cunumpy.interop import julia + +jl = juliacall.Main + + +def test_c_order_arrives_with_reversed_dimensions(): + a = np.arange(6.0).reshape(2, 3) + with assert_no_transfers(): + array = julia.to_julia(a, copy=False) + back = julia.from_julia(array, copy=False) + assert jl.isa(array, jl.Matrix) and tuple(array.shape) == (3, 2) + jl.seval("A -> (A[3, 1] = -1.0; nothing)")(array) + assert a[0, 2] == -1.0 + assert back.shape == (2, 3) and back.flags.c_contiguous + assert np.shares_memory(back, a) + + +def test_f_order_keeps_its_dimensions(): + a = np.asfortranarray(np.arange(6.0).reshape(2, 3)) + array = julia.to_julia(a, copy=False) + assert tuple(array.shape) == (2, 3) + back = julia.from_julia(array, order="F") + assert back.flags.f_contiguous and np.shares_memory(back, a) + with pytest.raises(ValueError, match="order must be"): + julia.from_julia(array, order="K") + + +def test_strided_and_read_only_arrays_are_copied(): + strided = np.arange(6.0).reshape(2, 3)[:, ::2] + read_only = np.arange(3.0) + read_only.flags.writeable = False + with count_transfers() as counter: + array = julia.to_julia(strided) + julia.to_julia(read_only) + julia.to_julia(np.arange(3.0), copy=True) + assert counter.count("device_copy") == 3 + np.testing.assert_array_equal(julia.from_julia(array), strided) + with pytest.raises(ValueError, match="not contiguous or read-only"): + julia.to_julia(strided, copy=False) + + +def test_dtypes_and_scalars(): + array = julia.to_julia(np.arange(3, dtype=np.int32), dtype=np.float32) + assert jl.eltype(array) == jl.Float32 + assert jl.eltype(julia.to_julia(np.ones(2, dtype=bool))) == jl.Bool + assert julia.from_julia(julia.to_julia(np.float64(3.0))) == 3.0 + assert tuple(julia.to_julia(np.zeros((0, 3))).shape) == (3, 0) + assert julia.from_julia(array, dtype=np.float64).dtype == np.float64 + with pytest.raises(TypeError, match="no Julia element type"): + julia.to_julia(np.array(["a"])) + + +def test_julia_arrays_pass_through(): + array = jl.seval("zeros(2, 2)") + assert julia.to_julia(array) is array + assert julia.to_julia(array, copy=True) is not array + with pytest.raises(TypeError, match="cast Julia arrays in Julia"): + julia.to_julia(array, dtype=np.float32) + + +def test_other_julia_arrays_are_collected(): + with count_transfers() as counter: + result = julia.from_julia(jl.seval("reshape(1:6, 2, 3)")) + np.testing.assert_array_equal(result, [[1, 2], [3, 4], [5, 6]]) + assert counter.count("device_copy") == 1 + with pytest.raises(ValueError, match="not a Julia Array"): + julia.from_julia(jl.seval("view(zeros(3, 4), :, 1:2)"), copy=False) + with pytest.raises(TypeError, match="expected a Julia array"): + julia.from_julia(np.ones(2)) + + +def test_the_julia_array_keeps_the_numpy_array_alive(): + a = np.ones(4) + array = julia.to_julia(a) + del a + gc.collect() + jl.GC.gc() + np.testing.assert_array_equal(julia.from_julia(array), np.ones(4)) + + +def test_device_arrays_are_not_supported_yet(): + class OnCuda: + def __dlpack_device__(self): + return (2, 0) + + with pytest.raises(ValueError, match="Julia cannot take a OnCuda on CUDA device 0"): + julia.to_julia(OnCuda())