From 32586b393e85714ea0d2c3ff47caff841c950885 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 6 Oct 2026 14:43:58 +0200 Subject: [PATCH 1/8] Add C-contiguous array views CArray1D-4D (#69) --- CHANGELOG.md | 5 + docs/source/kernels/arguments.md | 34 +++++ src/cunumpy/_cuda_kernel.py | 110 ++++++++++++---- src/cunumpy/_emulation.py | 15 ++- .../cuda/include/cunumpy/array_view.cuh | 116 ++++++++++++++++- tests/unit/test_cuda_kernel.py | 123 ++++++++++++++++++ tests/unit/test_emulation.py | 31 +++++ 7 files changed, 400 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 94dd615..7d0caef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 The former function names are removed without compatibility aliases. ### Added +- C-contiguous array views `CArray1D` to `CArray4D` in + `cunumpy/array_view.cuh`. They hold a pointer and shape only, so `a(i, j)` is + `data[i * shape[1] + j]`. As kernel parameters or struct fields, they reject + non-contiguous arrays and never copy them. `CudaStruct.from_signature` and + `from_pyccel_class` take `contiguous=True` or field names to generate them. - Device implementation selection via `kernels.set_device_kernel_implementation`, `get_device_kernel_implementation`, `use_device_kernel_implementation`, and `CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION`. Accept `"cuda"` or automatic selection diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md index 5e4e80f..87cf523 100644 --- a/docs/source/kernels/arguments.md +++ b/docs/source/kernels/arguments.md @@ -325,6 +325,40 @@ extern "C" __global__ void push(MarkerArgs m, double dt) { } ``` +### Contiguous views: `CArray2D` + +`Array2D` takes any view, so its strides are only known at run time and a +kernel cannot tell which index is the fast one. When an array is always +C-contiguous (a marker array, a grid), declare it as `CArray1D` to +`CArray4D` instead. The view holds a pointer and the shape, no strides, and +`m.markers(ip, 0)` is `data[ip * shape[1] + 0]`: the last index is always the +fast one, as in the row-major memory the host code uses. + +```python +MarkerArgs = xp.cuda.CudaStruct.from_pyccel_class( + "my_sim/kernel_arguments/pusher_args_kernels.py", + "MarkerArguments", + "MarkerArgs", + contiguous=["markers"], # or True for every array field +) +``` + +```c +struct MarkerArgs { + CArray2D markers; + long long n_markers; + Array1D valid; +}; +``` + +* A non-contiguous array (e.g. `markers[:, 0:3]`) raises `TypeError` at + packing or launch. It is **not** copied with `ascontiguousarray`, since + the kernel would write into the copy and the result would be lost. + Use `Array2D` for arguments that are sometimes views. +* A `CArrayND` converts to an `ArrayND`, so device helper functions + written for strided views also take it. +* `contiguous=` works the same in `CudaStruct.from_signature`. + ### Write the struct to a header Kernels in `.cu` files include the struct from a header. Generate it from the diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index c3f931f..2d0eb56 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -21,7 +21,9 @@ (from the shipped header ``cunumpy/array_view.cuh``, see :func:`cuda_include_dir`) takes a 2D CuPy array, contiguous or not, and receives its pointer, shape and strides, so that kernels index ``a(i, j)`` - like the pyccel kernels they are ported from; + like the pyccel kernels they are ported from; ``CArray2D`` is the + C-contiguous variant (shape only, ``a(i, j)`` is ``data[i * shape[1] + j]``), + which takes C-contiguous arrays only and raises for other views; * C++ function templates are instantiated with ``template_args``, and generated kernels (one source per variant) are compiled once per variant by :class:`CudaKernelVariants`. @@ -187,7 +189,11 @@ class CudaParameter(NamedTuple): The struct type, for a struct passed by value. view_ndim : int | None The number of dimensions, for an array view (``Array1D`` to - ``Array4D``, see :func:`cuda_include_dir`) passed by value. + ``Array4D`` or ``CArray1D`` to ``CArray4D``, see + :func:`cuda_include_dir`) passed by value. + contiguous : bool + Whether the array view is C-contiguous (``CArray2D``): packed + without strides, and only C-contiguous arrays are accepted. """ name: str @@ -196,6 +202,7 @@ class CudaParameter(NamedTuple): pointer: bool struct: CudaStruct | None = None view_ndim: int | None = None + contiguous: bool = False # C types (after removing qualifiers) and their NumPy dtypes. ``long`` is 64 bit, @@ -259,30 +266,30 @@ class CudaParameter(NamedTuple): _QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"} _COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>") -# Array1D to Array4D (cunumpy/array_view.cuh), T a scalar type of _CTYPES -_VIEW = re.compile(r"\bArray([1234])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>") +# Array1D to Array4D and CArray1D to CArray4D (cunumpy/array_view.cuh), +# T a scalar type of _CTYPES +_VIEW = re.compile(r"\b(C?)Array([1234])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>") _TOKEN = re.compile( - r"Array[1234]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>" + r"C?Array[1234]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>" r"|[A-Za-z_]\w*|\*|\[\s*\]", ) def _normalize_view(match: re.Match) -> str: """``Array2D< const double >`` -> ``Array2D``.""" - words = [w for w in match.group(2).split() if w not in _QUALIFIERS] - return f"Array{match.group(1)}D<{' '.join(words)}>" - - -def _view_dtype(ndim: int) -> np.dtype: - """The structured dtype with the C layout of ``ArrayD``.""" - return np.dtype( - [ - ("data", np.uint64), - ("shape", np.int64, (ndim,)), - ("strides", np.int64, (ndim,)), - ], - align=True, - ) + words = [w for w in match.group(3).split() if w not in _QUALIFIERS] + return f"{match.group(1)}Array{match.group(2)}D<{' '.join(words)}>" + + +def _view_dtype(ndim: int, contiguous: bool = False) -> np.dtype: + """The structured dtype with the C layout of ``ArrayD``. + + ``CArrayD`` (`contiguous`) has no strides. + """ + fields = [("data", np.uint64), ("shape", np.int64, (ndim,))] + if not contiguous: + fields.append(("strides", np.int64, (ndim,))) + return np.dtype(fields, align=True) def ctype_of(dtype: Any) -> str: @@ -515,14 +522,15 @@ def _parse_parameter( return CudaParameter(name, ctype, struct.dtype, False, struct) view = _VIEW.fullmatch(ctype) if view is not None: - element = view.group(2) + element = view.group(3) if pointers or element not in _CTYPES: raise ValueError( f"cannot check the kernel parameter {text.strip()!r}: array views " f"take a scalar element type and are passed by value", ) - ndim = int(view.group(1)) - return CudaParameter(name, ctype, np.dtype(_CTYPES[element]), False, None, ndim) + ndim = int(view.group(2)) + dtype = np.dtype(_CTYPES[element]) + return CudaParameter(name, ctype, dtype, False, None, ndim, bool(view.group(1))) if ctype == "void" and pointers == 1: return CudaParameter(name, ctype, None, True) if pointers > 1 or ctype not in _CTYPES: @@ -714,10 +722,13 @@ def _view_checker(param: CudaParameter, index: int) -> Callable[[Any], Any]: """Checker for an array view: packs (pointer, shape, strides) of a device array. Strides are converted from bytes to elements; the array need not be - contiguous. + contiguous. A contiguous view (``CArray2D``) packs (pointer, shape) and + takes C-contiguous arrays only: it is never copied, since what the kernel + writes into a copy would be lost. """ ndim = param.view_ndim - dtype = _view_dtype(ndim) + contiguous = param.contiguous + dtype = _view_dtype(ndim, contiguous) def check(value: Any) -> Any: _check_device_array(param, index, value) @@ -725,6 +736,19 @@ def check(value: Any) -> Any: raise TypeError( f"{_describe(param, index)} must be a {ndim}D array, got {value.ndim}D", ) + if contiguous: + flags = getattr(value, "flags", None) + if flags is not None and not flags.c_contiguous: + raise TypeError( + f"{_describe(param, index)} must be C-contiguous: got a view " + f"with strides {tuple(value.strides)}; declare the parameter " + f"as Array{ndim}D to take strided views, or pass " + f"cupy.ascontiguousarray(...) and copy the result back", + ) + packed = np.zeros((), dtype=dtype) + packed["data"] = value.data.ptr + packed["shape"] = value.shape + return packed[()] itemsize = value.dtype.itemsize strides = [s // itemsize for s in value.strides] if any(s * itemsize != stride for s, stride in zip(strides, value.strides)): @@ -835,7 +859,7 @@ def _field_dtype(field: CudaParameter) -> np.dtype: if field.pointer: return np.dtype(np.uint64) if field.view_ndim is not None: - return _view_dtype(field.view_ndim) + return _view_dtype(field.view_ndim, field.contiguous) return field.dtype @@ -891,6 +915,27 @@ def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str return f"Array{ndim}D<{ctype}>" +def _apply_contiguous( + fields: list[tuple[str, str]], + contiguous: bool | Iterable[str], + what: str, +) -> list[tuple[str, str]]: + """Make array fields ``CArrayD``: all of them, or those named.""" + if contiguous is True: + names = {f for f, ctype in fields if ctype.startswith("Array")} + elif contiguous is False: + return fields + else: + names = {contiguous} if isinstance(contiguous, str) else set(contiguous) + arrays = {f for f, ctype in fields if ctype.startswith("Array")} + if names - arrays: + raise ValueError( + f"{what}: contiguous names {sorted(names - arrays)} that are not " + f"array fields (array fields: {sorted(arrays)})", + ) + return [(f, f"C{ctype}" if f in names else ctype) for f, ctype in fields] + + def _header_guard(name: str) -> str: guard = re.sub(r"\W", "_", name).upper().strip("_") return guard if re.match(r"[A-Z_]", guard) else f"_{guard}" @@ -1084,6 +1129,7 @@ def from_signature( *, int_type: str = "long long", scalar_names: Mapping[str, str] | None = None, + contiguous: bool | Iterable[str] = False, ) -> CudaStruct: """Build the struct from the annotated parameters of a Python function. @@ -1109,6 +1155,10 @@ def from_signature( scalar_names : Mapping[str, str] | None Additional (or changed) mappings from annotation scalar names to C types, e.g. ``{"float": "float"}`` for single precision. + contiguous : bool | Iterable[str] + Array fields that become C-contiguous views (``CArray2D`` + instead of ``Array2D``): ``True`` for all of them, or their + names. Packing such a field raises for a non-contiguous array. Raises ------ @@ -1141,7 +1191,8 @@ def from_signature( continue what = f"parameter {param.name!r} of {getattr(func, '__qualname__', func)}" fields.append((param.name, _pyccel_ctype(param.annotation, scalars, what))) - return cls(name, fields) + what = f"{getattr(func, '__qualname__', func)}" + return cls(name, _apply_contiguous(fields, contiguous, what)) @classmethod def from_pyccel_class( @@ -1154,6 +1205,7 @@ def from_pyccel_class( scalar_names: Mapping[str, str] | None = None, exclude: Sequence[str] = (), attribute_names: bool = True, + contiguous: bool | Iterable[str] = False, ) -> CudaStruct: """Build the struct from the ``__init__`` of a class in a Python source file. @@ -1177,8 +1229,9 @@ def from_pyccel_class( Name of the class in the source. name : str | None Name of the struct type in C; `class_name` by default. - int_type, scalar_names - As for :meth:`from_signature`. + int_type, scalar_names, contiguous + As for :meth:`from_signature` (`contiguous` takes field names, i.e. + attribute names when `attribute_names` is set). exclude : Sequence[str] Parameter or attribute names that do not become fields, e.g. scratch arrays the host class allocates for itself. @@ -1217,6 +1270,7 @@ def from_pyccel_class( raise ValueError(f"{what} has no type annotation") annotation = _annotation_text(arg.annotation) fields.append((field, _pyccel_ctype(annotation, scalars, what))) + fields = _apply_contiguous(fields, contiguous, f"{class_name} in {where}") return cls(class_name if name is None else name, fields) def __repr__(self) -> str: diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 499b888..d3e2de4 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -343,12 +343,18 @@ def emulate_cuda_kernel( f"argument {i} ({param.name}) must be a {param.view_ndim}D " f"array, got {value.ndim}D", ) + if param.contiguous and not value.flags.c_contiguous: + # as on the GPU: a copy would drop what the kernel writes + raise TypeError( + f"argument {i} ({param.name}) must be C-contiguous for " + f"{param.ctype}", + ) buffer = np.ascontiguousarray(value) path = tmp_path / f"{name}.bin" buffer.tofile(path) ctype = param.ctype if param.dtype is not None else "unsigned char" element = ( - re.match(r"Array\dD<(.*)>", ctype).group(1) + re.match(r"C?Array\dD<(.*)>", ctype).group(1) if param.view_ndim is not None else ctype ) @@ -362,10 +368,11 @@ def emulate_cuda_kernel( strides = ", ".join( f"{s // buffer.itemsize}LL" for s in buffer.strides ) + members = f"{{{shape}}}" + if not param.contiguous: + members += f", {{{strides}}}" globals_.append(f"static {ctype} {name}_view;") - inits.append( - f" {name}_view = {ctype}{{{name}, {{{shape}}}, {{{strides}}}}};", - ) + inits.append(f" {name}_view = {ctype}{{{name}, {members}}};") call_args.append(f"{name}_view") else: call_args.append(name) diff --git a/src/cunumpy/cuda/include/cunumpy/array_view.cuh b/src/cunumpy/cuda/include/cunumpy/array_view.cuh index c1ebe6a..d93667f 100644 --- a/src/cunumpy/cuda/include/cunumpy/array_view.cuh +++ b/src/cunumpy/cuda/include/cunumpy/array_view.cuh @@ -1,4 +1,4 @@ -// Strided array views for CUDA kernels, passed by value from Python. +// Array views for CUDA kernels, passed by value from Python. // // Array1D to Array4D describe a (possibly non-contiguous) // device array the way NumPy/CuPy do: a data pointer, a shape and strides. @@ -20,6 +20,19 @@ // the end of this file check the size. Member functions do not change the // layout. // +// CArray1D to CArray4D are C-contiguous (row-major) views: a data +// pointer and a shape, no strides. `a(i, j)` is `data[i * shape[1] + j]`, so +// the last index is always the fast one and the compiler knows it has unit +// stride. A parameter or field of these types only takes C-contiguous arrays; +// cunumpy raises for a non-contiguous view instead of copying it (a copy would +// silently drop what the kernel writes). Their layout, for ndim = 1 to 4: +// +// T* data; // 8 bytes +// long long shape[ndim]; // ndim * 8 bytes +// +// (sizeof == 8 * (1 + ndim)). A CArrayND converts to an ArrayND, so it +// can be passed to device functions written for strided views. +// // Bounds checks: compile with -DCUNUMPY_BOUNDS_CHECK to check every index // against the shape (an out-of-bounds index prints a message and traps the // kernel, which CuPy reports as a CUDA error). Without the macro, indexing is @@ -117,12 +130,111 @@ struct Array4D { } }; -// The layout the Python side packs: pointer, shape, strides, 8-byte aligned. +// C-contiguous views: shape only, the strides follow from it. + +template +struct CArray1D { + T* data; + long long shape[1]; + + __device__ __forceinline__ T& operator()(long long i) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + return data[i]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { return shape[0]; } + + __host__ __device__ operator Array1D() const { + return Array1D{data, {shape[0]}, {1}}; + } +}; + +template +struct CArray2D { + T* data; + long long shape[2]; + + __device__ __forceinline__ T& operator()(long long i, long long j) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + CUNUMPY_CHECK_INDEX(j, 1, shape[1]); + return data[i * shape[1] + j]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { + return shape[0] * shape[1]; + } + + __host__ __device__ operator Array2D() const { + return Array2D{data, {shape[0], shape[1]}, {shape[1], 1}}; + } +}; + +template +struct CArray3D { + T* data; + long long shape[3]; + + __device__ __forceinline__ T& operator()(long long i, long long j, + long long k) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + CUNUMPY_CHECK_INDEX(j, 1, shape[1]); + CUNUMPY_CHECK_INDEX(k, 2, shape[2]); + return data[(i * shape[1] + j) * shape[2] + k]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { + return shape[0] * shape[1] * shape[2]; + } + + __host__ __device__ operator Array3D() const { + return Array3D{data, + {shape[0], shape[1], shape[2]}, + {shape[1] * shape[2], shape[2], 1}}; + } +}; + +template +struct CArray4D { + T* data; + long long shape[4]; + + __device__ __forceinline__ T& operator()(long long i, long long j, + long long k, long long l) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + CUNUMPY_CHECK_INDEX(j, 1, shape[1]); + CUNUMPY_CHECK_INDEX(k, 2, shape[2]); + CUNUMPY_CHECK_INDEX(l, 3, shape[3]); + return data[((i * shape[1] + j) * shape[2] + k) * shape[3] + l]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { + return shape[0] * shape[1] * shape[2] * shape[3]; + } + + __host__ __device__ operator Array4D() const { + return Array4D{data, + {shape[0], shape[1], shape[2], shape[3]}, + {shape[1] * shape[2] * shape[3], shape[2] * shape[3], + shape[3], 1}}; + } +}; + +// The layouts the Python side packs: pointer, shape (and strides), 8-byte +// aligned. static_assert(sizeof(Array1D) == 24, "unexpected Array1D layout"); static_assert(sizeof(Array2D) == 40, "unexpected Array2D layout"); static_assert(sizeof(Array3D) == 56, "unexpected Array3D layout"); static_assert(sizeof(Array4D) == 72, "unexpected Array4D layout"); static_assert(sizeof(Array1D) == 24, "unexpected Array1D layout"); static_assert(alignof(Array2D) == 8, "unexpected Array2D alignment"); +static_assert(sizeof(CArray1D) == 16, "unexpected CArray1D layout"); +static_assert(sizeof(CArray2D) == 24, "unexpected CArray2D layout"); +static_assert(sizeof(CArray3D) == 32, "unexpected CArray3D layout"); +static_assert(sizeof(CArray4D) == 40, "unexpected CArray4D layout"); +static_assert(alignof(CArray2D) == 8, "unexpected CArray2D alignment"); #endif // CUNUMPY_ARRAY_VIEW_CUH diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 1309007..582164e 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -756,6 +756,129 @@ def test_view_parameter_errors(): ) +SCALE_COLUMN_CONTIGUOUS = r""" +#include "cunumpy/array_view.cuh" +#include +extern "C" __global__ +void scale_column(CArray2D a, long long column, double factor) { + CUNUMPY_THREAD_1D(i, a.shape[0]); + a(i, column) *= factor; +} +""" + + +def test_parse_contiguous_view_parameters(): + source = ( + "__global__ void f(CArray1D a, const CArray3D< float > b, " + "Array2D c, CArray4D> d) {}" + ) + params = parse_cuda_signature(source, "f") + assert [(p.ctype, p.view_ndim, p.contiguous, p.dtype) for p in params] == [ + ("CArray1D", 1, True, np.dtype(np.float64)), + ("CArray3D", 3, True, np.dtype(np.float32)), + ("Array2D", 2, False, np.dtype(np.float64)), + ("CArray4D>", 4, True, np.dtype(np.complex128)), + ] + with pytest.raises(ValueError, match="array views"): + parse_cuda_signature("__global__ void f(CArray2D* a) {}", "f") + + +def test_contiguous_view_packs_pointer_and_shape_only(): + kernel = CudaKernel(SCALE_COLUMN_CONTIGUOUS, "scale_column") + a = FakeDeviceArray(np.float64, ptr=0xF00, shape=(4, 3)) + packed, _, _ = kernel.prepare_args(a, 1, 2.0) + assert packed.dtype.names == ("data", "shape") + assert packed.dtype.itemsize == 8 + 2 * 8 # data, shape[2] + assert packed["data"] == 0xF00 + assert packed["shape"].tolist() == [4, 3] + + +def test_contiguous_view_rejects_strided_arrays_instead_of_copying(): + kernel = CudaKernel(SCALE_COLUMN_CONTIGUOUS, "scale_column") + every_second_row = FakeDeviceArray(np.float64, shape=(4, 3), strides=(48, 8)) + with pytest.raises(TypeError, match=r"CArray2D a\) must be C-contiguous"): + kernel.prepare_args(every_second_row, 1, 2.0) + with pytest.raises(TypeError, match="must be a 2D array, got 1D"): + kernel.prepare_args(FakeDeviceArray(np.float64, shape=(3,)), 1, 2.0) + + +def test_struct_with_contiguous_view_fields(): + struct = CudaStruct( + "Markers", + [("markers", "CArray2D"), ("valid", "Array1D"), ("n", "int")], + ) + assert [f.contiguous for f in struct.fields] == [True, False, False] + assert struct.has_views + assert " CArray2D markers;" in struct.declaration + # 24 (CArray2D) + 24 (Array1D) + 4 (int) + 4 padding + assert struct.dtype.itemsize == 56 + value = struct( + markers=FakeDeviceArray(np.float64, ptr=0xA0, shape=(10, 7)), + valid=FakeDeviceArray(np.bool_, shape=(10,)), + n=10, + ) + assert value.packed["markers"]["shape"].tolist() == [10, 7] + assert value.packed["markers"].dtype.names == ("data", "shape") + with pytest.raises(TypeError, match="must be C-contiguous"): + struct( + markers=FakeDeviceArray(np.float64, shape=(10, 3), strides=(56, 8)), + valid=FakeDeviceArray(np.bool_, shape=(10,)), + n=10, + ) + + +def test_from_signature_contiguous(): + class Args: + def __init__(self, markers: "float[:, :]", valid: "bool[:]", n: "int"): ... + + every = CudaStruct.from_signature(Args.__init__, "A", contiguous=True) + assert [f.ctype for f in every.fields] == [ + "CArray2D", + "CArray1D", + "long long", + ] + some = CudaStruct.from_signature(Args.__init__, "A", contiguous=["markers"]) + assert [f.ctype for f in some.fields][:2] == ["CArray2D", "Array1D"] + with pytest.raises(ValueError, match=r"\['n'\] that are not array fields"): + CudaStruct.from_signature(Args.__init__, "A", contiguous=["n"]) + + source = ( + "class MarkerArguments:\n" + " def __init__(self, mks: 'float[:, :]', n: 'int'):\n" + " self.markers = mks\n" + ) + struct = CudaStruct.from_pyccel_class( + source, "MarkerArguments", contiguous=["markers"] + ) + assert [(f.name, f.ctype) for f in struct.fields] == [ + ("markers", "CArray2D"), + ("n", "long long"), + ] + + +def test_scale_column_of_contiguous_view_on_gpu(): + _skip_without_cupy() + import cupy as cp + + a = cp.arange(12, dtype=cp.float64).reshape(4, 3) + expected = a.get() + kernel = CudaKernel(SCALE_COLUMN_CONTIGUOUS, "scale_column", block_size=2) + kernel(a, 1, 10.0, n_threads=4) + expected[:, 1] *= 10.0 + assert np.array_equal(a.get(), expected) + with pytest.raises(TypeError, match="must be C-contiguous"): + kernel(a[::2], 1, 10.0, n_threads=2) + + +def test_contiguous_view_layout_on_gpu(): + _skip_without_cupy() + struct = CudaStruct( + "ContiguousLayout", + [("a", "CArray3D"), ("n", "int"), ("b", "Array2D")], + ) + struct.verify_layout() + + MARKERS = CudaStruct( "Markers", [("markers", "Array2D"), ("valid", "Array1D"), ("n", "int")], diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index 1da8a7d..cd16a75 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -69,6 +69,37 @@ def test_strided_view_written_back_into_the_callers_array(): np.testing.assert_array_equal(markers, expected) +CONTIGUOUS = r""" +#include "cunumpy/array_view.cuh" +#include +__device__ double sum_row(Array2D a, long long i) { + double total = 0.0; + for (long long j = 0; j < a.shape[1]; ++j) total += a(i, j); + return total; +} +extern "C" __global__ +void row_sums(CArray2D a, CArray3D b, CArray1D out) { + CUNUMPY_GRID_STRIDE_1D(i, a.shape[0]) { + out(i) = sum_row(a, i) + b(i, 1, 2); // CArray2D converts to Array2D + } +} +""" + + +def test_contiguous_views(): + a = np.arange(12.0).reshape(4, 3) + b = np.arange(4 * 2 * 3.0).reshape(4, 2, 3) + out = np.zeros(4) + kernel = CudaKernel(CONTIGUOUS, "row_sums") + emulate_cuda_kernel(kernel, a, b, out, grid=1, block=2) + np.testing.assert_array_equal(out, a.sum(axis=1) + b[:, 1, 2]) + + with pytest.raises(TypeError, match="must be C-contiguous"): + emulate_cuda_kernel( + kernel, a[:, :2].copy()[::2], b[::2], out[:2], grid=1, block=2 + ) + + TEMPLATE = r""" template __global__ void power(T* x, int n) { From 0ec730e2198ab1378a62539b575ababd06ff8601 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 6 Oct 2026 15:26:39 +0200 Subject: [PATCH 2/8] Remove pyccel args kernel (#70) * Add C-contiguous array views CArray1D-4D * Remove KernelArguments, resolve_host_args and PyccelStructArguments * Move arguments and kernels to new folders * Bump version number to 0.6.0 --- CHANGELOG.md | 14 ++ README.md | 45 ++--- docs/source/api.md | 173 +++++-------------- docs/source/examples/particle-pusher.md | 2 +- docs/source/kernels/accumulation.md | 2 +- docs/source/kernels/arguments.md | 193 +++++++++------------- docs/source/kernels/cuda-kernel.md | 18 +- docs/source/kernels/debugging.md | 4 +- docs/source/kernels/dispatch.md | 12 +- docs/source/kernels/overview.md | 6 +- docs/source/kernels/pyccel-kernel.md | 2 - docs/source/quickstart.md | 2 +- pyproject.toml | 2 +- src/cunumpy/LLM_GUIDE.md | 62 +++---- src/cunumpy/__init__.py | 15 +- src/cunumpy/__init__.pyi | 1 + src/cunumpy/_cuda_kernel.py | 138 ---------------- src/cunumpy/_dispatch.py | 37 ++--- src/cunumpy/_fake_cupy.py | 4 +- src/cunumpy/_kernel.py | 133 +-------------- src/cunumpy/arguments.py | 42 +++++ src/cunumpy/cuda/__init__.py | 61 ++++--- src/cunumpy/cuda_kernel.py | 4 +- src/cunumpy/kernel_testing.py | 8 +- src/cunumpy/kernels.py | 17 +- tests/unit/test_automatic_launch.py | 3 +- tests/unit/test_cuda_collectives.py | 2 +- tests/unit/test_cuda_kernel.py | 29 ++-- tests/unit/test_cuda_launch_limits.py | 3 +- tests/unit/test_emulation.py | 2 +- tests/unit/test_kernel_dispatch.py | 141 ++-------------- tests/unit/test_kernel_dispatch_arrays.py | 19 ++- tests/unit/test_kernel_testing.py | 9 +- tests/unit/test_mirror.py | 2 +- tests/unit/test_morton.py | 3 +- tests/unit/test_namespaces.py | 23 +++ tests/unit/test_philox.py | 3 +- tests/unit/test_porting_helpers.py | 106 ++---------- 38 files changed, 419 insertions(+), 923 deletions(-) create mode 100644 src/cunumpy/arguments.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d0caef..de369be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- `CudaKernel` and `CudaKernelVariants` moved to `cunumpy.kernels`, next to + `Kernel` and `PyccelKernel`. The argument classes `CudaArguments`, + `CudaStruct`, `CudaStructArguments`, `CudaStructValue` and + `write_cuda_header` moved to the new `cunumpy.arguments`. `cunumpy.cuda` + keeps the device runtime and the CUDA source tools. The old + `cunumpy.cuda.` imports still work with a `DeprecationWarning` and + will be removed in cunumpy 0.6. +- **Removed** `kernels.KernelArguments`, `kernels.resolve_host_args`, + `kernels.PyccelStructArguments` and the `__host_args__()` protocol, with no + replacement. Kernels receive argument objects as they are. Write the host + argument class (e.g. pyccel) and a `CudaStructArguments` with the same + constructor, and let the owner of the arrays build the one for its backend. + `Kernel(dispatch="arrays")` now treats every object with `__cuda_args__()` as + a device argument. - Rename `CUNUMPY_KERNEL_IMPLEMENTATION` to `CUNUMPY_HOST_KERNEL_IMPLEMENTATION` to make its host-only scope explicit. The former environment variable is no longer read; update job scripts. diff --git a/README.md b/README.md index 92d709d..d84a081 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,9 @@ never hide a NumPy name: | Submodule | Contents | |---|---| -| `xp.cuda` | CUDA only: `CudaKernel`, `CudaStruct`, CUDA headers, devices, streams | -| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, host implementations, `fuse` | +| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, host implementations, `fuse` | +| `xp.arguments` | CUDA only: `CudaStruct`, `CudaStructArguments`, `CudaArguments` | +| `xp.cuda` | CUDA only: devices, streams, debug mode, CUDA headers | | `xp.rng` | `random_streams`, `get_rng`, `philox_*` | | `xp.algorithms` | `morton_*`, `sort_by_key`, `cell_offsets`, `segment_boundaries`, `segment_sum`, `SegmentPlan` | | `xp.mpi` | `mpi_buffer`, reusable `MPIStaging`, CUDA-aware MPI | @@ -34,7 +35,7 @@ never hide a NumPy name: | `xp.petsc` | `petsc_vec` | | `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs | -Everything except `xp.cuda` works on both backends. +Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends. ## Install @@ -333,7 +334,7 @@ def axpy(a, x, y, n): # host version, e.g. compiled with Pyccel y[i] += a * x[i] -kernel = xp.kernels.Kernel(axpy, xp.cuda.CudaKernel(AXPY, "axpy")) +kernel = xp.kernels.Kernel(axpy, xp.kernels.CudaKernel(AXPY, "axpy")) with xp.use_backend("cupy"): x = xp.arange(1000, dtype=xp.float64) @@ -356,8 +357,8 @@ the matching memory layout) and packs values into it, which the kernel takes as one parameter: ```python -Vec = xp.cuda.CudaStruct("Vec", [("data", "double*"), ("n", "int")]) -scale = xp.cuda.CudaKernel( +Vec = xp.arguments.CudaStruct("Vec", [("data", "double*"), ("n", "int")]) +scale = xp.kernels.CudaKernel( Vec.declaration + r""" extern "C" __global__ void scale(Vec v, double a) { @@ -379,29 +380,15 @@ device copies. Call it once when the object is built, not per kernel call; on the NumPy backend it raises, so host data is never copied to the device implicitly. When the host kernel takes such a group as one object too (e.g. a Pyccel class -holding NumPy arrays), give the group both forms with `KernelArguments`: -`__host_args__()` returns the object for the host kernel, `__cuda_args__()` -the flattened device arguments. `Kernel` and `PyccelKernel` resolve -`__host_args__()` on the host path and `CudaKernel` flattens `__cuda_args__()` -on the CUDA path, so the call site is the same on both backends and each form -can be built lazily on first access (a CPU run never builds device arguments): +holding NumPy arrays), write a CUDA class with the same constructor and +attributes (a `CudaStructArguments`, see below) and let the owner of the arrays +build the one for the active backend. CuNumpy passes argument objects through +as they are and never converts one form into the other: ```python -class ParticleArguments(xp.kernels.KernelArguments): - def __init__(self, markers): - self.markers = markers - self._host = None - - def __host_args__(self): - if self._host is None: - self._host = MarkerArguments(self.markers) # Pyccel class - return self._host - - def __cuda_args__(self): - return (self.markers, self.markers.shape[0]) - - -kernel(particles.kernel_args, dt, n_threads=n) # host or CUDA kernel +args_class = CudaMarkerArguments if xp.is_gpu(markers) else MarkerArguments +particles.args_markers = args_class(markers, markers.shape[0]) +kernel(particles.args_markers, dt) # host or CUDA kernel ``` Kernels ported from pyccel index arrays like `markers[ip, j]`, which needs @@ -417,11 +404,11 @@ class MarkerArguments: def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... -MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") +MarkerArgs = xp.arguments.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") MarkerArgs.to_header( "marker_args.cuh" ) # Array2D markers; long long n_markers; ... -push = xp.cuda.CudaKernel( +push = xp.kernels.CudaKernel( r""" #include "marker_args.cuh" #include diff --git a/docs/source/api.md b/docs/source/api.md index b172027..ad52468 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -33,15 +33,16 @@ function that needs to follow an input array's location should use The top level of `cunumpy` is the NumPy (or CuPy) namespace plus the functions that select the backend and convert arrays. Everything else is in a submodule, -imported with `cunumpy` (`xp.cuda.CudaKernel`, `xp.rng.random_streams`, ...). +imported with `cunumpy` (`xp.kernels.CudaKernel`, `xp.rng.random_streams`, ...). The submodules are named so that they do not hide a NumPy name (`rng`, not `random`): | Submodule | Backends | Contents | |---|---|---| | `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `scipy`, `require_version` | -| `cunumpy.cuda` | CUDA only | `CudaKernel`, `CudaKernelVariants`, `CudaStruct`, `CudaStructArguments`, `CudaArguments`, CUDA headers, debug mode, device selection and memory, `stream`, `pin_memory` | -| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `KernelArguments`, `PyccelStructArguments`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | +| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | +| `cunumpy.arguments` | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | +| `cunumpy.cuda` | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | | `cunumpy.rng` | both | `random_streams`, `get_rng`, `philox_*` | | `cunumpy.algorithms` | both | `morton_*`, `sort_by_key`, `segment_sum` | | `cunumpy.mpi` | both | `mpi_buffer`, CUDA-aware MPI detection, `local_rank`, `synchronize_for_mpi` | @@ -365,7 +366,7 @@ running on CuPy, and host data is never copied to the device implicitly. If `ValueError`; `name` is the argument name used in error messages. ```python -class DeviceParticles(xp.cuda.CudaArguments): +class DeviceParticles(xp.arguments.CudaArguments): def __init__(self, markers, degree): self.markers = xp.as_device_array(markers, np.float64, ndim=2, name="markers") self.degree = xp.as_device_array(degree, np.int32, ndim=1, name="degree") @@ -853,12 +854,12 @@ lists) are converted back using `is_array`; dictionaries in return values are not recursively converted. On the NumPy path, the original return value and normal Python mutation and exception behavior are preserved. -## `cuda.CudaKernel` +## `kernels.CudaKernel` ### Constructor ```python -xp.cuda.CudaKernel( +xp.kernels.CudaKernel( source, name, *, @@ -873,8 +874,8 @@ xp.cuda.CudaKernel( n_threads_from="auto", check_finite=False, ) -xp.cuda.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) -xp.cuda.CudaKernel.all_from_file(path, **kwargs) +xp.kernels.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) +xp.kernels.CudaKernel.all_from_file(path, **kwargs) ``` Wraps the `__global__` function `name` in the CUDA C `source` (declared @@ -900,7 +901,7 @@ compiles the file once. `xp.cuda.cuda_kernel_names(source)` lists the `__global_ functions of a source string (ignoring comments). ```python -kernels = xp.cuda.CudaKernel.all_from_file("small_kernels.cu", block_size=64) +kernels = xp.kernels.CudaKernel.all_from_file("small_kernels.cu", block_size=64) kernels["scale"](x, 2.0, x.size, n_threads=x.size) kernels["shift"](x, 1.0, x.size, n_threads=x.size) ``` @@ -941,7 +942,7 @@ resolves the quoted includes of its source when it compiles and adds a define with a hash of their contents to the options: ```python -kernel = xp.cuda.CudaKernel.from_file("push/push_cuda.cu", include_dirs=[src_root]) +kernel = xp.kernels.CudaKernel.from_file("push/push_cuda.cu", include_dirs=[src_root]) kernel.included_headers # (Path('push/helpers.cuh'), Path('.../common.cuh')) kernel.options # ('-Ipush', '-I') kernel.compile_options() # options + ('-DCUNUMPY_INCLUDE_HASH=0x3f9a...',) @@ -1004,7 +1005,7 @@ Missing arrays or insufficient array dimensions raise with instructions to give an explicit size. ```python -push = xp.cuda.CudaKernel.from_file("push_cuda.cu") +push = xp.kernels.CudaKernel.from_file("push_cuda.cu") push(positions, velocities, e_field, dt) # n_threads = positions.shape[0] ``` @@ -1037,7 +1038,7 @@ extern "C" __global__ void block_sum(const double* x, double* out, int n) { if (threadIdx.x == 0) out[blockIdx.x] = buffer[0]; } """ -block_sum = xp.cuda.CudaKernel(BLOCK_SUM, "block_sum", block_size=128) +block_sum = xp.kernels.CudaKernel(BLOCK_SUM, "block_sum", block_size=128) (n_blocks,), _ = block_sum.launch_shape(x.size) partial = xp.zeros(n_blocks) block_sum(x, partial, x.size, n_threads=x.size, shared_mem=128 * 8) @@ -1103,7 +1104,7 @@ void scale_column(Array2D a, long long column, double factor) { a(i, column) *= factor; } """ -scale_column = xp.cuda.CudaKernel(SCALE_COLUMN, "scale_column") +scale_column = xp.kernels.CudaKernel(SCALE_COLUMN, "scale_column") view = markers[::2, 1:5] # non-contiguous is fine scale_column(view, 1, 10.0, n_threads=view.shape[0]) ``` @@ -1149,7 +1150,7 @@ __global__ void scale(T* x, T factor, int n) { if (i < n) x[i] = factor * x[i] * (T)N; } """ -scale_f64 = xp.cuda.CudaKernel(SCALE, "scale", template_args=(np.float64, 3)) +scale_f64 = xp.kernels.CudaKernel(SCALE, "scale", template_args=(np.float64, 3)) scale_f64(x, 2.0, x.size, n_threads=x.size) # instantiation scale ``` @@ -1158,8 +1159,8 @@ dimensions and dtype), `CudaKernelVariants` creates and caches one kernel per key: ```python -matvec = xp.cuda.CudaKernelVariants( - lambda ndim, dtype: xp.cuda.CudaKernel( +matvec = xp.kernels.CudaKernelVariants( + lambda ndim, dtype: xp.kernels.CudaKernel( make_source(ndim, xp.cuda.ctype_of(dtype)), "matvec" ) ) @@ -1178,7 +1179,7 @@ threads (see `KernelCatalog.compile_all`). xp.cuda.set_cuda_debug(enabled) xp.cuda.get_cuda_debug() xp.cuda.cuda_debug(enabled=True) # context manager -xp.cuda.CudaKernel(..., debug=None) +xp.kernels.CudaKernel(..., debug=None) kernel.debug_active() kernel.compile_options() xp.cuda.DEBUG_OPTIONS # ("-lineinfo", "-DCUNUMPY_BOUNDS_CHECK") @@ -1214,7 +1215,7 @@ now would use. ```python with xp.cuda.cuda_debug(): - kernel = xp.cuda.CudaKernel(SOURCE, "kernel") + kernel = xp.kernels.CudaKernel(SOURCE, "kernel") kernel( x, y, n, n_threads=n ) # RuntimeError: CUDA error after launching kernel 'kernel' ... @@ -1231,10 +1232,10 @@ CUNUMPY_CUDA_DEBUG=1 compute-sanitizer python -m pytest tests/unit/test_my_kerne Note that after an illegal memory access the CUDA context is unusable; the process (or the pytest run) has to be restarted. -## `cuda.CudaStruct` +## `arguments.CudaStruct` ```python -Particles = xp.cuda.CudaStruct( +Particles = xp.arguments.CudaStruct( "Particles", [("x", "double*"), ("v", "double*"), ("n", "int"), ("charge", "double")], ) @@ -1247,7 +1248,7 @@ extern "C" __global__ void push(Particles p, double dt) { } """ ) -push = xp.cuda.CudaKernel(source, "push", structs=[Particles]) +push = xp.kernels.CudaKernel(source, "push", structs=[Particles]) push(Particles(x=x, v=v, n=x.size, charge=-1.0), 0.1, n_threads=x.size) ``` @@ -1296,7 +1297,7 @@ class MarkerArguments: # the pyccel argument class, e.g. in struphy def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... -MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") +MarkerArgs = xp.arguments.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") print(MarkerArgs.declaration) # struct MarkerArgs { # Array2D markers; @@ -1333,14 +1334,14 @@ or an annotation cannot be mapped. ### Generating headers ```python -xp.cuda.write_cuda_header("pusher_args.cuh", [MarkerArgs, DomainArgs]) +xp.arguments.write_cuda_header("pusher_args.cuh", [MarkerArgs, DomainArgs]) ``` `struct.to_header(path=None, *, guard=None, includes=())` returns the struct definition as a header: an include guard (`_CUH` by default), `#include "cunumpy/array_view.cuh"` if the struct has array view fields, the `includes` (file names or `#include` lines), and the definition. With `path` -the header is also written. `xp.cuda.write_cuda_header(path, structs, guard=None, +the header is also written. `xp.arguments.write_cuda_header(path, structs, guard=None, *, includes=())` writes several structs to one header (the guard defaults to the file name, `pusher_args.cuh` -> `PUSHER_ARGS_CUH`) and returns the source. @@ -1349,7 +1350,7 @@ to the kernels that `#include` it, and keep it in sync with a test: ```python def test_pusher_args_header_is_up_to_date(): - generated = xp.cuda.write_cuda_header( + generated = xp.arguments.write_cuda_header( tmp_path / "pusher_args.cuh", [MarkerArgs, DomainArgs] ) assert Path("kernels/pusher_args.cuh").read_text() == generated @@ -1358,10 +1359,10 @@ def test_pusher_args_header_is_up_to_date(): Kernels created with `structs=[MarkerArgs, ...]` also check a definition in their own source against the Python definition (`check_source`). -## `cuda.CudaStructArguments` +## `arguments.CudaStructArguments` ```python -class MarkerArguments(xp.cuda.CudaStructArguments): +class MarkerArguments(xp.arguments.CudaStructArguments): struct_name = "MarkerArgs" fields = (("markers", "Array2D"), ("valid", "bool*"), ("n_markers", "int")) @@ -1372,7 +1373,7 @@ class MarkerArguments(xp.cuda.CudaStructArguments): self.pack() -push = xp.cuda.CudaKernel(source, "push", structs=[MarkerArguments.struct]) +push = xp.kernels.CudaKernel(source, "push", structs=[MarkerArguments.struct]) push(MarkerArguments(markers, valid), dt, n_threads=markers.shape[0]) ``` @@ -1402,29 +1403,10 @@ built once per subclass when the class is defined and is the class attribute base class (its instances cannot be packed); setting only one raises `TypeError`. Subclasses of a complete class inherit its struct. -## `kernels.PyccelStructArguments` - -`CudaStructArguments` with a host form (the `KernelArguments` protocol). Class -attributes, besides `struct_name` and `fields`: - -* `host_class`: the class of the host argument object, e.g. the - pyccel-compiled class (which cannot inherit from anything). -* `host_fields`: the attributes passed to `host_class(...)`, positionally and - in this order; by default the struct fields. -* `host_copies`: whether `__host_args__()` may build the host object from host - copies of device arrays (default `False`: it raises on the CuPy backend). - Copies are counted by `count_transfers()`; results are not copied back. - -`__host_args__()` builds `host_class(*host_fields)` once, and again when one of -the attributes was replaced (array identity or address, scalar value). -`__cuda_args__()` is the packed struct. `has_device_arrays()` tells whether the -array fields are device arrays; objects holding host arrays are copied and -pickled without packing, and the host object is never pickled. - -## `cuda.CudaArguments` +## `arguments.CudaArguments` ```python -class Particles(xp.cuda.CudaArguments): +class Particles(xp.arguments.CudaArguments): def __init__(self, positions, velocities): self.positions = positions super().__init__(positions, velocities, positions.shape[0]) @@ -1436,83 +1418,11 @@ kernel(dt, Particles(x, v), n_threads=x.shape[0]) Base class for objects passed to a `CudaKernel` as one argument that stands for several kernel parameters. `CudaArguments(*values)` stores the values; `__cuda_args__()` returns them. Subclassing is optional: any object with a -`__cuda_args__()` method returning a tuple is flattened. This lets an -application keep its host argument objects (e.g. Pyccel classes holding NumPy -arrays) and matching device argument objects that reference the same data on -the device, and pass either to the same call. A `CudaArguments` object may also -return struct values (`CudaStructValue.packed`) among its values. - -## `kernels.KernelArguments` - -```python -class ParticleArguments(xp.kernels.KernelArguments): - def __init__(self, particles): - self._particles = particles - self._host = None - self._cuda = None - - def __host_args__(self): - if self._host is None: # e.g. a Pyccel class holding NumPy arrays - self._host = MarkerArguments(self._particles.markers) - return self._host - - def __cuda_args__(self): - if self._cuda is None: # device arrays and scalars, flattened - markers = self._particles.markers - self._cuda = (markers, markers.shape[0], markers.shape[1]) - return self._cuda - - -class Particles: - @property - def kernel_args(self): - if self._kernel_args is None: - self._kernel_args = ParticleArguments(self) - return self._kernel_args - - -push(particles.kernel_args, dt, n_threads=n) # same call on both backends -``` - -Base class for argument objects that have a host form and a device form. A -group of arrays, e.g. the marker data of a particle species, is typically -passed to the host kernel as one object holding NumPy arrays (a Pyccel class) -and to the CUDA kernel as several device arrays and scalars. `KernelArguments` -lets one object stand for both, so a `Kernel` call never branches on the -backend: - -* `__host_args__()` returns the single object the host kernel receives in that - position. `Kernel` (on the NumPy backend) and `PyccelKernel` (always, so the - `missing_cuda="fallback"` path works with the same objects) replace the - argument by this value. -* `__cuda_args__()` returns the tuple of CUDA kernel arguments the object - stands for, the `CudaArguments` protocol above; `CudaKernel` flattens it. - -Only top-level positional and keyword arguments are resolved, not objects -nested in tuples, lists or dicts. The check is made on the type, like for -`__cuda_args__`: an instance attribute named `__host_args__` (e.g. a stored -object) is not treated as the protocol. Subclassing is optional; both methods -of the base class raise `NotImplementedError`, so a subclass overrides the ones -it supports (a `KernelArguments` without `__cuda_args__` raises when it reaches -a `CudaKernel`). - -In the example above both forms are built lazily on first access and cached, -so a CPU run never builds device arguments and a GPU run never builds the host -object. The owner is responsible for invalidating the cache (setting the -stored forms to `None`, or replacing the `ParticleArguments` object) when its -arrays are replaced, e.g. after resizing, `deepcopy` or unpickling. - -### `kernels.resolve_host_args(args, kwargs=None)` - -Returns `(args, kwargs)` with every top-level argument whose type defines a -callable `__host_args__()` replaced by its result; everything else is passed -through untouched. `Kernel` and `PyccelKernel` call it before the host kernel; -it is exported for code that calls host kernels by other means: - -```python -args, kwargs = xp.kernels.resolve_host_args((particles.kernel_args, dt), {"out": out}) -host_push(*args, **kwargs) -``` +`__cuda_args__()` method returning a tuple is flattened. Host kernels +receive their own argument objects (e.g. Pyccel classes holding NumPy arrays) +unchanged: the caller passes the host or the CUDA object, cunumpy never +converts one into the other. A `CudaArguments` object may also return struct +values (`CudaStructValue.packed`) among its values. ## `kernels.Kernel` @@ -1540,10 +1450,9 @@ kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) calls the kernel of the active backend. The launch arguments are passed to the CUDA kernel and ignored by the host kernel. Omitted sizes use the CUDA kernel's -shape-based default or its configured callback. Arguments implementing -`KernelArguments` are replaced by their -`__host_args__()` on the host path and flattened via `__cuda_args__()` on the -CUDA path. `kernel.compile()` compiles the CUDA kernel now and returns whether +shape-based default or its configured callback. Arguments are passed to the +selected kernel as they are (a `CudaKernel` flattens `__cuda_args__()` +objects). `kernel.compile()` compiles the CUDA kernel now and returns whether there is one. Without a CUDA kernel on the CuPy backend, `missing_cuda="raise"` raises @@ -1562,8 +1471,8 @@ device. Passing `host_options` together with a `PyccelKernel` raises `dispatch` decides which kernel a call runs. `"backend"` (default): the CUDA kernel on the CuPy backend, the host kernel on the NumPy backend. `"arrays"`: the CUDA kernel if any top-level argument lives on the GPU (a CuPy -array, or a device-only argument object: one with `__cuda_args__()` but no -`__host_args__()`, such as a `CudaArguments` or a struct value), else the host +array, or a CUDA argument object: one with `__cuda_args__()`, such as a +`CudaArguments` or a struct value), else the host kernel, whatever the backend. Use `"arrays"` in codes that hand host arrays to kernels while CuPy is active (diagnostics, MPI staging, CPU fallbacks): those calls then run the host kernel instead of failing in the CUDA argument checks, diff --git a/docs/source/examples/particle-pusher.md b/docs/source/examples/particle-pusher.md index 8d8bf2b..6b76801 100644 --- a/docs/source/examples/particle-pusher.md +++ b/docs/source/examples/particle-pusher.md @@ -335,7 +335,7 @@ as here, pure Python) host kernels on the CPU. ## Where to go from here * The kernels take five to seven loose arguments. In a real code, group the - particle data with [`KernelArguments` or a `CudaStruct`](../kernels/arguments.md). + particle data with [a `CudaStruct`](../kernels/arguments.md). * If the charge density belongs to a host library (a distributed vector exchanged over MPI), deposit into a [`DeviceMirror`](../kernels/accumulation.md). * For several GPUs, split the particles over MPI ranks, bind one GPU per rank diff --git a/docs/source/kernels/accumulation.md b/docs/source/kernels/accumulation.md index bf0953c..1fead66 100644 --- a/docs/source/kernels/accumulation.md +++ b/docs/source/kernels/accumulation.md @@ -79,7 +79,7 @@ def deposit_host(x, w, rho, n_particles, n_cells, dx): np.add.at(rho, cells[inside], w[:n_particles][inside] / dx) -deposit = xp.kernels.Kernel(deposit_host, xp.cuda.CudaKernel(DEPOSIT, "deposit")) +deposit = xp.kernels.Kernel(deposit_host, xp.kernels.CudaKernel(DEPOSIT, "deposit")) # owned by a host library, e.g. a distributed vector rho_host = np.zeros(64) diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md index 87cf523..a243a71 100644 --- a/docs/source/kernels/arguments.md +++ b/docs/source/kernels/arguments.md @@ -3,18 +3,18 @@ Kernels of real simulations take many arguments: the marker array, its shape, the grid spacing, spline degrees, knot vectors, a dozen parameters. Listing them at every call site is error-prone, and every new field means editing every -kernel signature. CuNumpy offers three ways to pass a group of values as one -argument: +kernel signature. CuNumpy groups them for CUDA kernels: -| Tool | Host kernel sees | CUDA kernel sees | Use when | -| --- | --- | --- | --- | -| `CudaArguments` | (not used) | several parameters, flattened | grouping device arguments only | -| `KernelArguments` | one object (`__host_args__()`) | several parameters (`__cuda_args__()`) | the host kernel takes an argument class, the CUDA kernel flat parameters | -| `CudaStruct` | (not used) | one C struct parameter | many fields; one definition shared by all CUDA kernels | -| `CudaStructArguments` | (not used) | one C struct parameter | the struct as a class: an object with attributes, built once and reused | +| Tool | CUDA kernel sees | Use when | +| --- | --- | --- | +| `CudaArguments` | several parameters, flattened | grouping a few device arrays and scalars | +| `CudaStruct` | one C struct parameter | many fields; one definition shared by all CUDA kernels | +| `CudaStructArguments` | one C struct parameter | the struct as a class: an object with attributes, built once and reused | -They combine: a `KernelArguments` object can return a `CudaStruct` value from -`__cuda_args__()`. +CuNumpy never converts an argument object between a host form and a CUDA +form. A host kernel gets its own argument object (e.g. a pyccel class of NumPy +arrays), a CUDA kernel gets a CUDA one, and the code that owns the arrays +decides which one to build; see [Host and CUDA argument classes](#host-and-cuda-argument-classes). ## `CudaArguments`: flatten into several parameters @@ -28,7 +28,7 @@ import numpy as np import cunumpy as xp -class DeviceParticles(xp.cuda.CudaArguments): +class DeviceParticles(xp.arguments.CudaArguments): def __init__(self, positions, velocities): self.positions = xp.as_device_array( positions, np.float64, ndim=2, name="positions" @@ -43,7 +43,7 @@ PUSH = r""" extern "C" __global__ void push(double dt, double* x, const double* v, int n) { ... } """ -push = xp.cuda.CudaKernel(PUSH, "push") +push = xp.kernels.CudaKernel(PUSH, "push") particles = DeviceParticles(x, v) push(0.1, particles, n_threads=particles.positions.shape[0]) # -> push(0.1, x, v, n) ``` @@ -52,54 +52,6 @@ Build the object once and reuse it. `as_device_array()` references existing device arrays of the right dtype and copies everything else exactly once (see [Data movement](../guides/data-movement.md), section "Build device arguments once"). -## `KernelArguments`: one object, two forms - -A host kernel compiled with Pyccel often receives a group of arrays as an -instance of an argument class, while the CUDA kernel receives the arrays as -separate parameters. `KernelArguments` lets one object stand for both, so the -call site never branches on the backend: - -```python -class MarkerArguments: # the host argument class, e.g. compiled with Pyccel - def __init__(self, markers: "float[:, :]", n_markers: int): - self.markers = markers - self.n_markers = n_markers - - -class ParticleArguments(xp.kernels.KernelArguments): - def __init__(self, particles): - self._particles = particles - self._host = None - self._cuda = None - - def __host_args__(self): - if self._host is None: - markers = self._particles.markers - self._host = MarkerArguments(markers, markers.shape[0]) - return self._host - - def __cuda_args__(self): - if self._cuda is None: - markers = self._particles.markers - self._cuda = (markers, markers.shape[0], markers.shape[1]) - return self._cuda - - -args = ParticleArguments(particles) -push(args, dt, n_threads=n_markers) # a Kernel: same call on both backends -``` - -* `Kernel` (on the NumPy backend) and `PyccelKernel` replace the argument with - `__host_args__()`; `CudaKernel` flattens `__cuda_args__()`. -* Only top-level arguments are resolved, not objects inside lists or dicts. -* Build both forms lazily, as above: a CPU run never builds device arguments, a - GPU run never builds the host object. -* **Invalidate the cache** when the underlying arrays are replaced (resizing, - `deepcopy`, unpickling): reset the stored forms or create a new arguments - object. Stale cached arguments point at the old arrays. -* `xp.kernels.resolve_host_args(args, kwargs)` applies the host replacement, for code - that calls host kernels without `Kernel`. - ## `CudaStruct`: one C struct, defined once A `CudaStruct` defines a C struct in Python: its C declaration for the kernel @@ -107,7 +59,7 @@ source, its exact memory layout, and a packer for values. The kernel takes the struct by value as one parameter: ```python -Particles = xp.cuda.CudaStruct( +Particles = xp.arguments.CudaStruct( "Particles", [("x", "double*"), ("v", "double*"), ("n", "long long"), ("charge", "double")], ) @@ -121,7 +73,7 @@ extern "C" __global__ void push(Particles p, double dt) { } """ ) -push = xp.cuda.CudaKernel(PUSH, "push", structs=[Particles]) +push = xp.kernels.CudaKernel(PUSH, "push", structs=[Particles]) value = Particles(x=x, v=v, n=x.size, charge=-1.0) push(value, 0.1, n_threads=x.size) @@ -150,7 +102,7 @@ species, domain or grid, and kept next to the host argument object), subclass field values as attributes and is passed to kernels as it is: ```python -class CudaMarkerArguments(xp.cuda.CudaStructArguments): +class CudaMarkerArguments(xp.arguments.CudaStructArguments): struct_name = "MarkerArgs" fields = ( ("markers", "double*"), @@ -170,8 +122,8 @@ class CudaMarkerArguments(xp.cuda.CudaStructArguments): self.pack() -xp.cuda.write_cuda_header("kernels/marker_args.cuh", [CudaMarkerArguments.struct]) -push = xp.cuda.CudaKernel.from_file( +xp.arguments.write_cuda_header("kernels/marker_args.cuh", [CudaMarkerArguments.struct]) +push = xp.kernels.CudaKernel.from_file( "kernels/push_cuda.cu", structs=[CudaMarkerArguments.struct] ) @@ -189,7 +141,7 @@ push(args, dt, n_threads=args.n_markers) array: ```python - class CudaMarkerArguments(xp.cuda.CudaStructArguments): + class CudaMarkerArguments(xp.arguments.CudaStructArguments): struct_name = "MarkerArgs" fields = (("markers", "Array2D"), ("n_markers", "int")) @@ -210,50 +162,8 @@ push(args, dt, n_threads=args.n_markers) keep working on it: assign the new array to the attribute instead. * Copies and unpickled objects are packed again from their own arrays, so a `deepcopy` never points at the device memory of the original. -* When the same call site must also reach a host kernel, pair it with the host - argument object in a `KernelArguments`: `__host_args__()` returns the host - object, `__cuda_args__()` returns `cuda_args.__cuda_args__()`. - -### `PyccelStructArguments`: a pyccel host class and a struct - -When the host kernels take a pyccel-compiled argument class, that class cannot -inherit from `CudaStructArguments` (or anything else). `PyccelStructArguments` -holds it instead: the same object is passed to a `Kernel` on both backends, and -arrives as the pyccel object on the host path and as the struct on the device -path: - -```python -from my_sim.kernel_arguments import pusher_args_kernels # compiled by pyccel - - -class MarkerArguments(xp.kernels.PyccelStructArguments): - struct_name = "MarkerArgs" - fields = (("markers", "Array2D"), ("Np", "long long"), ("n_markers", "int")) - host_class = pusher_args_kernels.MarkerArguments - host_fields = ("markers", "Np") # its constructor arguments, in order - - def __init__(self, markers, Np): - self.markers = markers # NumPy or CuPy, whatever the owner has - self.Np = Np - self.n_markers = markers.shape[0] - if self.has_device_arrays(): - self.pack() # fail early on a bad device array - - -args = MarkerArguments(particles.markers, Np) -push(args, dt, n_threads=args.n_markers) # Kernel: same call on both backends -``` - -* `__host_args__()` builds `host_class(*host_fields)` once and again when one of - those attributes was replaced (a resized array, a changed scalar). The host - object is not pickled; a copy or an unpickled object rebuilds it. -* On the CuPy backend the attributes are device arrays, and there is no host - form: `__host_args__()` raises. Set `host_copies = True` on a class whose host - kernels only *read* the arrays (an evaluation, not a push): the host object - is then built from host copies, which `count_transfers()` reports, and what - the host kernel writes is not copied back. -* Objects holding host arrays are copied and pickled without packing; the - struct is only built from device arrays. +* For the host kernels, write the matching host class and let the owner choose, + see [Host and CUDA argument classes](#host-and-cuda-argument-classes). ### Check the layout against the compiler @@ -281,7 +191,7 @@ class MarkerArguments: def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... -MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") +MarkerArgs = xp.arguments.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") print(MarkerArgs.declaration) ``` @@ -307,7 +217,7 @@ the parameter is stored in (`self.first_init_idx = first_pusher_idx` gives a field `first_init_idx`), and skips the parameters in `exclude=`: ```python -MarkerArgs = xp.cuda.CudaStruct.from_pyccel_class( +MarkerArgs = xp.arguments.CudaStruct.from_pyccel_class( "my_sim/kernel_arguments/pusher_args_kernels.py", "MarkerArguments", "MarkerArgs" ) ``` @@ -335,7 +245,7 @@ C-contiguous (a marker array, a grid), declare it as `CArray1D` to fast one, as in the row-major memory the host code uses. ```python -MarkerArgs = xp.cuda.CudaStruct.from_pyccel_class( +MarkerArgs = xp.arguments.CudaStruct.from_pyccel_class( "my_sim/kernel_arguments/pusher_args_kernels.py", "MarkerArguments", "MarkerArgs", @@ -365,7 +275,7 @@ Kernels in `.cu` files include the struct from a header. Generate it from the Python definition and commit it: ```python -xp.cuda.write_cuda_header("kernels/marker_args.cuh", [MarkerArgs, DomainArgs]) +xp.arguments.write_cuda_header("kernels/marker_args.cuh", [MarkerArgs, DomainArgs]) ``` The header gets an include guard (`MARKER_ARGS_CUH`), the `array_view.cuh` @@ -377,19 +287,64 @@ from pathlib import Path def test_marker_args_header_is_up_to_date(tmp_path): - generated = xp.cuda.write_cuda_header( + generated = xp.arguments.write_cuda_header( tmp_path / "marker_args.cuh", [MarkerArgs, DomainArgs] ) assert Path("kernels/marker_args.cuh").read_text() == generated ``` +## Host and CUDA argument classes + +A host kernel compiled with Pyccel often receives a group of arrays as an +instance of an argument class. Write its CUDA counterpart as a +`CudaStructArguments` subclass with the same constructor and attributes, and +let the owner of the arrays create the one that matches the backend: + +```python +from my_sim.kernel_arguments.pusher_args_kernels import MarkerArguments # pyccel + + +class CudaMarkerArguments(xp.arguments.CudaStructArguments): + """CUDA version of MarkerArguments: same constructor, same attributes.""" + + struct_name = "MarkerArgs" + fields = (("markers", "Array2D"), ("n_markers", "int")) + + def __init__(self, markers, n_markers): + self.markers = markers # a CuPy array + self.n_markers = n_markers + self.pack() + + +class Particles: + def __init__(self, markers): + self.markers = markers + args_class = CudaMarkerArguments if xp.is_gpu(markers) else MarkerArguments + self.args_markers = args_class(markers, markers.shape[0]) + + +push = xp.kernels.Kernel.from_folder("my_sim.kernels.push") +push(particles.args_markers, dt) # the pyccel object on NumPy, the struct on CuPy +``` + +* `Kernel` and `PyccelKernel` pass the host object to the host kernel as it + is; `CudaKernel` flattens a CUDA argument object with `__cuda_args__()`. +* Pass the CUDA object to the host kernel, or the pyccel object to the CUDA + kernel, and the kernel's own argument checks raise. Nothing is copied + behind your back. +* A test can check that the two classes of a pair stay in sync, by comparing + `CudaStruct.from_pyccel_class(...)` (the pyccel constructor) with + `CudaMarkerArguments.struct.fields`. + ## Choosing * Start with plain arguments. Group when the same set of five or more values appears in several kernels. -* Use `KernelArguments` when the host kernels already take argument objects; - it keeps the call sites identical. * Use a `CudaStruct` when many CUDA kernels take the same group, or when signatures become too long to read and keep in sync. -* Generate the struct with `from_signature` when a host argument class exists, - so the two cannot drift apart. +* When the host kernels take an argument class, write a `CudaStructArguments` + with the same constructor and attributes, and build one of the two per + backend. +* Generate the struct with `from_signature` or `from_pyccel_class` when a host + argument class exists, or test that the two field lists match, so the two + cannot drift apart. diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index 6e77948..7552ffa 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -19,7 +19,7 @@ void axpy(double a, const double* x, double* y, int n) { } """ -axpy = xp.cuda.CudaKernel(AXPY, "axpy") +axpy = xp.kernels.CudaKernel(AXPY, "axpy") xp.set_backend("cupy") x = xp.arange(10_000, dtype=xp.float64) @@ -138,7 +138,7 @@ extern "C" __global__ void block_sum(const double* x, double* out, int n) { if (threadIdx.x == 0) out[blockIdx.x] = buffer[0]; } """ -block_sum = xp.cuda.CudaKernel(BLOCK_SUM, "block_sum", block_size=128) +block_sum = xp.kernels.CudaKernel(BLOCK_SUM, "block_sum", block_size=128) (n_blocks,), (threads,) = block_sum.launch_shape(x.size) partial = xp.zeros(n_blocks) block_sum(x, partial, x.size, n_threads=x.size, shared_mem=threads * 8) @@ -151,7 +151,7 @@ Keeping CUDA source in `.cu` files gives editor support and lets kernels share headers: ```python -push = xp.cuda.CudaKernel.from_file("kernels/push/push_cuda.cu") # kernel name "push" +push = xp.kernels.CudaKernel.from_file("kernels/push/push_cuda.cu") # kernel name "push" ``` `from_file` derives the kernel name from the file name minus the `_cuda.cu` @@ -162,7 +162,7 @@ A file with several small kernels is loaded at once with `all_from_file`, which returns a dict by name; the kernels share one compilation: ```python -ops = xp.cuda.CudaKernel.all_from_file("kernels/vector_ops.cu", block_size=256) +ops = xp.kernels.CudaKernel.all_from_file("kernels/vector_ops.cu", block_size=256) ops["scale"](x, 2.0, x.size, n_threads=x.size) ops["shift"](x, 1.0, x.size, n_threads=x.size) ``` @@ -210,7 +210,7 @@ void scale_column(Array2D a, long long column, double factor) { a(i, column) *= factor; } """ -scale_column = xp.cuda.CudaKernel(SCALE_COLUMN, "scale_column") +scale_column = xp.kernels.CudaKernel(SCALE_COLUMN, "scale_column") markers = xp.zeros((1000, 7)) view = markers[::2, 1:5] # non-contiguous view is fine @@ -247,8 +247,8 @@ __global__ void scale(T* x, T factor, int n) { if (i < n) x[i] *= factor; } """ -scale_f64 = xp.cuda.CudaKernel(SCALE, "scale", template_args=(np.float64,)) -scale_f32 = xp.cuda.CudaKernel(SCALE, "scale", template_args=(np.float32,)) +scale_f64 = xp.kernels.CudaKernel(SCALE, "scale", template_args=(np.float64,)) +scale_f32 = xp.kernels.CudaKernel(SCALE, "scale", template_args=(np.float32,)) ``` When the source itself is generated per variant (unrolled loops per dimension, @@ -257,10 +257,10 @@ on first use and caches it: ```python def make_matvec(ndim, dtype): - return xp.cuda.CudaKernel(generate_source(ndim, xp.cuda.ctype_of(dtype)), "matvec") + return xp.kernels.CudaKernel(generate_source(ndim, xp.cuda.ctype_of(dtype)), "matvec") -matvec = xp.cuda.CudaKernelVariants(make_matvec) +matvec = xp.kernels.CudaKernelVariants(make_matvec) matvec.get(3, np.float64)(mat, x, out, n_threads=out.size) matvec.compile_all([(3, np.float64), (3, np.complex128)], jobs=4) # at setup ``` diff --git a/docs/source/kernels/debugging.md b/docs/source/kernels/debugging.md index 4333eb1..5a2493e 100644 --- a/docs/source/kernels/debugging.md +++ b/docs/source/kernels/debugging.md @@ -19,7 +19,7 @@ CUNUMPY_CUDA_DEBUG=1 python simulate.py # whole process xp.cuda.set_cuda_debug(True) # globally, from now on with xp.cuda.cuda_debug(): # for a block ... -xp.cuda.CudaKernel( +xp.kernels.CudaKernel( src, "push", debug=True ) # one kernel, regardless of the global setting ``` @@ -39,7 +39,7 @@ In debug mode a `CudaKernel`: ```python with xp.cuda.cuda_debug(): - push = xp.cuda.CudaKernel(SOURCE, "push") + push = xp.kernels.CudaKernel(SOURCE, "push") push(markers, dt, n, n_threads=n) # RuntimeError: CUDA error after launching kernel 'push' with grid (79,) and block (128,): ... ``` diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index ff43991..9cd99dd 100644 --- a/docs/source/kernels/dispatch.md +++ b/docs/source/kernels/dispatch.md @@ -24,7 +24,7 @@ def axpy_host(a, x, y, n): y[i] += a * x[i] -axpy = xp.kernels.Kernel(axpy_host, xp.cuda.CudaKernel(AXPY, "axpy"), name="axpy") +axpy = xp.kernels.Kernel(axpy_host, xp.kernels.CudaKernel(AXPY, "axpy"), name="axpy") axpy(2.0, x, y, x.size) # infer n_threads = x.shape[0] ``` @@ -38,8 +38,9 @@ axpy(2.0, x, y, x.size) # infer n_threads = x.shape[0] * A plain host function is wrapped in a [`PyccelKernel`](pyccel-kernel.md); pass `host_options={"outputs": (2,)}` to configure that wrapper, or pass a `PyccelKernel` you built yourself. -* Kernels take positional arguments only. [`KernelArguments`](arguments.md) - objects are resolved per backend. +* Kernels take positional arguments only, passed to the selected kernel as + they are: pass the host argument objects on the host and the CUDA ones on + the device (see [Kernel arguments and structs](arguments.md)). * `kernel.compile()` compiles the CUDA kernel now; `kernel.has_cuda` tells whether there is one. @@ -261,9 +262,8 @@ catalog["gather"](host_positions, host_field, host_result) # host kernel, also ``` The CUDA kernel runs if any top-level argument is on the GPU: a CuPy array, or -a device-only argument object (`CudaArguments`, a struct value). A -`KernelArguments` object has both forms and does not decide on its own. Host -arguments go to the host function directly, without conversion. +a CUDA argument object (`CudaArguments`, `CudaStructArguments`, a struct +value). Host arguments go to the host function directly, without conversion. The arguments of one call must then all be on one side, with the dtype and layout the kernels take. `xp.kernels.as_kernel_array(value, like, dtype)` brings an diff --git a/docs/source/kernels/overview.md b/docs/source/kernels/overview.md index 2ad1ddf..4c892fd 100644 --- a/docs/source/kernels/overview.md +++ b/docs/source/kernels/overview.md @@ -18,7 +18,7 @@ unchanged. | `CudaKernel` | wraps a CUDA C kernel, checks every call against its signature | [Writing CUDA kernels](cuda-kernel.md) | | `Kernel` | a host kernel plus its CUDA kernel; calls the one matching the backend | [Pairing host and CUDA kernels](dispatch.md) | | `KernelCatalog` | all `Kernel`s of a package, found by folder convention | [Pairing host and CUDA kernels](dispatch.md) | -| `CudaArguments`, `KernelArguments`, `CudaStruct` | pass a group of arrays and scalars as one argument | [Kernel arguments and structs](arguments.md) | +| `CudaArguments`, `CudaStruct`, `CudaStructArguments` | pass a group of arrays and scalars as one argument | [Kernel arguments and structs](arguments.md) | | `DeviceMirror`, `cunumpy/atomic.cuh` | scatter-add into a buffer owned by a host library | [Accumulation kernels](accumulation.md) | | `cunumpy.kernel_testing` | check that host and CUDA kernels compute the same | [Testing kernels](testing.md) | @@ -34,8 +34,8 @@ unchanged. host/CUDA pair in a `Kernel`, collect them in a `KernelCatalog`, and test each pair with `assert_kernels_agree`. Unported kernels either raise or fall back to the host version. -* **"My kernels take ten arrays each."** Group them with `KernelArguments` (host - form and device form of the same data) or a `CudaStruct`. +* **"My kernels take ten arrays each."** Group them in a `CudaStruct`, or a + `CudaStructArguments` class next to the host argument class. ## A typical porting workflow diff --git a/docs/source/kernels/pyccel-kernel.md b/docs/source/kernels/pyccel-kernel.md index 918fcf7..22db221 100644 --- a/docs/source/kernels/pyccel-kernel.md +++ b/docs/source/kernels/pyccel-kernel.md @@ -83,8 +83,6 @@ Other details: (for NumPy subclasses). * **`use_cupy`** forces conversion on (`True`) or off (`False`); the default `None` decides per call. -* **Argument objects with two forms** (`KernelArguments`) are replaced by their - `__host_args__()` first; see [Kernel arguments and structs](arguments.md). ## Costs and when to move on diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index 0ccde3c..22d054e 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -91,7 +91,7 @@ def axpy(a, x, y, n): # the host version y[:n] += a * x[:n] -kernel = xp.kernels.Kernel(axpy, xp.cuda.CudaKernel(AXPY, "axpy")) +kernel = xp.kernels.Kernel(axpy, xp.kernels.CudaKernel(AXPY, "axpy")) x = xp.arange(1000, dtype=xp.float64) y = xp.zeros(1000) diff --git a/pyproject.toml b/pyproject.toml index 60f12bf..1ac81b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.5.0" +version = "0.6.0" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 630a4e3..34c26ef 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -18,7 +18,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. with one rank per GPU, profiling. * A kernel layer for porting compiled CPU kernels (Pyccel, Numba, Python loops) to CUDA one at a time: `PyccelKernel`, `CudaKernel`, `Kernel`, - `KernelCatalog`, `KernelArguments`, `CudaStruct`, `DeviceMirror`, and test + `KernelCatalog`, `CudaStruct`, `CudaStructArguments`, `DeviceMirror`, and test helpers in `cunumpy.kernel_testing`. ## Hard rules @@ -54,11 +54,13 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. An argument the kernel writes but that is not declared leaves stale device data, silently. If unsure, leave `outputs=None` (copies everything back). 11. **Helpers are in submodules; the top level is NumPy plus backend control.** - `xp.cuda.CudaKernel`, `xp.kernels.Kernel`, `xp.rng.random_streams`, + `xp.kernels.CudaKernel`, `xp.kernels.Kernel`, `xp.arguments.CudaStructArguments`, + `xp.rng.random_streams`, `xp.algorithms.morton_keys`, `xp.mpi.mpi_buffer`, `xp.profiling.timed_region`, `xp.memory.HostStaging`, `xp.petsc.petsc_vec`; kernel test helpers in `cunumpy.kernel_testing`. The old top-level names (`xp.CudaKernel`) and - `cunumpy.testing` are deprecated (removed in 0.6); do not write new code + `cunumpy.testing`, and the kernel and argument classes in `xp.cuda` + (`xp.cuda.CudaKernel`), are deprecated (removed in 0.6); do not write new code with them. Modules starting with `_` (`cunumpy._cuda_kernel`, ...) are private; never import from them. @@ -71,7 +73,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | normalize inputs at an API boundary | `xp.to_cunumpy(a)` once | | hand data to SciPy/matplotlib/h5py | `xp.to_numpy(a)` | | call an existing NumPy-only kernel with GPU arrays (slow, correct) | `xp.kernels.PyccelKernel(fn, outputs=(...))` | -| launch a hand-written CUDA C kernel | `xp.cuda.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` | +| launch a hand-written CUDA C kernel | `xp.kernels.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` | | host kernel + CUDA port, chosen by backend | `xp.kernels.Kernel(host_fn, cuda_kernel_or_None)` | | many kernels in a package, ported incrementally | `xp.kernels.KernelCatalog.from_package(__name__, missing_cuda="fallback")` | | host kernels compiled at first call (your compile function), NumPy fallback | `from_package(..., host_suffix="_pyccel", compile_host=my_compile, host_fallback={...})` -> `xp.kernels.CompiledHostKernel` | @@ -89,8 +91,8 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | copy device arrays to the host for output without stalling | `xp.memory.HostStaging(shape, dtype)`: `c = staging.copy(a)` ... `c.result()` | | PIC recipes (compaction, sort by cell, MPI exchange, graphs) | docs guide "Particle codes" | | reproducible random numbers per MPI rank | `xp.rng.random_streams.seed(seed, rank=rank)`, then `xp.rng.random_streams.normal(...)` / `.generator()` | -| group arrays/scalars into one kernel argument | `xp.cuda.CudaArguments` (device only), `xp.kernels.KernelArguments` (host object + device tuple), `xp.cuda.CudaStruct` (C struct), `xp.cuda.CudaStructArguments` (C struct as a class) | -| CUDA struct from a Pyccel argument class | `xp.cuda.CudaStruct.from_signature(Cls.__init__, "Name")`, `xp.cuda.write_cuda_header(...)` | +| group arrays/scalars into one kernel argument | `xp.arguments.CudaArguments` (flattened), `xp.arguments.CudaStruct` (C struct), `xp.arguments.CudaStructArguments` (C struct as a class); host kernels take their own argument objects, the caller picks one per backend | +| CUDA struct from a Pyccel argument class | `xp.arguments.CudaStruct.from_signature(Cls.__init__, "Name")`, `xp.arguments.write_cuda_header(...)` | | SciPy (sparse, sparse.linalg, fft, special, ndimage, ...) on either backend | `xp.scipy..` (SciPy or `cupyx.scipy`); `xp.scipy.special.available(name)` | | chain of elementwise operations as one GPU kernel | `@xp.kernels.fuse` (`cupy.fuse` for CuPy arrays, plain call otherwise) | | PETSc solve on device arrays without copies | `xp.petsc.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls | @@ -263,17 +265,17 @@ CuPy. Does not compile anything. `CudaKernel`: ```python -k = xp.cuda.CudaKernel(source, name, *, block_size=128, options=(), include_dirs=(), +k = xp.kernels.CudaKernel(source, name, *, block_size=128, options=(), include_dirs=(), source_dir=None, structs=(), template_args=None, check_signature=True, debug=None) -k = xp.cuda.CudaKernel.from_file("push/push_cuda.cu") # name "push" -ks = xp.cuda.CudaKernel.all_from_file("ops.cu") # dict name -> kernel +k = xp.kernels.CudaKernel.from_file("push/push_cuda.cu") # name "push" +ks = xp.kernels.CudaKernel.all_from_file("ops.cu") # dict name -> kernel k(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) k.compile(log_stream=None); k.recompile(log_stream=None) k.is_compiled # successful compilation on the current CUDA device k.launch_shape(n_threads=None, *, grid=None, block=None, args=None) -> (grid, block) k.included_headers; k.compile_options(); k.debug_active() -xp.cuda.CudaKernelVariants(factory).get(*key); .compile_all(keys, jobs=1) +xp.kernels.CudaKernelVariants(factory).get(*key); .compile_all(keys, jobs=1) xp.cuda.ctype_of(np.float64) == "double" xp.cuda.cuda_kernel_names(source); xp.cuda.parse_cuda_signature(source, name) xp.cuda.cuda_include_dir() @@ -339,20 +341,12 @@ function `` (host); optional `pkg//_cuda.cu` defines Argument objects: ```python -class Dev(xp.cuda.CudaArguments): # flattened into several CUDA params +class Dev(xp.arguments.CudaArguments): # flattened into several CUDA params def __init__(self, x, n): super().__init__(x, n) -class Args(xp.kernels.KernelArguments): # one object, host form + device form - def __host_args__(self): - return host_object # host kernel gets this - - def __cuda_args__(self): - return (arr, n, ...) # CUDA kernel gets these, flattened - - -S = xp.cuda.CudaStruct( +S = xp.arguments.CudaStruct( "S", [("x", "double*"), ("n", "long long"), ("a", "Array2D")] ) S.declaration @@ -360,12 +354,12 @@ S.dtype S.to_header(path) value = S(x=..., n=..., a=...) S.verify_layout() # GPU test: compiler layout == S.dtype (also verify_layout("hdr.cuh")) -S = xp.cuda.CudaStruct.from_signature(Cls.__init__, "S", int_type="long long") -xp.cuda.write_cuda_header("args.cuh", [S1, S2]) +S = xp.arguments.CudaStruct.from_signature(Cls.__init__, "S", int_type="long long") +xp.arguments.write_cuda_header("args.cuh", [S1, S2]) class A( - xp.cuda.CudaStructArguments + xp.arguments.CudaStructArguments ): # the struct as a class; A.struct is the CudaStruct struct_name = "A" fields = (("x", "double*"), ("n", "int")) @@ -375,12 +369,14 @@ class A( self.pack() # repacks itself when a field changes; copies repack -xp.cuda.CudaKernel(S.declaration + src, "k", structs=[S]) -xp.kernels.resolve_host_args(args, kwargs) +xp.kernels.CudaKernel(S.declaration + src, "k", structs=[S]) + +# host kernels: a pyccel class and a CudaStructArguments with the same +# constructor; the owner builds one per backend, cunumpy never converts them +args = (CudaMarkerArguments if xp.is_gpu(markers) else MarkerArguments)(markers, n) ``` -Only top-level arguments are resolved. Cache both forms lazily and invalidate -them when the underlying arrays are replaced. A packed struct holds device +A packed struct holds device addresses: re-pack after replacing an array (a `CudaStructArguments` does this itself at the next launch; make its fields properties to follow an owner's arrays). @@ -409,7 +405,7 @@ Debugging: ```python xp.cuda.set_cuda_debug(True); xp.cuda.get_cuda_debug(); with xp.cuda.cuda_debug(): ... -xp.cuda.CudaKernel(..., debug=True) # env: CUNUMPY_CUDA_DEBUG=1 +xp.kernels.CudaKernel(..., debug=True) # env: CUNUMPY_CUDA_DEBUG=1 ``` Debug mode adds `-lineinfo -DCUNUMPY_BOUNDS_CHECK` at compile time and @@ -443,12 +439,8 @@ def test_parity(kernel): check_parity(kernel) # without a GPU: CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy pytest (fake CuPy: strict host # stand-in, no kernel launches; fake_cupy_active(); requires_cupy skips) -# argument classes with a pyccel host class: one object on both backends -class MarkerArguments(xp.kernels.PyccelStructArguments): - struct_name = "MarkerArgs"; fields = (("markers", "Array2D"), ("Np", "long long")) - host_class = pusher_args_kernels.MarkerArguments # pyccel class; cannot inherit - host_fields = ("markers", "Np") # its constructor args, in order -MarkerArgs = xp.cuda.CudaStruct.from_pyccel_class("pusher_args_kernels.py", "MarkerArguments", "MarkerArgs") +# struct fields from a pyccel argument class (contiguous=True or names -> CArray2D) +MarkerArgs = xp.arguments.CudaStruct.from_pyccel_class("pusher_args_kernels.py", "MarkerArguments", "MarkerArgs") kernel.n_threads_from = lambda args: args[0].n_markers # launch size from an argument kernel.check_finite = True # NaN/inf after each launch (debug) ``` @@ -497,7 +489,7 @@ def scale_host(x: "float[:]", a: float, n: int): scale = xp.kernels.Kernel( - scale_host, xp.cuda.CudaKernel(SRC, "scale"), host_options={"outputs": (0,)} + scale_host, xp.kernels.CudaKernel(SRC, "scale"), host_options={"outputs": (0,)} ) scale(x, 2.0, x.size, n_threads=x.size) ``` diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 935ac6d..b6d8b70 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -3,7 +3,18 @@ import warnings as _warnings from importlib.metadata import PackageNotFoundError, version -from cunumpy import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp +from cunumpy import ( + algorithms, + arguments, + cuda, + kernels, + memory, + mpi, + petsc, + profiling, + rng, + xp, +) from cunumpy._scipy_backend import scipy from cunumpy.xp import ( as_device_array, @@ -29,6 +40,7 @@ # moved to. They still resolve (with a DeprecationWarning) until cunumpy 0.6. _MOVED = { **dict.fromkeys(cuda.__all__, "cuda"), + **dict.fromkeys(arguments.__all__, "arguments"), **dict.fromkeys( ( name @@ -97,6 +109,7 @@ def require_version(minimum: str) -> None: __all__ = [ "__version__", "algorithms", + "arguments", "as_device_array", "assert_same_backend", "backend_info", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index ad45ff0..52ab855 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -9,6 +9,7 @@ import numpy as np from numpy import * from cunumpy import algorithms as algorithms +from cunumpy import arguments as arguments from cunumpy import cuda as cuda from cunumpy import kernels as kernels from cunumpy import memory as memory diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 2d0eb56..1ca92e2 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -75,7 +75,6 @@ class is the one definition of the arguments. "CudaStruct", "CudaStructArguments", "CudaStructValue", - "PyccelStructArguments", "ctype_of", "cuda_include_dir", "cuda_kernel_names", @@ -1760,143 +1759,6 @@ def _is_device_array(value: Any) -> bool: ) -def _host_state(value: Any) -> Any: - """What a host argument object depends on, to detect replaced values.""" - if _is_device_array(value) or isinstance(value, np.ndarray): - ptr = getattr(getattr(value, "data", None), "ptr", None) - if ptr is None: - ptr = getattr(getattr(value, "ctypes", None), "data", None) - return (id(value), ptr, tuple(getattr(value, "shape", ()))) - if isinstance(value, (bool, int, float, complex, str, np.generic)): - return (type(value), value) - return ("id", id(value)) - - -class PyccelStructArguments(CudaStructArguments): - """Argument object with a pyccel host class and a C struct for CUDA kernels. - - The :class:`~cunumpy.kernels.KernelArguments` form of :class:`CudaStructArguments`: - the same object is passed to a :class:`~cunumpy.kernels.Kernel` on both backends. - On the device path it arrives as the packed struct (``__cuda_args__()``); - on the host path, ``__host_args__()`` builds an instance of - :attr:`host_class` (typically a pyccel-compiled argument class, which - cannot inherit from anything) from the attributes named in - :attr:`host_fields`, once, and again when one of them was replaced. - - A subclass sets :attr:`struct_name`, :attr:`fields` and :attr:`host_class`, - stores every field as an attribute (NumPy or CuPy arrays, whatever the - owner has) and, on the CuPy backend, calls :meth:`pack` at the end of its - constructor so that invalid arrays raise there (:meth:`has_device_arrays` - tells). Objects holding host arrays are copied and pickled without - packing; the struct is only built from device arrays. - - On the CuPy backend there is no host form: the arrays are device arrays, - and a host kernel would have to copy them. ``__host_args__()`` raises - then, unless :attr:`host_copies` is True, in which case the host object is - built from host copies (counted by :func:`~cunumpy.profiling.count_transfers`) and - what the host kernel writes is **not** copied back; use it for read-only - evaluations only. - - Attributes - ---------- - host_class : type - The class of the host argument object, e.g. the pyccel class. - host_fields : Sequence[str] | None - The attributes passed to ``host_class(...)``, positionally and in this - order; by default the struct fields in declaration order. - host_copies : bool - Whether ``__host_args__()`` may copy device arrays to the host - (default False). - - Examples - -------- - >>> from my_kernels import pusher_args_kernels # pyccel-compiled module - >>> class MarkerArguments(PyccelStructArguments): - ... struct_name = "MarkerArgs" - ... fields = (("markers", "Array2D"), ("Np", "long long")) - ... host_class = pusher_args_kernels.MarkerArguments - ... - ... def __init__(self, markers, Np): - ... self.markers = markers - ... self.Np = Np - ... if xp.is_gpu(markers): - ... self.pack() - >>> push(MarkerArguments(markers, Np), dt, n_threads=markers.shape[0]) # doctest: +SKIP - """ - - host_class: type | None = None - host_fields: Sequence[str] | None = None - host_copies: bool = False - - def _host_field_names(self) -> tuple[str, ...]: - if self.host_fields is not None: - return tuple(self.host_fields) - struct = getattr(type(self), "struct", None) - if struct is None: - raise TypeError( - f"{type(self).__qualname__} does not define struct_name and fields", - ) - return tuple(field.name for field in struct.fields) - - def __host_args__(self) -> Any: - """The host argument object, built from the current attributes.""" - host_class = self.host_class - if host_class is None: - raise TypeError( - f"{type(self).__qualname__}.host_class is not set: the class of " - "the host argument object (e.g. the pyccel class) is required", - ) - names = self._host_field_names() - values = [getattr(self, name) for name in names] - state = tuple(_host_state(value) for value in values) - if ( - self.__dict__.get("_host_value") is not None - and self.__dict__.get("_host_state") == state - ): - return self._host_value - host_values = [] - for name, value in zip(names, values): - if _is_device_array(value): - if not self.host_copies: - raise RuntimeError( - f"{type(self).__qualname__}.{name} is a device array: there " - "is no host form on the CuPy backend. Call the CUDA kernel, " - "or set host_copies = True for a read-only host evaluation " - "from host copies", - ) - from cunumpy.xp import to_numpy - - value = to_numpy(value) - host_values.append(value) - self._host_value = host_class(*host_values) - self._host_state = state - return self._host_value - - def has_device_arrays(self) -> bool: - """Whether any array field is a device array (then the struct can be packed).""" - struct = getattr(type(self), "struct", None) - if struct is None: - return False - return any( - _is_device_array(getattr(self, field.name, None)) - for field in struct.fields - if field.pointer or field.view_ndim is not None - ) - - def __getstate__(self) -> dict[str, Any]: - # the host object is rebuilt from the copied or restored attributes - state = super().__getstate__() - state.pop("_host_value", None) - state.pop("_host_state", None) - return state - - def __setstate__(self, state: dict[str, Any]) -> None: - # on the NumPy backend the fields are host arrays: nothing to pack - self.__dict__.update(state) - if self.has_device_arrays(): - self.pack() - - def _array_shapes_in(args: Sequence[Any]) -> Iterator[tuple[int, ...]]: """Array shapes in argument order, including supported argument objects.""" for arg in args: diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 2dfc41b..ed3e388 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -2,7 +2,7 @@ A :class:`Kernel` holds a host kernel (a :class:`~cunumpy.kernels.PyccelKernel`, e.g. a Pyccel-compiled function) and, optionally, its 1:1 corresponding CUDA kernel -(:class:`~cunumpy.cuda.CudaKernel`). It calls the host kernel on the NumPy backend and +(:class:`~cunumpy.kernels.CudaKernel`). It calls the host kernel on the NumPy backend and the CUDA kernel on the CuPy backend (or, with ``dispatch="arrays"``, the CUDA kernel for device arguments and the host kernel for host arguments), so a code base can port its kernels to CUDA one by one. @@ -40,7 +40,6 @@ HostImplementations, PyccelKernel, get_device_kernel_implementation, - resolve_host_args, ) from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _record @@ -58,17 +57,10 @@ def _on_device(arg: Any) -> bool: """Whether a kernel argument lives on the GPU. - A CuPy array, or a device-only argument object (one with ``__cuda_args__`` - but no ``__host_args__``, e.g. a :class:`~cunumpy.cuda.CudaArguments` or a struct - value). A :class:`~cunumpy.kernels.KernelArguments` object has both forms and does - not decide. + A CuPy array, or a CUDA argument object (one with ``__cuda_args__``, e.g. a + :class:`~cunumpy.arguments.CudaArguments` or a struct value). """ - if _is_device_array(arg): - return True - kind = type(arg) - return callable(getattr(kind, "__cuda_args__", None)) and not callable( - getattr(kind, "__host_args__", None), - ) + return _is_device_array(arg) or callable(getattr(type(arg), "__cuda_args__", None)) #: Longest name a Fortran compiler accepts. pyccel names the wrapper module of a @@ -179,7 +171,7 @@ class Kernel: the CUDA kernel on the CuPy backend, the host kernel on the NumPy backend. ``"arrays"``: the CUDA kernel if any top-level argument lives on the GPU (a CuPy array, or a device-only argument object such as a - :class:`~cunumpy.cuda.CudaArguments` or a struct value), else the host + :class:`~cunumpy.arguments.CudaArguments` or a struct value), else the host kernel, whatever the backend. Use ``"arrays"`` when a code deliberately hands host arrays to kernels while CuPy is active (diagnostics, MPI staging, CPU fallbacks): the host kernel then runs on the host arrays @@ -189,11 +181,11 @@ class Kernel: ----- Both kernels take the same arguments, except that the CUDA kernel gets the launch shape (``n_threads`` or ``grid``) and argument objects in their CUDA - form (see :class:`~cunumpy.cuda.CudaArguments` and :class:`~cunumpy.cuda.CudaStruct`). - An argument object implementing :class:`~cunumpy.kernels.KernelArguments` is - replaced by its ``__host_args__()`` on the host path and flattened via - ``__cuda_args__()`` on the CUDA path, so the call site is the same on both - backends. + form (see :class:`~cunumpy.arguments.CudaArguments` and :class:`~cunumpy.arguments.CudaStruct`). + Argument objects are passed as they are: the caller passes the host + argument object (e.g. a pyccel class) on the host and the CUDA one (a + :class:`~cunumpy.arguments.CudaStructArguments`) on the device; see + :doc:`/kernels/arguments`. """ def __init__( @@ -602,12 +594,10 @@ def __call__( Parameters ---------- *args - Kernel arguments. Objects implementing - :class:`~cunumpy.kernels.KernelArguments` are resolved per backend (see - :func:`~cunumpy.kernels.resolve_host_args`). + Kernel arguments, as the selected kernel takes them. n_threads, grid, block, shared_mem, stream Launch configuration of the CUDA kernel, see - :meth:`CudaKernel.__call__ `; + :meth:`CudaKernel.__call__ `; Thread counts are inferred from array shapes by default; `n_threads` or `grid` overrides that choice. Ignored by the host kernel. @@ -625,7 +615,6 @@ def __call__( f"Kernel {self._name!r} has no CUDA kernel: host kernel " "called on the CuPy backend", ) - args, _ = resolve_host_args(args) if ( self._dispatch == "arrays" and not on_device @@ -698,7 +687,7 @@ def from_package( is the CUDA kernel (with a ``__global__`` function ````). Other ``__global__`` functions in that file are ignored by the catalog; load them with :meth:`CudaKernel.all_from_file - `. + `. Parameters ---------- diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index f8d189d..6f84e4a 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -13,8 +13,8 @@ like CuPy does; mixing CuPy and NumPy arrays in arithmetic raises; * reductions and scalar indexing return 0-d arrays, not Python scalars; * arrays have ``data.ptr``, ``device`` and ``__cuda_array_interface__``, so - :class:`~cunumpy.cuda.CudaStruct` packing and the argument checks of - :class:`~cunumpy.cuda.CudaKernel` work; + :class:`~cunumpy.arguments.CudaStruct` packing and the argument checks of + :class:`~cunumpy.kernels.CudaKernel` work; * CUDA kernels cannot run: ``RawKernel`` and friends raise ``NotImplementedError`` when called, and :func:`cunumpy.kernel_testing.requires_cupy` skips tests while the fake is active. diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index ba34a33..a87c2de 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -11,13 +11,6 @@ Ordinary Python callables are supported, including in Pyodide. This module neither imports Pyccel nor compiles kernels; compilation, if desired, is the caller's responsibility. - -Argument objects that have a host and a device form implement the -:class:`KernelArguments` protocol: ``__host_args__()`` returns what the host -kernel receives in that position, ``__cuda_args__()`` what a -:class:`~cunumpy.cuda.CudaKernel` receives. :class:`PyccelKernel` and -:class:`~cunumpy.kernels.Kernel` resolve ``__host_args__()`` with -:func:`resolve_host_args` before calling the host kernel. """ from __future__ import annotations @@ -38,123 +31,7 @@ from cunumpy._transfers import _describe, _is_device_copy, _nbytes, _record from cunumpy.xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy -__all__ = ["CompiledHostKernel", "KernelArguments", "PyccelKernel", "resolve_host_args"] - - -class KernelArguments: - """Base class for argument objects with a host form and a device form. - - A kernel argument that stands for a group of arrays, e.g. all marker data - of a particle species, usually exists twice: as an object holding NumPy - arrays for the host kernel (e.g. a Pyccel class) and as device arrays for - the CUDA kernel. This protocol lets one object represent both, so that a - call to a :class:`~cunumpy.kernels.Kernel` never branches on the backend: - - * ``__host_args__()`` returns the single object that the host kernel - receives in that position; :class:`PyccelKernel` and - :class:`~cunumpy.kernels.Kernel` resolve it (top-level positional and keyword - arguments only) before calling the host kernel; - * ``__cuda_args__()`` returns the tuple of CUDA kernel arguments the object - stands for (the :class:`~cunumpy.cuda.CudaArguments` protocol); - :class:`~cunumpy.cuda.CudaKernel` flattens it into the kernel parameters. - - Subclassing is optional: any object whose *type* defines a callable - ``__host_args__`` is resolved (an instance attribute of that name is not). - Both methods of this base class raise ``NotImplementedError``; a subclass - overrides the ones it supports. - - Examples - -------- - Build each form lazily on first access, so that a CPU run never builds - device arguments and a GPU run never builds the host object: - - >>> class Particles: - ... def __init__(self, markers): - ... self.markers = markers # NumPy or CuPy array - ... self._kernel_args = None - ... - ... @property - ... def kernel_args(self): - ... if self._kernel_args is None: - ... self._kernel_args = ParticleArguments(self) - ... return self._kernel_args - ... - >>> class ParticleArguments(xp.kernels.KernelArguments): - ... def __init__(self, particles): - ... self._particles = particles - ... self._host = None - ... self._cuda = None - ... - ... def __host_args__(self): - ... if self._host is None: # e.g. a Pyccel class of NumPy arrays - ... self._host = MarkerArguments(self._particles.markers) - ... return self._host - ... - ... def __cuda_args__(self): - ... if self._cuda is None: # device arrays and scalars, flattened - ... markers = self._particles.markers - ... self._cuda = (markers, markers.shape[0], markers.shape[1]) - ... return self._cuda - ... - >>> push(particles.kernel_args, dt, n_threads=n) # doctest: +SKIP - - ``push`` receives ``MarkerArguments(...)`` on the NumPy backend and - ``(markers, n_rows, n_cols)`` on the CuPy backend. Invalidate the cached - forms (set them to ``None``) whenever the arrays are replaced, e.g. after - resizing, ``deepcopy`` or unpickling. - """ - - def __host_args__(self) -> Any: - """The object passed to the host kernel in place of this one.""" - raise NotImplementedError( - f"{type(self).__name__} does not provide host kernel arguments " - "(implement __host_args__)", - ) - - def __cuda_args__(self) -> tuple[Any, ...]: - """The CUDA kernel arguments this object stands for.""" - raise NotImplementedError( - f"{type(self).__name__} does not provide CUDA kernel arguments " - "(implement __cuda_args__)", - ) - - -def _host_args(value: Any) -> Any: - """`value.__host_args__()` if its type defines the method, else `value`.""" - # checked on the class, so that an instance attribute of that name (e.g. a - # stored object or None) is not mistaken for the protocol - if callable(getattr(type(value), "__host_args__", None)): - return value.__host_args__() - return value - - -def resolve_host_args( - args: Sequence[Any], - kwargs: Mapping[str, Any] | None = None, -) -> tuple[tuple[Any, ...], dict[str, Any]]: - """Replace :class:`KernelArguments` objects by their host form. - - Every top-level positional and keyword argument whose type defines a - callable ``__host_args__()`` is replaced by the value it returns; all other - arguments (including objects nested in tuples, lists or dicts) are passed - through untouched. - - Parameters - ---------- - args : Sequence - Positional arguments. - kwargs : Mapping[str, Any] | None - Keyword arguments. - - Returns - ------- - tuple[tuple, dict] - The resolved positional and keyword arguments. - """ - return ( - tuple(_host_args(arg) for arg in args), - {name: _host_args(value) for name, value in (kwargs or {}).items()}, - ) +__all__ = ["CompiledHostKernel", "PyccelKernel"] # The conversions between device and host arrays, as module attributes so that @@ -202,13 +79,6 @@ class PyccelKernel: none of its arguments. By default (``None``) every converted array is copied back, which is always correct but does more work. - Notes - ----- - Top-level arguments implementing :class:`KernelArguments` (a - ``__host_args__()`` method on their type) are replaced by their host form - before anything else, so the same argument objects can be passed to a - ``PyccelKernel`` and to a :class:`~cunumpy.cuda.CudaKernel`. - Examples -------- >>> from cunumpy.kernels import PyccelKernel @@ -460,7 +330,6 @@ def _needs_conversion(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bo return any(self._contains_cupy(value) for value in (*args, *kwargs.values())) def __call__(self, *args: Any, **kwargs: Any) -> Any: - args, kwargs = resolve_host_args(args, kwargs) if not self._needs_conversion(args, kwargs): return self._kernel(*args, **kwargs) diff --git a/src/cunumpy/arguments.py b/src/cunumpy/arguments.py new file mode 100644 index 0000000..5ef2de1 --- /dev/null +++ b/src/cunumpy/arguments.py @@ -0,0 +1,42 @@ +"""Argument objects for CUDA kernels: flattened groups and C structs. + +A host kernel (e.g. compiled with Pyccel) and its CUDA version each take their +own argument objects. Write the host argument class, and a +:class:`CudaStructArguments` with the same constructor and attributes for the +CUDA kernels; the code that owns the arrays builds the one for its backend. +Kernels receive argument objects as they are; cunumpy never converts one form +into the other:: + + import cunumpy as xp + + + class CudaMarkerArguments(xp.arguments.CudaStructArguments): + struct_name = "MarkerArgs" + fields = (("markers", "Array2D"), ("n_markers", "int")) + + def __init__(self, markers, n_markers): + self.markers = markers + self.n_markers = n_markers + self.pack() + +:class:`CudaArguments` flattens a group into several kernel parameters, +:class:`CudaStruct` defines a C struct passed by value, and +:func:`write_cuda_header` writes struct definitions to a header. The kernel +classes are in :mod:`cunumpy.kernels`. +""" + +from cunumpy._cuda_kernel import ( + CudaArguments, + CudaStruct, + CudaStructArguments, + CudaStructValue, + write_cuda_header, +) + +__all__ = [ + "CudaArguments", + "CudaStruct", + "CudaStructArguments", + "CudaStructValue", + "write_cuda_header", +] diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 412d548..95a0667 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -1,45 +1,42 @@ -"""CUDA kernels, device helpers, and reusable stream/event interfaces. +"""CUDA device runtime, reusable streams/events, and CUDA source tools. -The kernel classes -(:class:`CudaKernel`, :class:`CudaStruct`, ...) need CuPy to launch; the device -functions (:func:`set_device`, :func:`memory_info`, :func:`stream`, ...) do -nothing (or return ``None``/``0``) on the NumPy backend:: +The device functions (:func:`set_device`, :func:`memory_info`, :func:`stream`, +...) do nothing (or return ``None``/``0``) on the NumPy backend:: import cunumpy as xp - kernel = xp.cuda.CudaKernel(source, "push") + kernel = xp.kernels.CudaKernel(source, "push") with xp.cuda.stream(): kernel(positions, velocities, dt, n_threads=n) The CUDA headers shipped with cunumpy (``cunumpy/atomic.cuh``, -``cunumpy/random.cuh``, ...) are in :func:`cuda_include_dir`. +``cunumpy/random.cuh``, ...) are in :func:`cuda_include_dir`; +:func:`parse_cuda_signature` and the other source tools inspect CUDA sources. Reusable :func:`create_stream` and :func:`create_event` return synchronous host equivalents on NumPy, supporting the same recording and completion interface. -Backend-neutral kernel tools (:class:`~cunumpy.kernels.Kernel`, -:class:`~cunumpy.kernels.PyccelKernel`, ...) are in :mod:`cunumpy.kernels`. +The kernel classes (:class:`~cunumpy.kernels.CudaKernel`, +:class:`~cunumpy.kernels.PyccelKernel`, :class:`~cunumpy.kernels.Kernel`, ...) +are in :mod:`cunumpy.kernels`, the argument objects +(:class:`~cunumpy.arguments.CudaStructArguments`, ...) in :mod:`cunumpy.arguments`. Importing this module makes ``xp.cuda`` refer to it instead of ``cupy.cuda``; use ``import cupy; cupy.cuda`` for CuPy's module. """ +import importlib as _importlib +import warnings as _warnings + from cunumpy._cuda_kernel import ( DEBUG_OPTIONS, - CudaArguments, - CudaKernel, - CudaKernelVariants, CudaParameter, - CudaStruct, - CudaStructArguments, - CudaStructValue, ctype_of, cuda_include_dir, cuda_kernel_names, include_hash, parse_cuda_signature, resolve_includes, - write_cuda_header, ) from cunumpy._device import ( DEFAULT_SHARED_MEMORY_PER_BLOCK, @@ -68,13 +65,7 @@ __all__ = [ "DEBUG_OPTIONS", "DEFAULT_SHARED_MEMORY_PER_BLOCK", - "CudaArguments", - "CudaKernel", - "CudaKernelVariants", "CudaParameter", - "CudaStruct", - "CudaStructArguments", - "CudaStructValue", "HostEvent", "HostStream", "bind_local_device", @@ -99,5 +90,29 @@ "set_device_for_rank", "stream", "wait_event", - "write_cuda_header", ] + +# Names that were in this module before they moved to cunumpy.kernels and +# cunumpy.arguments. They still resolve (with a DeprecationWarning) until cunumpy 0.6. +_MOVED = { + "CudaKernel": "kernels", + "CudaKernelVariants": "kernels", + "CudaArguments": "arguments", + "CudaStruct": "arguments", + "CudaStructArguments": "arguments", + "CudaStructValue": "arguments", + "write_cuda_header": "arguments", +} + + +def __getattr__(name: str): + submodule = _MOVED.get(name) + if submodule is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + _warnings.warn( + f"cunumpy.cuda.{name} moved to cunumpy.{submodule}.{name}; the old name " + "is deprecated and will be removed in cunumpy 0.6", + DeprecationWarning, + stacklevel=2, + ) + return getattr(_importlib.import_module(f"cunumpy.{submodule}"), name) diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index 516e70d..02cd61d 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -1,8 +1,8 @@ """Deprecated alias of :mod:`cunumpy._cuda_kernel`, removed in cunumpy 0.6. -The public names are in :mod:`cunumpy.cuda`. +The public names are in :mod:`cunumpy.kernels` and :mod:`cunumpy.arguments`. """ from cunumpy._deprecated import alias_module -alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.cuda") +alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.kernels") diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 684b128..c26f3f5 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -201,7 +201,7 @@ def _collect_arrays( argument object (e.g. a ``CudaArguments`` object), are named ``"argument []"`` or ``"argument ."``, and arrays in a container attribute of an object ``"argument .[]"``. A - :class:`~cunumpy.cuda.CudaStructArguments` object or a struct value is read + :class:`~cunumpy.arguments.CudaStructArguments` object or a struct value is read through its struct fields, ``"argument ."``, so that its arrays get the names of the attributes of the host argument object it mirrors, also when the fields are properties. @@ -291,7 +291,7 @@ def assert_kernels_agree( arguments only. n_threads, grid, block Launch configuration of the CUDA kernel, see - :meth:`CudaKernel.__call__ `. `n_threads` + :meth:`CudaKernel.__call__ `. `n_threads` may also be a function of the tuple of arguments, e.g. ``lambda args: args[0].shape[0]``. Omitted sizes use the CUDA kernel's shape-based default or its configured ``n_threads_from``; explicit sizes @@ -548,7 +548,7 @@ def device_function_kernel( Name of the generated output array parameter. **kwargs Passed on to :class:`CudaKernel`, e.g. ``include_dirs``, ``options``, - ``block_size`` or ``structs`` (the :class:`~cunumpy.cuda.CudaStruct` types + ``block_size`` or ``structs`` (the :class:`~cunumpy.arguments.CudaStruct` types of struct parameters, whose definitions `header_source` or the `includes` must provide). @@ -563,7 +563,7 @@ def device_function_kernel( every call (an array shared by all threads); * a struct parameter (``DomainArgs d`` or ``const DomainArgs& d``) is taken by value and passed through unchanged to every call (pass a - :class:`~cunumpy.cuda.CudaStructArguments` object or a packed value); + :class:`~cunumpy.arguments.CudaStructArguments` object or a packed value); * a scalar parameter ``T x`` becomes a device array ``const T* x`` of length ``n``, and thread ``i`` calls the function with ``x[i]``; * the return value of thread ``i`` is stored in ``out[i]``, an array diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 7b3b27c..6d4d644 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -1,7 +1,8 @@ """Kernels that run on either backend: dispatch, host implementations, fusion. A :class:`Kernel` pairs a host implementation (Pyccel, numba, NumPy or plain -Python) with an optional CUDA version and runs the one matching the arrays it +Python, wrapped in a :class:`PyccelKernel`) with an optional CUDA version (a +:class:`CudaKernel`) and runs the one matching the arrays it is given; a :class:`KernelCatalog` loads every kernel of a package:: import cunumpy as xp @@ -16,12 +17,13 @@ Device dispatch can require CUDA with :func:`set_device_kernel_implementation` or temporarily with :func:`use_device_kernel_implementation`. -The CUDA-only classes (:class:`~cunumpy.cuda.CudaKernel`, ...) are in -:mod:`cunumpy.cuda`; the pytest helpers for kernel pairs are in +Argument objects for CUDA kernels (:class:`~cunumpy.arguments.CudaStruct`, ...) +are in :mod:`cunumpy.arguments`, the device runtime (streams, devices, debug +mode) in :mod:`cunumpy.cuda`, and the pytest helpers for kernel pairs in :mod:`cunumpy.kernel_testing`. """ -from cunumpy._cuda_kernel import PyccelStructArguments +from cunumpy._cuda_kernel import CudaKernel, CudaKernelVariants from cunumpy._dispatch import Kernel, KernelCatalog from cunumpy._fusion import fuse from cunumpy._kernel import ( @@ -29,13 +31,11 @@ HOST_IMPLEMENTATIONS, CompiledHostKernel, HostImplementations, - KernelArguments, PyccelKernel, as_kernel_array, get_device_kernel_implementation, get_host_kernel_implementation, kernel_output, - resolve_host_args, set_device_kernel_implementation, set_host_kernel_implementation, use_device_kernel_implementation, @@ -46,18 +46,17 @@ "DEVICE_IMPLEMENTATIONS", "HOST_IMPLEMENTATIONS", "CompiledHostKernel", + "CudaKernel", + "CudaKernelVariants", "HostImplementations", "Kernel", - "KernelArguments", "KernelCatalog", "PyccelKernel", - "PyccelStructArguments", "as_kernel_array", "fuse", "get_device_kernel_implementation", "get_host_kernel_implementation", "kernel_output", - "resolve_host_args", "set_device_kernel_implementation", "set_host_kernel_implementation", "use_device_kernel_implementation", diff --git a/tests/unit/test_automatic_launch.py b/tests/unit/test_automatic_launch.py index 36e73d0..4668a18 100644 --- a/tests/unit/test_automatic_launch.py +++ b/tests/unit/test_automatic_launch.py @@ -6,7 +6,8 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaArguments, CudaKernel, CudaStructArguments, CudaStructValue +from cunumpy.arguments import CudaArguments, CudaStructArguments, CudaStructValue +from cunumpy.kernels import CudaKernel SOURCE = 'extern "C" __global__ void work() {}' diff --git a/tests/unit/test_cuda_collectives.py b/tests/unit/test_cuda_collectives.py index 3abfdc1..899cefb 100644 --- a/tests/unit/test_cuda_collectives.py +++ b/tests/unit/test_cuda_collectives.py @@ -6,8 +6,8 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaKernel from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler +from cunumpy.kernels import CudaKernel requires_cuda = pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 582164e..15eedc6 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1,4 +1,4 @@ -"""Tests for `cunumpy.cuda.CudaKernel` and `cunumpy.cuda.parse_cuda_signature`. +"""Tests for `cunumpy.kernels.CudaKernel` and `cunumpy.cuda.parse_cuda_signature`. Signature parsing and argument checking run everywhere: a small stand-in for a device array (`FakeDeviceArray`) takes the place of CuPy arrays. Launching @@ -15,20 +15,21 @@ import cunumpy as xp import cunumpy._cuda_kernel as cuda_module from cunumpy import as_device_array -from cunumpy.cuda import ( +from cunumpy.arguments import ( CudaArguments, - CudaKernel, - CudaKernelVariants, CudaStruct, CudaStructValue, + write_cuda_header, +) +from cunumpy.cuda import ( ctype_of, cuda_include_dir, cuda_kernel_names, include_hash, parse_cuda_signature, resolve_includes, - write_cuda_header, ) +from cunumpy.kernels import CudaKernel, CudaKernelVariants AXPY = r""" // y = a * x + y @@ -992,7 +993,7 @@ def test_bounds_check_traps_on_gpu(): BOUNDS_TRAP = f""" import cupy as cp -from cunumpy.cuda import CudaKernel +from cunumpy.kernels import CudaKernel kernel = CudaKernel({SCALE_COLUMN!r}, "scale_column", options=("-DCUNUMPY_BOUNDS_CHECK",)) a = cp.ones((4, 3)) @@ -1526,7 +1527,7 @@ def test_debug_on_gpu_compiles_and_synchronizes(): if (i < n) y[((long long)i + 1) << 36] = 1.0; // 512 GB and more past y } ''' -kernel = xp.cuda.CudaKernel(SOURCE, "smash", debug=DEBUG) +kernel = xp.kernels.CudaKernel(SOURCE, "smash", debug=DEBUG) y = cp.zeros(64) try: kernel(y, 64, n_threads=64) @@ -1816,7 +1817,7 @@ def test_as_device_array_checks_ndim(): # --------------------------------------------------------------------------- -class ParticleArguments(xp.cuda.CudaStructArguments): +class ParticleArguments(xp.arguments.CudaStructArguments): """The class form of PARTICLES.""" struct_name = "Particles" @@ -1881,7 +1882,7 @@ def test_struct_arguments_check_their_fields(): def test_struct_arguments_need_every_field_attribute(): - class Incomplete(xp.cuda.CudaStructArguments): + class Incomplete(xp.arguments.CudaStructArguments): struct_name = "Incomplete" fields = (("x", "double*"), ("n", "int")) @@ -1899,16 +1900,16 @@ def __init__(self, x): def test_struct_arguments_class_definition(): with pytest.raises(TypeError, match="must define both struct_name and fields"): - class OnlyName(xp.cuda.CudaStructArguments): + class OnlyName(xp.arguments.CudaStructArguments): struct_name = "OnlyName" with pytest.raises(ValueError, match="unsupported type"): - class BadField(xp.cuda.CudaStructArguments): + class BadField(xp.arguments.CudaStructArguments): struct_name = "BadField" fields = (("a", "Other"),) - class Base(xp.cuda.CudaStructArguments): # intermediate base: no struct + class Base(xp.arguments.CudaStructArguments): # intermediate base: no struct def __init__(self): self.pack() @@ -1950,7 +1951,7 @@ def grow(self, n, ptr): self.markers = FakeDeviceArray(np.float64, ptr=ptr, shape=(n, 4)) -class OwnerArguments(xp.cuda.CudaStructArguments): +class OwnerArguments(xp.arguments.CudaStructArguments): struct_name = "OwnerArgs" fields = (("markers", "Array2D"), ("n_markers", "int")) @@ -2013,7 +2014,7 @@ def test_struct_arguments_as_kernel_arguments(): packed, dt, _, _ = kernel.prepare_args(args, 1, out, size) assert packed is args.packed and type(dt) is np.float64 - class Other(xp.cuda.CudaStructArguments): + class Other(xp.arguments.CudaStructArguments): struct_name = "Other" fields = (("x", "double*"),) diff --git a/tests/unit/test_cuda_launch_limits.py b/tests/unit/test_cuda_launch_limits.py index 4834bc5..cf30e07 100644 --- a/tests/unit/test_cuda_launch_limits.py +++ b/tests/unit/test_cuda_launch_limits.py @@ -9,8 +9,7 @@ import cunumpy as xp from cunumpy import _cuda_kernel as implementation -from cunumpy.cuda import CudaKernel -from cunumpy.kernels import Kernel, KernelCatalog +from cunumpy.kernels import CudaKernel, Kernel, KernelCatalog EMPTY = 'extern "C" __global__ void empty() {}' requires_gpu = pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index cd16a75..f9e5482 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -3,8 +3,8 @@ import numpy as np import pytest -from cunumpy.cuda import CudaKernel from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler +from cunumpy.kernels import CudaKernel pytestmark = pytest.mark.skipif( emulation_compiler() is None, diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index ed88deb..57801a3 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -14,14 +14,8 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaKernel -from cunumpy.kernels import ( - Kernel, - KernelArguments, - KernelCatalog, - PyccelKernel, - resolve_host_args, -) +from cunumpy.arguments import CudaArguments +from cunumpy.kernels import CudaKernel, Kernel, KernelCatalog, PyccelKernel SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { @@ -403,7 +397,7 @@ def test_catalog_kernel_includes_from_the_source_root(tmp_path, monkeypatch): # --------------------------------------------------------------------------- -# KernelArguments: one argument object with a host and a device form +# argument objects: the host object for the host kernel, the CUDA one for CUDA # --------------------------------------------------------------------------- @@ -435,28 +429,13 @@ def __init__(self, x, n): self.n = n -class VectorArguments(KernelArguments): - """Lazy owner of both forms; counts how often each is built.""" +class CudaVector(CudaArguments): + """The CUDA counterpart of `HostVector`: flattened into (x, n).""" def __init__(self, x, n): self.x = x self.n = n - self.host_builds = 0 - self.cuda_builds = 0 - self._host = None - self._cuda = None - - def __host_args__(self): - if self._host is None: - self.host_builds += 1 - self._host = HostVector(self.x, self.n) - return self._host - - def __cuda_args__(self): - if self._cuda is None: - self.cuda_builds += 1 - self._cuda = (self.x, self.n) - return self._cuda + super().__init__(x, n) def scale_vector(vector, factor): @@ -466,83 +445,21 @@ def scale_vector(vector, factor): vector.x[i] *= factor -def test_kernel_arguments_base_class(): - args = KernelArguments() - with pytest.raises(NotImplementedError, match="__host_args__"): - args.__host_args__() - with pytest.raises(NotImplementedError, match="__cuda_args__"): - args.__cuda_args__() - - -def test_resolve_host_args(): - args = VectorArguments(np.ones(3), 3) - plain = object() - received, kwargs = resolve_host_args((args, 2.0, plain), {"v": args, "n": 3}) - assert isinstance(received[0], HostVector) and received[0].x is args.x - assert received[1] == 2.0 and received[2] is plain - assert kwargs["v"] is received[0] and kwargs["n"] == 3 # built once, cached - assert args.host_builds == 1 and args.cuda_builds == 0 - - # only the top level is resolved, and kwargs are optional - received, kwargs = resolve_host_args(([args], (args,))) - assert received[0][0] is args and received[1][0] is args and kwargs == {} - - -def test_non_callable_host_args_attribute_is_not_the_protocol(): - """Like `__cuda_args__`, the protocol is looked up on the type.""" - - class Holder: - def __init__(self): - self.__host_args__ = "not a method" - - class ClassAttribute: - __host_args__ = None - - holder, other = Holder(), ClassAttribute() - assert resolve_host_args((holder, other)) == ((holder, other), {}) - - kernel = Kernel(lambda h: h) - with xp.use_backend("numpy"): - assert kernel(holder) is holder - assert PyccelKernel(lambda h: h)(other) is other - - -def test_kernel_arguments_reach_host_kernel_on_numpy(): +def test_host_argument_object_reaches_host_kernel_unchanged(): kernel = Kernel(scale_vector, CudaKernel(SCALE_VECTOR_CUDA, "scale_vector")) - args = VectorArguments(np.ones(4), 4) - with xp.use_backend("numpy"): - kernel(args, 3.0) - kernel(args, 2.0, n_threads=4) - assert np.all(args.x == 6.0) - assert args.host_builds == 1 and args.cuda_builds == 0 # lazy and cached - - -def test_kernel_arguments_with_pyccel_kernel(): - kernel = PyccelKernel(scale_vector) - args = VectorArguments(np.ones(4), 4) + vector = HostVector(np.ones(4), 4) with xp.use_backend("numpy"): - kernel(args, 3.0) - kernel(args, factor=2.0) - kernel(vector=args, factor=0.5) - assert np.all(args.x == 3.0) - assert args.host_builds == 1 and args.cuda_builds == 0 + kernel(vector, 3.0) + kernel(vector, 2.0, n_threads=4) + PyccelKernel(scale_vector)(vector, factor=0.5) + assert np.all(vector.x == 3.0) -def test_kernel_arguments_are_flattened_for_cuda_kernel(): - """`CudaKernel.prepare_args` uses `__cuda_args__`, `__host_args__` is not built.""" +def test_cuda_argument_object_is_flattened_for_cuda_kernel(): x = FakeDeviceArray(np.float64) - args = VectorArguments(x, 7) kernel = CudaKernel(SCALE_VECTOR_CUDA, "scale_vector") - x_out, n, factor = kernel.prepare_args(args, 2.0) + x_out, n, factor = kernel.prepare_args(CudaVector(x, 7), 2.0) assert x_out is x and n == 7 and type(n) is np.int32 and factor == 2.0 - assert args.cuda_builds == 1 and args.host_builds == 0 - - class HostOnly(KernelArguments): - def __host_args__(self): - return HostVector(x, 7) - - with pytest.raises(NotImplementedError, match="__cuda_args__"): - kernel.prepare_args(HostOnly(), 2.0) def test_objects_without_protocol_pass_through(): @@ -554,34 +471,12 @@ def test_objects_without_protocol_pass_through(): assert received[0] is holder and received[1][0] is holder and received[2] == 1 -def test_kernel_arguments_on_cupy(): - """On the CuPy backend the same object is flattened via `__cuda_args__`.""" +def test_cuda_argument_object_on_cupy(): _skip_without_cupy() import cupy as cp kernel = Kernel(scale_vector, CudaKernel(SCALE_VECTOR_CUDA, "scale_vector")) - args = VectorArguments(cp.ones(300), 300) + vector = CudaVector(cp.ones(300), 300) with xp.use_backend("cupy"): - kernel(args, 3.0, n_threads=300) - assert cp.all(args.x == 3.0) - assert args.cuda_builds == 1 and args.host_builds == 0 - - -def test_kernel_arguments_in_fallback_on_cupy(): - """`missing_cuda="fallback"` resolves `__host_args__` through PyccelKernel.""" - _skip_without_cupy() - import cupy as cp - - class DeviceVectorArguments(VectorArguments): - def __host_args__(self): # the host form holds host copies - if self._host is None: - self.host_builds += 1 - self._host = HostVector(xp.to_numpy(self.x), self.n) - return self._host - - kernel = Kernel(scale_vector, missing_cuda="fallback") - args = DeviceVectorArguments(cp.ones(4), 4) - with xp.use_backend("cupy"), pytest.warns(RuntimeWarning): - kernel(args, 3.0) - assert np.all(args.__host_args__().x == 3.0) - assert args.host_builds == 1 and args.cuda_builds == 0 + kernel(vector, 3.0, n_threads=300) + assert cp.all(vector.x == 3.0) diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index 7e74d07..a04f14b 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -12,12 +12,12 @@ import cunumpy as xp from cunumpy import _dispatch as dispatch_module -from cunumpy.cuda import CudaArguments, CudaKernel +from cunumpy.arguments import CudaArguments from cunumpy.kernels import ( CompiledHostKernel, + CudaKernel, HostImplementations, Kernel, - KernelArguments, KernelCatalog, ) @@ -137,12 +137,17 @@ class Device(CudaArguments): kernel(Device(np.ones(1)), 2.0, 1, n_threads=1) assert len(fake_gpu) == 1 - class Both(KernelArguments): # host and device form: does not decide - def __host_args__(self): - return np.ones(2) + class HostArgs: # e.g. a pyccel argument class: no __cuda_args__ + def __init__(self): + self.x = np.ones(2) - host = Kernel(lambda x, f, n: x.__setitem__(slice(None), f), dispatch="arrays") - host(Both(), 5.0, 2) # no device argument: the host kernel + def fill(args, value, n): + args.x[:n] = value + + host_args = HostArgs() + host = Kernel(fill, CudaKernel(SCALE_CUDA, "scale"), dispatch="arrays") + host(host_args, 5.0, 2) # no device argument: the host kernel + assert len(fake_gpu) == 1 and host_args.x.tolist() == [5.0, 5.0] def test_arrays_dispatch_without_cuda_kernel(fake_gpu): diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index c5147ef..ed01e9c 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -15,7 +15,8 @@ import cunumpy as xp import cunumpy.kernel_testing -from cunumpy.cuda import CudaArguments, CudaKernel, parse_cuda_signature +from cunumpy.arguments import CudaArguments +from cunumpy.cuda import parse_cuda_signature from cunumpy.kernel_testing import ( BACKENDS, _collect_arrays, @@ -25,7 +26,7 @@ device_function_kernel, requires_cupy, ) -from cunumpy.kernels import Kernel +from cunumpy.kernels import CudaKernel, Kernel SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { @@ -329,7 +330,7 @@ def __init__(self, owner): self.weights = owner.weights self.n = 3 - class DeviceArguments(xp.cuda.CudaStructArguments): + class DeviceArguments(xp.arguments.CudaStructArguments): struct_name = "OwnerArgs" fields = (("markers", "Array2D"), ("weights", "double*"), ("n", "int")) @@ -355,7 +356,7 @@ def weights(self): assert device["argument 1.markers"] is owner.markers struct = DeviceArguments.struct - value = xp.cuda.CudaStructValue( + value = xp.arguments.CudaStructValue( struct, np.zeros((), struct.dtype)[()], vars(owner) | {"n": 3}, diff --git a/tests/unit/test_mirror.py b/tests/unit/test_mirror.py index 4b67b6e..f9ef252 100644 --- a/tests/unit/test_mirror.py +++ b/tests/unit/test_mirror.py @@ -11,7 +11,7 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaKernel +from cunumpy.kernels import CudaKernel from cunumpy.memory import DeviceMirror requires_gpu = pytest.mark.skipif( diff --git a/tests/unit/test_morton.py b/tests/unit/test_morton.py index 89f3003..bb8282e 100644 --- a/tests/unit/test_morton.py +++ b/tests/unit/test_morton.py @@ -6,8 +6,9 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaKernel, cuda_include_dir +from cunumpy.cuda import cuda_include_dir from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler +from cunumpy.kernels import CudaKernel def interleave(cells, levels): diff --git a/tests/unit/test_namespaces.py b/tests/unit/test_namespaces.py index 28245af..2f679aa 100644 --- a/tests/unit/test_namespaces.py +++ b/tests/unit/test_namespaces.py @@ -13,6 +13,7 @@ SUBMODULES = ( "algorithms", + "arguments", "cuda", "kernels", "memory", @@ -52,6 +53,28 @@ def test_moved_names_warn_and_resolve(name): assert value is getattr(getattr(xp, submodule), name) +def test_kernel_classes_are_in_kernels_and_arguments(): + import cunumpy._cuda_kernel as impl + + assert xp.kernels.CudaKernel is impl.CudaKernel + assert xp.kernels.PyccelKernel is not None + assert xp.arguments.CudaStructArguments is impl.CudaStructArguments + assert not set(xp.cuda._MOVED) & set(xp.cuda.__all__) + # the pre-0.5 top-level names point to the new homes + assert xp._MOVED["CudaKernel"] == "kernels" + assert xp._MOVED["CudaStruct"] == "arguments" + + +@pytest.mark.parametrize("name", sorted(xp.cuda._MOVED)) +def test_names_moved_out_of_cuda_warn_and_resolve(name): + submodule = xp.cuda._MOVED[name] + with pytest.warns(DeprecationWarning, match=f"cunumpy.{submodule}.{name}"): + value = getattr(xp.cuda, name) + assert value is getattr(getattr(xp, submodule), name) + with pytest.raises(AttributeError): + _ = xp.cuda.no_such_name + + def test_numpy_names_do_not_warn(): with warnings.catch_warnings(), xp.use_backend("numpy"): warnings.simplefilter("error") diff --git a/tests/unit/test_philox.py b/tests/unit/test_philox.py index b9c13a3..9e322f2 100644 --- a/tests/unit/test_philox.py +++ b/tests/unit/test_philox.py @@ -6,8 +6,9 @@ import pytest import cunumpy as xp -from cunumpy.cuda import CudaKernel, cuda_include_dir +from cunumpy.cuda import cuda_include_dir from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler +from cunumpy.kernels import CudaKernel # Known-answer vectors of Philox4x32-10 (Random123, kat_vectors) KAT = [ diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index 20be042..c6ac552 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -10,7 +10,6 @@ """ import os -import pickle import subprocess import sys import textwrap @@ -25,9 +24,9 @@ import cunumpy._cuda_kernel as cuda_kernel_module import cunumpy.kernel_testing from cunumpy._dispatch import FORTRAN_NAME_LIMIT, _pyccel_stub_parameters -from cunumpy.cuda import CudaKernel, CudaStruct +from cunumpy.arguments import CudaStruct from cunumpy.kernel_testing import check_parity, device_function_kernel, parity_cases -from cunumpy.kernels import Kernel, KernelCatalog, PyccelStructArguments +from cunumpy.kernels import CudaKernel, Kernel, KernelCatalog class FakeDeviceArray: @@ -184,91 +183,6 @@ def test_from_pyccel_class_errors(): ) -# --------------------------------------------------------------------------- -# PyccelStructArguments -# --------------------------------------------------------------------------- - - -class HostMarkers: - """Stands in for a pyccel-compiled argument class.""" - - instances = 0 - - def __init__(self, markers, Np): - type(self).instances += 1 - self.markers = markers - self.Np = Np - - -class MarkerArguments(PyccelStructArguments): - struct_name = "MarkerArgs" - fields = (("markers", "Array2D"), ("Np", "long long"), ("n_markers", "int")) - host_class = HostMarkers - host_fields = ("markers", "Np") - - def __init__(self, markers, Np): - self.markers = markers - self.Np = Np - self.n_markers = markers.shape[0] - - -def test_pyccel_struct_arguments_host_form_is_built_once_and_follows_changes(): - HostMarkers.instances = 0 - markers = np.zeros((5, 3)) - args = MarkerArguments(markers, 5) - host = args.__host_args__() - assert isinstance(host, HostMarkers) and host.markers is markers and host.Np == 5 - assert args.__host_args__() is host and HostMarkers.instances == 1 - args.Np = 6 # a changed scalar: rebuilt - assert args.__host_args__().Np == 6 and HostMarkers.instances == 2 - args.markers = np.zeros((7, 3)) # a replaced array: rebuilt - assert args.__host_args__().markers is args.markers and HostMarkers.instances == 3 - # the host form is the Kernel's host argument - seen = {} - - def push(m, dt): - seen["host"] = m - - Kernel(push)(args, 0.1) - assert seen["host"] is args.__host_args__() - - -def test_pyccel_struct_arguments_device_form_and_pickling(): - args = MarkerArguments(FakeDeviceArray(np.float64, shape=(5, 3)), 5) - (packed,) = args.__cuda_args__() - assert packed.dtype == MarkerArguments.struct.dtype - assert int(packed["n_markers"]) == 5 - with pytest.raises(RuntimeError, match="no host form on the CuPy backend"): - args.__host_args__() - host_copy = MarkerArguments(np.ones((2, 3)), 2) - host_copy.__host_args__() - restored = pickle.loads(pickle.dumps(host_copy)) - assert "_host_value" not in restored.__dict__ - assert restored.__host_args__().markers.shape == (2, 3) - - -def test_pyccel_struct_arguments_host_copies(monkeypatch): - class Copying(MarkerArguments): - host_copies = True - - device = FakeDeviceArray(np.float64, shape=(2, 3)) - monkeypatch.setattr("cunumpy.xp.to_numpy", lambda a: np.full((2, 3), 7.0)) - host = Copying(device, 2).__host_args__() - assert isinstance(host.markers, np.ndarray) and host.markers[0, 0] == 7.0 - - -def test_pyccel_struct_arguments_requires_host_class(): - class NoHost(PyccelStructArguments): - struct_name = "NoHost" - fields = (("n", "int"),) - - def __init__(self): - self.n = 1 - - with pytest.raises(TypeError, match="host_class is not set"): - NoHost().__host_args__() - - # --------------------------------------------------------------------------- # n_threads_from and check_finite # --------------------------------------------------------------------------- @@ -332,8 +246,18 @@ def __array__(self, dtype=None, copy=None): kernel(Arr([1.0, np.nan]), 2.0, 2, n_threads=2) +class CudaMarkerArguments(xp.arguments.CudaStructArguments): + struct_name = "MarkerArgs" + fields = (("markers", "Array2D"), ("Np", "long long")) + + def __init__(self, markers, Np): + self.markers = markers + self.Np = Np + self.pack() + + def test_device_arrays_in_struct_arguments(): - args = MarkerArguments(FakeDeviceArray(np.float64, shape=(5, 3)), 5) + args = CudaMarkerArguments(FakeDeviceArray(np.float64, shape=(5, 3)), 5) found = dict( cuda_kernel_module._device_arrays_in((1.0, args, FakeDeviceArray("f8"))), ) @@ -566,8 +490,8 @@ def test_require_version(monkeypatch): import pytest import cunumpy as xp import cunumpy.kernel_testing as testing -from cunumpy.cuda import CudaKernel, CudaStruct -from cunumpy.kernels import KernelArguments +from cunumpy.arguments import CudaStruct +from cunumpy.kernels import CudaKernel assert testing.fake_cupy_active() assert xp.cupy_available() and xp.get_backend() == "cupy", xp.get_backend() From 30211c01df2bc0ee4639c8902764753ee1b88508 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 6 Oct 2026 23:56:12 +0200 Subject: [PATCH 3/8] use maybempi (#71) * use maybempi * Patch _mpi.py * Bump version of maybempi --- CHANGELOG.md | 5 + docs/source/api.md | 5 +- docs/source/installation.md | 2 +- pyproject.toml | 1 + src/cunumpy/LLM_GUIDE.md | 1 + src/cunumpy/_device.py | 2 +- src/cunumpy/_mpi.py | 23 +- src/cunumpy/_mpi_serial.py | 657 ---------------------------- src/cunumpy/mpi.py | 43 +- tests/unit/test_device_binding.py | 2 +- tests/unit/test_mpi_cuda_aware.py | 22 +- tests/unit/test_mpi_serial.py | 294 ++----------- tests/unit/test_particle_recipes.py | 8 +- 13 files changed, 99 insertions(+), 966 deletions(-) delete mode 100644 src/cunumpy/_mpi_serial.py diff --git a/CHANGELOG.md b/CHANGELOG.md index de369be..b41e1c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- The launcher detection and the serial MPI stand-in moved to the new package + [maybempi](https://github.com/max-models/maybempi), a dependency of cunumpy. + `xp.mpi` re-exports them (`get_mpi`, `launched_under_mpi`, `local_rank`, + `SerialMPI`, `SerialComm`, ...), plus the new `xp.mpi.is_serial`. The override + variable is now `MAYBEMPI=1`/`0`; `CUNUMPY_MPI` is no longer read. - `CudaKernel` and `CudaKernelVariants` moved to `cunumpy.kernels`, next to `Kernel` and `PyccelKernel`. The argument classes `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue` and diff --git a/docs/source/api.md b/docs/source/api.md index ad52468..98e70e5 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -480,7 +480,8 @@ Importing `mpi4py.MPI` starts MPI (`MPI_Init`), which takes time and makes every collective cost something even on one process. `launched_under_mpi()` tells, from the environment that `mpirun`/`mpiexec`/`srun` set up and without importing mpi4py, whether the process belongs to an MPI job -(`CUNUMPY_MPI=1`/`0` overrides it). `get_mpi()` returns `mpi4py.MPI` then, and +(`MAYBEMPI=1`/`0` overrides it). These functions and the stand-in come from +[maybempi](https://max-models.github.io/maybempi/) and are re-exported in `xp.mpi`. `get_mpi()` returns `mpi4py.MPI` then, and otherwise the serial stand-in, so that the same code runs with and without MPI: @@ -489,7 +490,7 @@ MPI = xp.mpi.get_mpi() # decided once per process comm = MPI.COMM_WORLD comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM) # nothing to do on one process n_total = comm.allreduce(n_local) # n_local itself -if isinstance(MPI, xp.mpi.SerialMPI): +if xp.mpi.is_serial(MPI): ... # a serial run ``` diff --git a/docs/source/installation.md b/docs/source/installation.md index 9d06516..e9f9248 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -71,7 +71,7 @@ Tests that need a GPU are skipped automatically where CuPy is not functional. | `CUNUMPY_CUDA_DEBUG=1` | enable [CUDA debug mode](kernels/debugging.md) for all kernels | | `CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy` | choose the host kernel implementation (read at import) | | `CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda` | require CUDA for device kernel dispatch (read at import); unset allows the kernel's configured fallback | -| `CUNUMPY_MPI=1` / `0` | require MPI / use serial MPI regardless of launcher detection | +| `MAYBEMPI=1` / `0` | require MPI / use serial MPI regardless of launcher detection (read by [maybempi](https://max-models.github.io/maybempi/)) | | `CUNUMPY_FAKE_CUPY=1` | install the strict CPU stand-in for CuPy for tests | | `CUNUMPY_REQUIRE_CUDA=1` | require a real usable GPU when starting the test suite (CI guard) | diff --git a/pyproject.toml b/pyproject.toml index 1ac81b9..b3bced2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ ] dependencies = [ "array-api-compat", + "maybempi>=0.1.1", "numpy", ] diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 34c26ef..d78f741 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -155,6 +155,7 @@ MPI = ( xp.mpi.get_mpi() ) # mpi4py.MPI under mpirun/srun, else a serial stand-in (no MPI_Init) xp.mpi.launched_under_mpi() # from the launcher env, without importing mpi4py +xp.mpi.is_serial(MPI) # True for the stand-in (re-exported from maybempi; MAYBEMPI=0/1) xp.mpi.mpi_is_cuda_aware(comm) # collective, once at startup; remembered with xp.mpi.mpi_buffer(a) as buf: comm.Send(buf, ...) # host array, CUDA-aware device diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index bb13907..b1f6cc5 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -8,8 +8,8 @@ from typing import Any import array_api_compat.numpy as np +from maybempi import local_rank -from cunumpy._mpi_serial import local_rank from cunumpy.xp import array_backend, cupy_available diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 3804744..8dfaf98 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -10,8 +10,8 @@ import array_api_compat import array_api_compat.numpy as np +from maybempi import get_mpi -from cunumpy._mpi_serial import _LOCAL_RANK_VARIABLES # noqa: F401 - re-exported from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _describe, _nbytes, _record from cunumpy.xp import array_backend, cupy_available, to_numpy @@ -265,18 +265,6 @@ def mpi_buffer( cp.cuda.get_current_stream().synchronize() -def _mpi_module() -> Any: - """Import and return ``mpi4py.MPI``, with a clear error if it is missing.""" - try: - from mpi4py import MPI - except ImportError as e: - raise ImportError( - "mpi4py is required for the CUDA-aware MPI check: install it, or " - "pass a communicator explicitly.", - ) from e - return MPI - - def _device_buffers_in_use() -> bool: """Whether MPI calls of this process would carry device (CuPy) buffers.""" return array_backend.backend == "cupy" and cupy_available() @@ -306,8 +294,9 @@ def mpi_is_cuda_aware(comm: Any = None, *, method: str = "probe") -> bool: Parameters ---------- comm - The communicator to check; ``None`` means ``mpi4py.MPI.COMM_WORLD`` - (``mpi4py`` is imported only then, and only on the CuPy backend). + The communicator to check; ``None`` means ``COMM_WORLD`` of + :func:`~cunumpy.mpi.get_mpi` (mpi4py under an MPI launcher, else the + serial stand-in; only on the CuPy backend). method Only ``"probe"`` is available: each rank sends a tiny device buffer to rank ``(rank + 1) % size`` and receives from ``(rank - 1) % size`` with @@ -341,7 +330,7 @@ def mpi_is_cuda_aware(comm: Any = None, *, method: str = "probe") -> bool: if not _device_buffers_in_use(): return False - MPI = _mpi_module() + MPI = get_mpi() if comm is None: comm = MPI.COMM_WORLD @@ -377,7 +366,7 @@ def require_cuda_aware_mpi(comm: Any = None) -> None: Parameters ---------- comm - The communicator to check; ``None`` means ``mpi4py.MPI.COMM_WORLD``. + The communicator to check; ``None`` means ``get_mpi().COMM_WORLD``. """ if not _device_buffers_in_use(): return diff --git a/src/cunumpy/_mpi_serial.py b/src/cunumpy/_mpi_serial.py deleted file mode 100644 index 855f21f..0000000 --- a/src/cunumpy/_mpi_serial.py +++ /dev/null @@ -1,657 +0,0 @@ -"""MPI only when launched under MPI, and a serial stand-in otherwise (see :mod:`cunumpy.mpi`). - -Importing ``mpi4py.MPI`` calls ``MPI_Init``, which can take close to a second -and makes every collective cost something, even on one process. A plain -``python script.py`` should therefore not touch MPI, even when mpi4py is -installed. :func:`launched_under_mpi` tells, from the environment the launcher -sets up and without importing mpi4py, whether the process was started by -``mpirun``/``mpiexec``/``srun``; :func:`get_mpi` returns ``mpi4py.MPI`` then, -and :class:`SerialMPI` otherwise, so that one code path serves both:: - - MPI = xp.mpi.get_mpi() - comm = MPI.COMM_WORLD - comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM) # a no-op on one process - total = comm.allreduce(local_total) # local_total itself - -:class:`SerialComm` behaves like a communicator of size 1: collectives return -or copy what they would on one rank (never ``None`` in place of a value), and -methods it does not implement raise ``AttributeError`` instead of silently -doing nothing. - -This module does not depend on the rest of cunumpy (arrays are handled by duck -typing), so that it could become a package of its own. -""" - -from __future__ import annotations - -import os -import socket -import sys -import time -import warnings -from types import MappingProxyType -from typing import Any - -# Per-rank variables exported by the process managers behind common launchers. -# Each is set only for processes started *by* a launcher. SLURM_PROCID is -# deliberately absent: it is also set for the batch script of a plain `sbatch` -# job, which is not an MPI launch (`srun` exports the PMI/PMIx variables). -_LAUNCHER_VARIABLES = ( - "OMPI_COMM_WORLD_RANK", # Open MPI (and derivatives) - "PMI_RANK", # MPICH, Intel MPI, MS-MPI, Cray, srun --mpi=pmi2 - "PMIX_RANK", # PMIx: srun --mpi=pmix, Open MPI 5 - "MV2_COMM_WORLD_RANK", # MVAPICH2 - "MPI_LOCALRANKID", # Hydra (mpiexec.hydra) - "ALPS_APP_PE", # Cray ALPS aprun - "PALS_RANKID", # Cray PALS -) - -# Node-local rank of the process, as exported by common MPI launchers. They are -# set before ``MPI_Init``, so the device can be chosen before MPI starts. -_LOCAL_RANK_VARIABLES = ( - "OMPI_COMM_WORLD_LOCAL_RANK", # Open MPI - "MV2_COMM_WORLD_LOCAL_RANK", # MVAPICH2 - "MPI_LOCALRANKID", # Intel MPI, MPICH (Hydra) - "PMI_LOCAL_RANK", # MPICH / PMI - "PALS_LOCAL_RANKID", # Cray PALS - "SLURM_LOCALID", # Slurm (srun) - "LOCAL_RANK", # torchrun and others -) - -#: Environment variable that forces the decision of :func:`launched_under_mpi` -#: (``1``/``true``/``yes``/``on`` or ``0``/``false``/``no``/``off``). -OVERRIDE_VARIABLE = "CUNUMPY_MPI" - -_TRUE = ("1", "true", "yes", "on") -_FALSE = ("0", "false", "no", "off") - - -def _env_flag(name: str) -> bool | None: - """The boolean value of the environment variable `name`, or None if unset/unknown.""" - value = os.environ.get(name) - if value is None: - return None - value = value.strip().lower() - if value in _TRUE: - return True - if value in _FALSE: - return False - return None - - -def local_rank() -> int: - """Rank of this process within its node, from the MPI launcher's environment. - - Reads the node-local rank that common launchers export (Open MPI, MVAPICH2, - Intel MPI/MPICH, PMI, Cray PALS, Slurm, ``LOCAL_RANK``). These variables are - set before ``MPI_Init``, so this works before MPI is initialized, and - without importing ``mpi4py``. Returns 0 if none is set (e.g. a serial run). - """ - for variable in _LOCAL_RANK_VARIABLES: - value = os.environ.get(variable) - if value is None: - continue - try: - return int(value) - except ValueError: - continue - return 0 - - -def launched_under_mpi() -> bool: - """Whether this process was started by an MPI launcher (without importing mpi4py). - - True if a per-rank variable of a common launcher is set (Open MPI, MPICH, - Intel MPI, PMIx/``srun``, MVAPICH2, Hydra, Cray ALPS/PALS), or if mpi4py is - already imported and MPI initialized (then using it costs nothing more). - ``CUNUMPY_MPI=1``/``0`` overrides the detection, e.g. for a launcher whose - variables are not known here. - """ - override = _env_flag(OVERRIDE_VARIABLE) - if override is not None: - return override - if any(variable in os.environ for variable in _LAUNCHER_VARIABLES): - return True - # only look at mpi4py if the application imported it: importing it here - # is what must be avoided - mpi = sys.modules.get("mpi4py.MPI") - if mpi is not None: - try: - return bool(mpi.Is_initialized()) - except AttributeError: - return False - return False - - -_AUTO_MPI: Any = None # the result of get_mpi(None), decided once - - -def get_mpi(use_mpi: bool | None = None) -> Any: - """``mpi4py.MPI`` for an MPI run, else the serial stand-in (a :class:`SerialMPI`). - - Parameters - ---------- - use_mpi : bool | None - ``None`` (the default) decides with :func:`launched_under_mpi`, once - per process; ``True`` imports mpi4py (``ImportError`` if it is not - installed); ``False`` returns the stand-in without importing it. - - Returns - ------- - module or SerialMPI - ``mpi4py.MPI``, or the one :class:`SerialMPI` object, which has the - attributes of the module that a serial run needs. - ``isinstance(MPI, xp.mpi.SerialMPI)`` tells which one it is. - - Warns - ----- - RuntimeWarning - Launched under MPI but mpi4py is not installed: every rank then runs - as if it were alone, with the same rank 0. - """ - global _AUTO_MPI - if use_mpi is True: - from mpi4py import MPI - - return MPI - if use_mpi is False: - return _SERIAL_MPI - if _AUTO_MPI is None: - if launched_under_mpi(): - try: - from mpi4py import MPI - except ImportError: - warnings.warn( - "launched under an MPI launcher, but mpi4py is not installed: " - "every process runs serially as rank 0 of 1 (pip install mpi4py)", - RuntimeWarning, - stacklevel=2, - ) - MPI = _SERIAL_MPI - _AUTO_MPI = MPI - else: - _AUTO_MPI = _SERIAL_MPI - return _AUTO_MPI - - -class _Constant: - """A named placeholder for an MPI constant (an op, a datatype, ``IN_PLACE``, ...).""" - - __slots__ = ("name",) - - def __init__(self, name: str) -> None: - self.name = name - - def __repr__(self) -> str: - return f"SerialMPI.{self.name}" - - -class _Datatype(_Constant): - """Placeholder for an MPI datatype (``isinstance(t, MPI.Datatype)`` holds).""" - - __slots__ = () - - -class _Op(_Constant): - """Placeholder for an MPI reduction operation.""" - - __slots__ = () - - -class _Null(_Constant): - """A null handle (``COMM_NULL``, ``DATATYPE_NULL``, ...): false, like mpi4py's.""" - - __slots__ = () - - def __bool__(self) -> bool: - return False - - -_IN_PLACE = _Constant("IN_PLACE") -_COMM_NULL = _Null("COMM_NULL") - - -def _buffer(spec: Any) -> Any: - """The array of an mpi4py buffer specification (``buf`` or ``[buf, ...]``).""" - if isinstance(spec, (list, tuple)): - return spec[0] - return spec - - -def _displacement(spec: Any) -> int: - """The displacement of rank 0 in a vector buffer spec ``[buf, counts, displs, type]``.""" - if isinstance(spec, (list, tuple)) and len(spec) >= 3: - displs = spec[2] - if displs is not None and not isinstance(displs, _Constant): - return int(displs[0]) - return 0 - - -def _copy(source: Any, target: Any, offset: int = 0) -> None: - """Copy the elements of the array `source` into `target`, starting at `offset`. - - Both are flattened (C order); NumPy and CuPy arrays mix (a device source is - copied to the host with ``.get()`` for a host target). - """ - if source is _IN_PLACE or source is None or target is None: - return - if hasattr(source, "get") and not hasattr(target, "get"): - from cunumpy._transfers import _ACTIVE, _nbytes, _record - - source = source.get() - if _ACTIVE: - _record("to_host", "SerialComm receive from device", nbytes=_nbytes(source)) - if not target.flags.c_contiguous: - raise ValueError("the receive buffer must be C-contiguous") - flat = target.reshape(-1) - source = source.reshape(-1) - if offset + source.size > flat.size: - raise ValueError( - f"receive buffer too small: {flat.size} elements for {source.size} " - f"at offset {offset}", - ) - flat[offset : offset + source.size] = source - if not hasattr(source, "get") and hasattr(target, "get"): - from cunumpy._transfers import _ACTIVE, _nbytes, _record - - if _ACTIVE: - _record("to_device", "SerialComm receive from host", nbytes=_nbytes(source)) - - -def _check_rank(rank: int, what: str) -> None: - if rank not in (0, SerialMPI.PROC_NULL, SerialMPI.ANY_SOURCE): - raise ValueError(f"{what}={rank}: a serial communicator has only rank 0") - - -class SerialRequest: - """A completed request, returned by the non-blocking calls of :class:`SerialComm`.""" - - def __init__(self, result: Any = None) -> None: - self._result = result - - def Wait(self, status: Any = None) -> None: - return None - - def Test(self, status: Any = None) -> bool: - return True - - def wait(self, status: Any = None) -> Any: - return self._result - - def test(self, status: Any = None) -> tuple[bool, Any]: - return True, self._result - - def Free(self) -> None: - return None - - def Cancel(self) -> None: - return None - - @staticmethod - def Waitall(requests: Any, statuses: Any = None) -> None: - return None - - @staticmethod - def waitall(requests: Any, statuses: Any = None) -> list[Any]: - return [request.wait() for request in requests] - - @staticmethod - def Testall(requests: Any, statuses: Any = None) -> bool: - return True - - @staticmethod - def Waitany(requests: Any, status: Any = None) -> int: - return 0 if requests else SerialMPI.UNDEFINED - - -class SerialPrequest(SerialRequest): - """A persistent request (``Send_init``/``Recv_init``), for ``Startall``/``Waitall``.""" - - def Start(self) -> None: - return None - - @staticmethod - def Startall(requests: Any) -> None: - return None - - -class SerialComm: - """A communicator of size 1, with the mpi4py ``Comm`` methods a serial run needs. - - Collectives return (object methods) or copy (buffer methods) what they - would on one rank: ``allreduce(x)`` is ``x``, ``gather(x)`` is ``[x]``, - ``Allreduce(send, recv)`` copies `send` into `recv` (nothing with - ``IN_PLACE``), ``Bcast`` does nothing. Point-to-point calls are only - supported to and from rank 0 itself (``sendrecv``, ``Sendrecv``) or - ``PROC_NULL``. Buffers may be NumPy or CuPy arrays, or mpi4py buffer specs - (``[array, MPI.DOUBLE]``). Other methods raise ``AttributeError``. - """ - - rank = 0 - size = 1 - - def __init__(self, name: str = "COMM_WORLD") -> None: - self._name = name - - def __repr__(self) -> str: - return f"SerialComm({self._name})" - - # ---------------------------------------------------------------- queries - def Get_rank(self) -> int: - return 0 - - def Get_size(self) -> int: - return 1 - - def Get_name(self) -> str: - return self._name - - def Is_inter(self) -> bool: - return False - - def Is_intra(self) -> bool: - return True - - # -------------------------------------------------- communicator creation - def Dup(self, info: Any = None) -> SerialComm: - return SerialComm(self._name) - - Clone = Dup - - def Split(self, color: int = 0, key: int = 0) -> Any: - if color == SerialMPI.UNDEFINED: - return SerialMPI.COMM_NULL - return SerialComm(self._name) - - def Free(self) -> None: - return None - - def Abort(self, errorcode: int = 0) -> None: - raise SystemExit(errorcode) - - # ------------------------------------------------------ synchronization - def Barrier(self) -> None: - return None - - barrier = Barrier - - def Ibarrier(self) -> SerialRequest: - return SerialRequest() - - # ------------------------------------------------- collectives, objects - def bcast(self, obj: Any, root: int = 0) -> Any: - _check_rank(root, "root") - return obj - - def reduce(self, sendobj: Any, op: Any = None, root: int = 0) -> Any: - _check_rank(root, "root") - return sendobj - - def allreduce(self, sendobj: Any, op: Any = None) -> Any: - return sendobj - - def scan(self, sendobj: Any, op: Any = None) -> Any: - return sendobj - - def exscan(self, sendobj: Any, op: Any = None) -> None: - return None # undefined on rank 0, None in mpi4py - - def gather(self, sendobj: Any, root: int = 0) -> list[Any]: - _check_rank(root, "root") - return [sendobj] - - def allgather(self, sendobj: Any) -> list[Any]: - return [sendobj] - - def scatter(self, sendobj: Any, root: int = 0) -> Any: - _check_rank(root, "root") - items = list(sendobj) - if len(items) != 1: - raise ValueError(f"scatter on 1 process needs 1 item, got {len(items)}") - return items[0] - - def alltoall(self, sendobj: Any) -> list[Any]: - items = list(sendobj) - if len(items) != 1: - raise ValueError(f"alltoall on 1 process needs 1 item, got {len(items)}") - return items - - def sendrecv( - self, - sendobj: Any, - dest: int = 0, - sendtag: int = 0, - recvbuf: Any = None, - source: int = 0, - recvtag: int = 0, - status: Any = None, - ) -> Any: - _check_rank(dest, "dest") - _check_rank(source, "source") - if source == SerialMPI.PROC_NULL: - return None - return sendobj if dest != SerialMPI.PROC_NULL else None - - def ibcast(self, obj: Any, root: int = 0) -> SerialRequest: - return SerialRequest(self.bcast(obj, root)) - - def iallreduce(self, sendobj: Any, op: Any = None) -> SerialRequest: - return SerialRequest(sendobj) - - # ------------------------------------------------- collectives, buffers - def Bcast(self, buf: Any, root: int = 0) -> None: - _check_rank(root, "root") - - def Reduce(self, sendbuf: Any, recvbuf: Any, op: Any = None, root: int = 0) -> None: - _check_rank(root, "root") - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Allreduce(self, sendbuf: Any, recvbuf: Any, op: Any = None) -> None: - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Scan(self, sendbuf: Any, recvbuf: Any, op: Any = None) -> None: - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Exscan(self, sendbuf: Any, recvbuf: Any, op: Any = None) -> None: - return None # the receive buffer of rank 0 is undefined - - def Gather(self, sendbuf: Any, recvbuf: Any, root: int = 0) -> None: - _check_rank(root, "root") - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Gatherv(self, sendbuf: Any, recvbuf: Any, root: int = 0) -> None: - _check_rank(root, "root") - _copy(_buffer(sendbuf), _buffer(recvbuf), _displacement(recvbuf)) - - def Allgather(self, sendbuf: Any, recvbuf: Any) -> None: - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Allgatherv(self, sendbuf: Any, recvbuf: Any) -> None: - _copy(_buffer(sendbuf), _buffer(recvbuf), _displacement(recvbuf)) - - def Scatter(self, sendbuf: Any, recvbuf: Any, root: int = 0) -> None: - _check_rank(root, "root") - if recvbuf is not _IN_PLACE: - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Scatterv(self, sendbuf: Any, recvbuf: Any, root: int = 0) -> None: - _check_rank(root, "root") - if recvbuf is _IN_PLACE: - return - source = _buffer(sendbuf).reshape(-1) - start = _displacement(sendbuf) - target = _buffer(recvbuf) - _copy(source[start : start + target.size], target) - - def Alltoall(self, sendbuf: Any, recvbuf: Any) -> None: - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Sendrecv( - self, - sendbuf: Any, - dest: int = 0, - sendtag: int = 0, - recvbuf: Any = None, - source: int = 0, - recvtag: int = 0, - status: Any = None, - ) -> None: - _check_rank(dest, "dest") - _check_rank(source, "source") - if SerialMPI.PROC_NULL in (dest, source): - return - _copy(_buffer(sendbuf), _buffer(recvbuf)) - - def Ibcast(self, buf: Any, root: int = 0) -> SerialRequest: - self.Bcast(buf, root) - return SerialRequest() - - def Iallreduce(self, sendbuf: Any, recvbuf: Any, op: Any = None) -> SerialRequest: - self.Allreduce(sendbuf, recvbuf, op) - return SerialRequest() - - def Iallgather(self, sendbuf: Any, recvbuf: Any) -> SerialRequest: - self.Allgather(sendbuf, recvbuf) - return SerialRequest() - - -_COMM_WORLD = SerialComm("COMM_WORLD") -_COMM_SELF = SerialComm("COMM_SELF") - - -class SerialStatus: - """Stand-in for ``MPI.Status`` (source 0, tag 0).""" - - source = 0 - tag = 0 - error = 0 - - def Get_source(self) -> int: - return 0 - - def Get_tag(self) -> int: - return 0 - - def Get_count(self, datatype: Any = None) -> int: - return 0 - - -class SerialMPI: - """Stand-in for the ``mpi4py.MPI`` module in a serial run (see :func:`get_mpi`). - - :func:`get_mpi` returns one instance; ``isinstance(MPI, SerialMPI)`` tells a - serial run from an MPI one. ``COMM_WORLD`` and ``COMM_SELF`` are :class:`SerialComm` objects; the - reduction operations, datatypes and other constants are placeholders that - :class:`SerialComm` accepts. ``Is_initialized()`` is False: MPI itself is - never started. - """ - - COMM_WORLD = _COMM_WORLD - COMM_SELF = _COMM_SELF - COMM_NULL = _COMM_NULL - Comm = Intracomm = SerialComm - Request = SerialRequest - Status = SerialStatus - - IN_PLACE = _IN_PLACE - BOTTOM = _Constant("BOTTOM") - DATATYPE_NULL = _Null("DATATYPE_NULL") - REQUEST_NULL = _Null("REQUEST_NULL") - OP_NULL = _Null("OP_NULL") - PROC_NULL = -2 - ANY_SOURCE = -1 - ANY_TAG = -1 - ROOT = -3 - UNDEFINED = -32766 - SUCCESS = 0 - - Datatype = _Datatype - Op = _Op - Prequest = SerialPrequest - - # reduction operations - SUM = _Op("SUM") - PROD = _Op("PROD") - MAX = _Op("MAX") - MIN = _Op("MIN") - LAND = _Op("LAND") - LOR = _Op("LOR") - LXOR = _Op("LXOR") - BAND = _Op("BAND") - BOR = _Op("BOR") - BXOR = _Op("BXOR") - MAXLOC = _Op("MAXLOC") - MINLOC = _Op("MINLOC") - REPLACE = _Op("REPLACE") - - # datatypes - BYTE = _Datatype("BYTE") - CHAR = _Datatype("CHAR") - BOOL = _Datatype("BOOL") - C_BOOL = _Datatype("C_BOOL") - INT = _Datatype("INT") - LONG = _Datatype("LONG") - LONG_LONG = _Datatype("LONG_LONG") - UNSIGNED = _Datatype("UNSIGNED") - UNSIGNED_LONG = _Datatype("UNSIGNED_LONG") - INT8_T = _Datatype("INT8_T") - INT16_T = _Datatype("INT16_T") - INT32_T = _Datatype("INT32_T") - INT64_T = _Datatype("INT64_T") - UINT8_T = _Datatype("UINT8_T") - UINT16_T = _Datatype("UINT16_T") - UINT32_T = _Datatype("UINT32_T") - UINT64_T = _Datatype("UINT64_T") - FLOAT = _Datatype("FLOAT") - DOUBLE = _Datatype("DOUBLE") - LONG_DOUBLE = _Datatype("LONG_DOUBLE") - C_FLOAT_COMPLEX = _Datatype("C_FLOAT_COMPLEX") - C_DOUBLE_COMPLEX = _Datatype("C_DOUBLE_COMPLEX") - COMPLEX = _Datatype("COMPLEX") - DOUBLE_COMPLEX = _Datatype("DOUBLE_COMPLEX") - - # NumPy type characters to datatypes, like mpi4py's (private) MPI._typedict - _typedict = MappingProxyType({ - "b": INT8_T, "h": INT16_T, "i": INT32_T, "l": LONG, "q": INT64_T, - "B": UINT8_T, "H": UINT16_T, "I": UINT32_T, "L": UNSIGNED_LONG, "Q": UINT64_T, - "f": FLOAT, "d": DOUBLE, "g": LONG_DOUBLE, "?": C_BOOL, - "F": C_FLOAT_COMPLEX, "D": C_DOUBLE_COMPLEX, - }) # fmt: skip - - def __repr__(self) -> str: - return "" - - @staticmethod - def Wtime() -> float: - return time.time() - - @staticmethod - def Wtick() -> float: - return time.get_clock_info("time").resolution - - @staticmethod - def Is_initialized() -> bool: - return False - - @staticmethod - def Is_finalized() -> bool: - return False - - @staticmethod - def Init() -> None: - return None - - @staticmethod - def Finalize() -> None: - return None - - @staticmethod - def Get_processor_name() -> str: - return socket.gethostname() - - @staticmethod - def Query_thread() -> int: - return 0 - - -_SERIAL_MPI = SerialMPI() diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index f5a82b7..ed9a9e8 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -4,7 +4,9 @@ launcher (:func:`launched_under_mpi`, decided from the environment without importing mpi4py), and :class:`SerialMPI` otherwise: a stand-in whose ``COMM_WORLD`` is a :class:`SerialComm` of size 1, so that the same code runs -serially without starting MPI:: +serially without starting MPI. These come from the `maybempi +`_ package and are re-exported here; +``MAYBEMPI=1``/``0`` overrides the launcher detection:: import cunumpy as xp @@ -21,26 +23,44 @@ mpi4py is imported only by the functions that need it. """ -from cunumpy._mpi import ( - MPIStaging, - get_mpi_cuda_aware, - mpi_buffer, - mpi_is_cuda_aware, - require_cuda_aware_mpi, - set_mpi_cuda_aware, - synchronize_for_mpi, -) -from cunumpy._mpi_serial import ( +from maybempi import ( OVERRIDE_VARIABLE, SerialComm, SerialMPI, SerialRequest, SerialStatus, get_mpi, + is_serial, launched_under_mpi, local_rank, + set_copy_hook, ) +from cunumpy._mpi import ( + MPIStaging, + get_mpi_cuda_aware, + mpi_buffer, + mpi_is_cuda_aware, + require_cuda_aware_mpi, + set_mpi_cuda_aware, + synchronize_for_mpi, +) +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _nbytes, _record + + +def _count_serial_copy(kind: str, array: object) -> None: + """Count the host/device copies of the serial stand-in's buffer collectives.""" + if _COUNTERS: + _record( + kind, + f"SerialComm receive ({kind.replace('_', ' ')})", + nbytes=_nbytes(array), + ) + + +set_copy_hook(_count_serial_copy) + __all__ = [ "OVERRIDE_VARIABLE", "MPIStaging", @@ -50,6 +70,7 @@ "SerialStatus", "get_mpi", "get_mpi_cuda_aware", + "is_serial", "launched_under_mpi", "local_rank", "mpi_buffer", diff --git a/tests/unit/test_device_binding.py b/tests/unit/test_device_binding.py index e229b5a..4396dd9 100644 --- a/tests/unit/test_device_binding.py +++ b/tests/unit/test_device_binding.py @@ -3,9 +3,9 @@ import numpy as np import pytest +from maybempi import LOCAL_RANK_VARIABLES as _LOCAL_RANK_VARIABLES import cunumpy as xp -from cunumpy._mpi import _LOCAL_RANK_VARIABLES @pytest.fixture diff --git a/tests/unit/test_mpi_cuda_aware.py b/tests/unit/test_mpi_cuda_aware.py index 9e2c377..d90726b 100644 --- a/tests/unit/test_mpi_cuda_aware.py +++ b/tests/unit/test_mpi_cuda_aware.py @@ -2,7 +2,6 @@ `require_cuda_aware_mpi`. They need neither MPI nor a GPU: the communicator is a fake and the probe buffers are host arrays. The last test needs both.""" -import sys from types import SimpleNamespace import numpy as np @@ -43,19 +42,21 @@ def allreduce(self, value, op=None): @pytest.fixture -def no_mpi4py(monkeypatch): - """Make any import of mpi4py fail.""" - monkeypatch.setitem(sys.modules, "mpi4py", None) - monkeypatch.setitem(sys.modules, "mpi4py.MPI", None) +def no_mpi(monkeypatch): + """Make any attempt to get the MPI module fail.""" + + def fail(): + raise AssertionError("get_mpi() must not be called") + + monkeypatch.setattr(mpi_module, "get_mpi", fail) @pytest.fixture def fake_mpi(monkeypatch): - """A fake `mpi4py.MPI` module with `COMM_WORLD` and `LAND`.""" + """A fake `get_mpi()` module with `COMM_WORLD` and `LAND`.""" comm = FakeComm() MPI = SimpleNamespace(COMM_WORLD=comm, LAND="LAND") - monkeypatch.setitem(sys.modules, "mpi4py", SimpleNamespace(MPI=MPI)) - monkeypatch.setitem(sys.modules, "mpi4py.MPI", MPI) + monkeypatch.setattr(mpi_module, "get_mpi", lambda: MPI) return comm @@ -70,11 +71,11 @@ def buffers(rank, n=4): monkeypatch.setattr(mpi_module, "_mpi_probe_buffers", buffers) -def test_numpy_backend_returns_false_without_mpi(no_mpi4py): +def test_numpy_backend_returns_false_without_mpi(no_mpi): with xp.use_backend("numpy"): assert xp.mpi.mpi_is_cuda_aware() is False assert xp.mpi.mpi_is_cuda_aware(FakeComm()) is False - assert xp.mpi.require_cuda_aware_mpi() is None # no-op, mpi4py not imported + assert xp.mpi.require_cuda_aware_mpi() is None # no-op, MPI not touched def test_unknown_method(): @@ -126,6 +127,5 @@ def test_require_raises(device_buffers, fake_mpi): def test_probe_on_comm_world(): if not xp.cupy_available(): pytest.skip("CuPy not installed or not functional") - pytest.importorskip("mpi4py") with xp.use_backend("cupy"): assert isinstance(xp.mpi.mpi_is_cuda_aware(), bool) diff --git a/tests/unit/test_mpi_serial.py b/tests/unit/test_mpi_serial.py index f441411..87484a9 100644 --- a/tests/unit/test_mpi_serial.py +++ b/tests/unit/test_mpi_serial.py @@ -1,219 +1,35 @@ -"""Tests for `xp.mpi.get_mpi`, `launched_under_mpi` and the serial stand-in `SerialComm`.""" +"""`xp.mpi` re-exports maybempi's launcher detection and serial stand-in. -import sys -import types +The stand-in itself is tested in maybempi; these tests check the cunumpy side: +the re-exports, transfer counting and NumPy/CuPy buffers. +""" +import maybempi import numpy as np import pytest import cunumpy as xp -from cunumpy import _mpi_serial -from cunumpy.mpi import SerialMPI, get_mpi, launched_under_mpi -MPI = get_mpi(False) +MPI = xp.mpi.get_mpi(False) comm = MPI.COMM_WORLD -@pytest.fixture -def clean_env(monkeypatch): - """No launcher variables, no override, no mpi4py imported, no cached decision.""" - for variable in (*_mpi_serial._LAUNCHER_VARIABLES, _mpi_serial.OVERRIDE_VARIABLE): - monkeypatch.delenv(variable, raising=False) - monkeypatch.delitem(sys.modules, "mpi4py.MPI", raising=False) - monkeypatch.setattr(_mpi_serial, "_AUTO_MPI", None) - return monkeypatch - - -@pytest.fixture -def fake_mpi4py(clean_env): - """A stand-in for an installed mpi4py (without starting MPI).""" - module = types.ModuleType("mpi4py.MPI") - module.Is_initialized = lambda: True - package = types.ModuleType("mpi4py") - package.MPI = module - clean_env.setitem(sys.modules, "mpi4py", package) - clean_env.setitem(sys.modules, "mpi4py.MPI", module) - return module - - -@pytest.fixture -def no_mpi4py(clean_env): - clean_env.setitem(sys.modules, "mpi4py", None) - clean_env.setitem(sys.modules, "mpi4py.MPI", None) - return clean_env - - -# ------------------------------------------------------------ the decision - - -def test_serial_by_default(clean_env): - assert launched_under_mpi() is False - assert isinstance(get_mpi(), SerialMPI) - - -@pytest.mark.parametrize("variable", _mpi_serial._LAUNCHER_VARIABLES) -def test_launcher_variables(clean_env, variable): - clean_env.setenv(variable, "3") - assert launched_under_mpi() is True - - -def test_slurm_batch_script_is_not_an_mpi_launch(clean_env): - clean_env.setenv("SLURM_PROCID", "0") - assert launched_under_mpi() is False - - -@pytest.mark.parametrize( - ("value", "launcher", "expected"), - [ - ("1", False, True), ("on", False, True), ("0", True, False), ("no", True, False), - ("maybe", True, True), ("maybe", False, False), - ], -) # fmt: skip -def test_override(clean_env, value, launcher, expected): - clean_env.setenv("CUNUMPY_MPI", value) - if launcher: - clean_env.setenv("PMI_RANK", "0") - assert launched_under_mpi() is expected - - -def test_initialized_mpi4py_counts_as_mpi(fake_mpi4py): - assert launched_under_mpi() is True - fake_mpi4py.Is_initialized = lambda: False - assert launched_under_mpi() is False - - -def test_get_mpi_under_a_launcher(fake_mpi4py, clean_env): - fake_mpi4py.Is_initialized = lambda: False - clean_env.setenv("OMPI_COMM_WORLD_RANK", "0") - assert get_mpi() is fake_mpi4py - clean_env.delenv("OMPI_COMM_WORLD_RANK") - assert get_mpi() is fake_mpi4py # decided once per process - - -def test_get_mpi_explicit(fake_mpi4py): - assert get_mpi(True) is fake_mpi4py - assert get_mpi(False) is MPI - - -def test_launcher_without_mpi4py_warns(no_mpi4py): - no_mpi4py.setenv("PMI_RANK", "0") - with pytest.warns(RuntimeWarning, match="mpi4py is not installed"): - assert isinstance(get_mpi(), SerialMPI) - - -def test_explicit_mpi_without_mpi4py_raises(no_mpi4py): - with pytest.raises(ImportError): - get_mpi(True) - - -def test_exported(): - assert xp.mpi.get_mpi is get_mpi - assert xp.mpi.local_rank is _mpi_serial.local_rank - assert "get_mpi" not in xp._MOVED # never was at the top level - - -# -------------------------------------------------- the serial communicator - - -def test_queries(): - assert comm.rank == 0 and comm.size == 1 - assert comm.Get_rank() == 0 and comm.Get_size() == 1 - assert isinstance(comm, MPI.Comm) and isinstance(comm, MPI.Intracomm) - assert comm.Is_intra() and not comm.Is_inter() - assert MPI.COMM_SELF.Get_size() == 1 - assert repr(comm) == "SerialComm(COMM_WORLD)" - - -def test_object_collectives_return_the_value(): - value = {"a": 1} - assert comm.bcast(value) is value - assert comm.allreduce(5, op=MPI.SUM) == 5 - assert comm.allreduce(2.5, op=MPI.MAX) == 2.5 - assert comm.reduce(7) == 7 - assert comm.scan(3) == 3 - assert comm.exscan(3) is None # as mpi4py on rank 0 - assert comm.gather(value) == [value] - assert comm.allgather(value) == [value] - assert comm.scatter([value]) is value - assert comm.alltoall([value]) == [value] - assert comm.sendrecv(value, dest=0, source=0) is value - assert comm.ibcast(value).wait() is value - assert comm.iallreduce(4).wait() == 4 - assert comm.Barrier() is None and comm.barrier() is None - - -def test_object_collectives_check_sizes_and_ranks(): - with pytest.raises(ValueError, match="1 item"): - comm.scatter([1, 2]) - with pytest.raises(ValueError, match="only rank 0"): - comm.bcast(1, root=1) - with pytest.raises(ValueError, match="only rank 0"): - comm.sendrecv(1, dest=1) - assert comm.sendrecv(1, dest=MPI.PROC_NULL, source=MPI.PROC_NULL) is None - - -def test_unknown_methods_raise(): - # MockComm returned None for everything, which hid missing support - with pytest.raises(AttributeError): - _ = comm.Create_cart - with pytest.raises(AttributeError): - _ = MPI.Win - - -def test_buffer_collectives_copy(): - send = np.arange(6.0).reshape(2, 3) - for call in ( - lambda r: comm.Allreduce(send, r, op=MPI.SUM), - lambda r: comm.Reduce(send, r, op=MPI.SUM, root=0), - lambda r: comm.Allgather(send, r), - lambda r: comm.Gather(send, r), - lambda r: comm.Scatter(send, r), - lambda r: comm.Alltoall(send, r), - lambda r: comm.Scan(send, r), - lambda r: comm.Sendrecv(send, dest=0, recvbuf=r, source=0), - lambda r: comm.Iallreduce(send, r).Wait(), - lambda r: comm.Iallgather(send, r).Wait(), +def test_reexported_from_maybempi(): + for name in ( + "OVERRIDE_VARIABLE", + "SerialComm", + "SerialMPI", + "SerialRequest", + "SerialStatus", + "get_mpi", + "is_serial", + "launched_under_mpi", + "local_rank", ): - recv = np.zeros_like(send) - call(recv) - np.testing.assert_array_equal(recv, send) - - -def test_buffer_specs_and_in_place(): - data = np.arange(4.0) - recv = np.zeros(4) - comm.Allreduce([data, MPI.DOUBLE], [recv, MPI.DOUBLE], op=MPI.SUM) - np.testing.assert_array_equal(recv, data) - before = data.copy() - comm.Allreduce(MPI.IN_PLACE, data, op=MPI.SUM) - comm.Bcast(data, root=0) - comm.Ibcast(data).Wait() - np.testing.assert_array_equal(data, before) - comm.Exscan(data, recv) # undefined on rank 0: untouched - np.testing.assert_array_equal(recv, before) - - -def test_vector_collectives_use_the_displacement(): - send = np.array([1.0, 2.0]) - recv = np.zeros(5) - comm.Allgatherv(send, [recv, [2], [3], MPI.DOUBLE]) - np.testing.assert_array_equal(recv, [0, 0, 0, 1, 2]) - recv[:] = 0 - comm.Gatherv(send, [recv, [2], [1]]) - np.testing.assert_array_equal(recv, [0, 1, 2, 0, 0]) - part = np.zeros(2) - comm.Scatterv([np.arange(5.0), [2], [2], MPI.DOUBLE], part) - np.testing.assert_array_equal(part, [2, 3]) - - -def test_buffer_errors(): - with pytest.raises(ValueError, match="too small"): - comm.Allreduce(np.ones(3), np.zeros(2)) - with pytest.raises(ValueError, match="C-contiguous"): - comm.Allreduce(np.ones(2), np.zeros((2, 2))[:, 0]) - with pytest.raises(ValueError, match="only rank 0"): - comm.Sendrecv(np.ones(2), dest=1, recvbuf=np.zeros(2)) - comm.Sendrecv(np.ones(2), dest=MPI.PROC_NULL, recvbuf=np.zeros(2)) + assert getattr(xp.mpi, name) is getattr(maybempi, name), name + assert xp.mpi.OVERRIDE_VARIABLE == "MAYBEMPI" + assert xp.mpi.is_serial(MPI) and xp.mpi.is_serial(comm) + assert "get_mpi" not in xp._MOVED # never was at the top level class _DeviceArray: @@ -222,6 +38,7 @@ class _DeviceArray: def __init__(self, data): self._data = np.asarray(data) self.flags = self._data.flags + self.nbytes = self._data.nbytes def get(self): return self._data.copy() @@ -229,52 +46,23 @@ def get(self): def reshape(self, *shape): return _DeviceArray(self._data.reshape(*shape)) + def __setitem__(self, key, value): + self._data[key] = value + @property def size(self): return self._data.size -def test_device_source_to_host_target(): - recv = np.zeros(3) - comm.Allreduce(_DeviceArray([1.0, 2.0, 3.0]), recv) - np.testing.assert_array_equal(recv, [1, 2, 3]) - - -def test_split_dup_and_requests(): - assert comm.Split(0, 0).Get_size() == 1 - assert comm.Split(MPI.UNDEFINED) is MPI.COMM_NULL - assert comm.Dup().Get_rank() == 0 and comm.Clone().Get_size() == 1 - requests = [comm.Ibarrier(), comm.Ibcast(np.zeros(1))] - assert MPI.Request.Waitall(requests) is None - assert MPI.Request.Testall(requests) is True - assert MPI.Request.waitall([comm.ibcast(1)]) == [1] - status = MPI.Status() - assert status.Get_source() == 0 and status.Get_tag() == 0 - - -def test_module_functions(): - assert MPI.Is_initialized() is False and MPI.Is_finalized() is False - assert MPI.Wtime() > 0 and MPI.Wtick() > 0 - assert isinstance(MPI.Get_processor_name(), str) - assert repr(MPI.SUM) == "SerialMPI.SUM" - assert MPI.PROC_NULL == -2 and MPI.ANY_SOURCE == -1 and MPI.ROOT == -3 - assert isinstance(MPI, SerialMPI) and get_mpi(False) is MPI - - -def test_constants_used_by_struphy_and_feectools(): - assert isinstance(MPI.DOUBLE, MPI.Datatype) and not isinstance( - MPI.SUM, - MPI.Datatype, - ) - assert isinstance(MPI.LOR, MPI.Op) - assert MPI._typedict[np.dtype(np.float64).char] is MPI.DOUBLE - # null handles are false, communicators true, like in mpi4py - assert not MPI.COMM_NULL and not MPI.DATATYPE_NULL - assert comm and comm != MPI.COMM_NULL - MPI.Prequest.Startall([]) - assert MPI.Prequest.Waitall([]) is None - # usable in annotations evaluated at definition time - assert (MPI.Intracomm | None) is not None +def test_serial_copies_between_host_and_device_are_counted(): + with xp.profiling.count_transfers() as counter: + comm.Allreduce(_DeviceArray([1.0, 2.0, 3.0]), np.zeros(3)) + comm.Allgather(np.ones(2), _DeviceArray(np.zeros(2))) + comm.Allreduce(np.ones(2), np.zeros(2)) # host to host: not a transfer + assert [(e.kind, e.nbytes) for e in counter.events] == [ + ("to_host", 24), + ("to_device", 16), + ] @pytest.mark.parametrize( @@ -295,17 +83,3 @@ def test_buffers_on_either_backend(backend): host = np.zeros(4) comm.Allgather(send, host) # a device array into a host buffer np.testing.assert_array_equal(host, np.arange(4.0)) - - -def test_matches_mpi4py_on_one_process(): - """The same calls on mpi4py's COMM_SELF give the same results (if mpi4py is there).""" - real = pytest.importorskip("mpi4py.MPI") - self_comm = real.COMM_SELF - assert self_comm.allreduce(5) == comm.allreduce(5) - assert self_comm.gather(3) == comm.gather(3) - assert self_comm.scatter([4]) == comm.scatter([4]) - assert self_comm.exscan(3) == comm.exscan(3) - send, recv_real, recv_serial = np.arange(3.0), np.zeros(3), np.zeros(3) - self_comm.Allreduce(send, recv_real, op=real.SUM) - comm.Allreduce(send, recv_serial, op=MPI.SUM) - np.testing.assert_array_equal(recv_real, recv_serial) diff --git a/tests/unit/test_particle_recipes.py b/tests/unit/test_particle_recipes.py index c447513..72be764 100644 --- a/tests/unit/test_particle_recipes.py +++ b/tests/unit/test_particle_recipes.py @@ -4,7 +4,8 @@ from pathlib import Path import numpy as np -import pytest + +import cunumpy as xp _PATH = Path(__file__).resolve().parents[2] / "docs/source/examples/particle_recipes.py" _spec = importlib.util.spec_from_file_location("particle_recipes", _PATH) @@ -50,10 +51,7 @@ def test_pack_for_ranks(): def test_exchange_on_one_rank(): - mpi4py = pytest.importorskip("mpi4py") - from mpi4py import MPI - - del mpi4py + MPI = xp.mpi.get_mpi() markers = np.arange(8.0).reshape(4, 2) received = recipes.exchange(MPI.COMM_SELF, markers, np.zeros(4, dtype=np.int64)) np.testing.assert_array_equal(received, markers) From f35dd965fea77e61a29ee8f54588f0b4d289d79d Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 6 Oct 2026 23:59:17 +0200 Subject: [PATCH 4/8] Bump version to 0.7.0 (#73) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b3bced2..cdaae1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.6.0" +version = "0.7.0" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] From d60decd2c0183cca6322cdcde493cd1144783d1a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 07:49:21 +0200 Subject: [PATCH 5/8] Prepare 0 7 0 (#74) * Bump version to 0.7.0 * Bumb maybempi to 0.1.2 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index cdaae1c..6f11be9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] dependencies = [ "array-api-compat", - "maybempi>=0.1.1", + "maybempi>=0.1.2", "numpy", ] From dca25227ebac6a536d782726b24c17f23118645f Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 08:09:33 +0200 Subject: [PATCH 6/8] Cleanup of deprecated code (#75) --- CHANGELOG.md | 12 ++++++-- docs/source/api.md | 13 +++++---- src/cunumpy/LLM_GUIDE.md | 6 ++-- src/cunumpy/__init__.py | 55 ++--------------------------------- src/cunumpy/_deprecated.py | 19 ------------ src/cunumpy/cuda/__init__.py | 28 ------------------ src/cunumpy/cuda_kernel.py | 8 ----- src/cunumpy/dispatch.py | 8 ----- src/cunumpy/kernel.py | 8 ----- src/cunumpy/main.py | 8 ----- src/cunumpy/testing.py | 9 ------ tests/unit/test_mpi_serial.py | 1 - tests/unit/test_namespaces.py | 52 ++++----------------------------- 13 files changed, 27 insertions(+), 200 deletions(-) delete mode 100644 src/cunumpy/_deprecated.py delete mode 100644 src/cunumpy/cuda_kernel.py delete mode 100644 src/cunumpy/dispatch.py delete mode 100644 src/cunumpy/kernel.py delete mode 100644 src/cunumpy/main.py delete mode 100644 src/cunumpy/testing.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b41e1c5..3273399 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,8 +18,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `CudaStruct`, `CudaStructArguments`, `CudaStructValue` and `write_cuda_header` moved to the new `cunumpy.arguments`. `cunumpy.cuda` keeps the device runtime and the CUDA source tools. The old - `cunumpy.cuda.` imports still work with a `DeprecationWarning` and - will be removed in cunumpy 0.6. + `cunumpy.cuda.` names are removed. +- **Removed** the deprecated names that were kept for one release after the + 0.5 reorganisation, with no replacement other than the submodules: + the top-level helpers (`xp.CudaKernel`, `xp.mpi_buffer`, `xp.fuse` as cunumpy's, + ...; use `xp.kernels`, `xp.arguments`, `xp.cuda`, `xp.mpi`, `xp.rng`, + `xp.algorithms`, `xp.profiling`, `xp.memory` and `xp.petsc`), and the modules + `cunumpy.testing` (now `cunumpy.kernel_testing`), `cunumpy.kernel`, + `cunumpy.dispatch` and `cunumpy.cuda_kernel` (public names are in + `cunumpy.kernels` and `cunumpy.arguments`). Also removed the placeholder + `cunumpy.main`. `xp.fuse` is now CuPy's own `fuse`. - **Removed** `kernels.KernelArguments`, `kernels.resolve_host_args`, `kernels.PyccelStructArguments` and the `__host_args__()` protocol, with no replacement. Kernels receive argument objects as they are. Write the host diff --git a/docs/source/api.md b/docs/source/api.md index 98e70e5..acfa555 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -54,9 +54,11 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not "Both" means the functions work on NumPy and CuPy arrays; the functions of `cunumpy.cuda` do nothing (or return `None`/`0`) on the NumPy backend. -Before cunumpy 0.5 these names were at the top level (`xp.CudaKernel`). The -old names still work until cunumpy 0.6 and raise a `DeprecationWarning` that -names the new place. +Each name has one import path, through its submodule: `xp.kernels.CudaKernel`, +`xp.arguments.CudaStruct`, `xp.mpi.mpi_buffer`, `xp.rng.random_streams`. The top +level of `cunumpy` is the NumPy/CuPy namespace plus backend selection and array +conversion, so it never hides a NumPy or CuPy name (`xp.fuse` is CuPy's `fuse`; +cunumpy's is `xp.kernels.fuse`). Modules starting with `_` are private. ## Version @@ -1713,9 +1715,8 @@ from cunumpy.kernel_testing import ( `cunumpy.kernel_testing` holds helpers for testing kernels with pytest. It is not imported by `import cunumpy`, and it imports pytest only when one of its pytest objects is used, so `device_function_kernel` works without pytest. -Before cunumpy 0.5 it was called `cunumpy.testing`, which replaced NumPy's -`xp.testing` once imported; that name still works, with a -`DeprecationWarning`, until cunumpy 0.6. +It is not called `cunumpy.testing`, because a submodule of that name would +replace NumPy's `xp.testing` once imported. ### `requires_cupy`, `BACKENDS`, `backend` diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index d78f741..a72563d 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -58,10 +58,8 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. `xp.rng.random_streams`, `xp.algorithms.morton_keys`, `xp.mpi.mpi_buffer`, `xp.profiling.timed_region`, `xp.memory.HostStaging`, `xp.petsc.petsc_vec`; kernel test helpers in - `cunumpy.kernel_testing`. The old top-level names (`xp.CudaKernel`) and - `cunumpy.testing`, and the kernel and argument classes in `xp.cuda` - (`xp.cuda.CudaKernel`), are deprecated (removed in 0.6); do not write new code - with them. Modules starting with `_` (`cunumpy._cuda_kernel`, ...) are + `cunumpy.kernel_testing`. There is no `xp.CudaKernel`, `xp.cuda.CudaKernel` + or `cunumpy.testing`; use the submodule names above. Modules starting with `_` (`cunumpy._cuda_kernel`, ...) are private; never import from them. ## Decision guide diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index b6d8b70..1125dcc 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -1,6 +1,5 @@ # cunumpy/__init__.py import re as _re -import warnings as _warnings from importlib.metadata import PackageNotFoundError, version from cunumpy import ( @@ -36,43 +35,6 @@ use_backend, ) -# Names that were at the top level before cunumpy 0.5, and the submodule each -# moved to. They still resolve (with a DeprecationWarning) until cunumpy 0.6. -_MOVED = { - **dict.fromkeys(cuda.__all__, "cuda"), - **dict.fromkeys(arguments.__all__, "arguments"), - **dict.fromkeys( - ( - name - for name in kernels.__all__ - if not name.endswith( - ("_host_kernel_implementation", "_device_kernel_implementation") - ) - and name != "DEVICE_IMPLEMENTATIONS" - ), - "kernels", - ), - **dict.fromkeys(rng.__all__, "rng"), - **dict.fromkeys(algorithms.__all__, "algorithms"), - # the names of cunumpy.mpi that were at the top level (not the later ones) - **dict.fromkeys( - ( - "get_mpi_cuda_aware", - "local_rank", - "mpi_buffer", - "mpi_is_cuda_aware", - "require_cuda_aware_mpi", - "set_mpi_cuda_aware", - "synchronize_for_mpi", - ), - "mpi", - ), - **dict.fromkeys(profiling.__all__, "profiling"), - **dict.fromkeys(memory.__all__, "memory"), - "petsc_vec": "petsc", -} -_MOVED.pop("BIT_GENERATORS") # never was at the top level - try: __version__ = version("cunumpy") except PackageNotFoundError: @@ -147,31 +109,19 @@ def __getattr__(name: str): The public names of the active backend are copied into this namespace (see `_sync_backend_namespace`), so this only runs for names missing from the - backend's ``__all__``, for ``numpy_backend``/``cupy_backend``, and for the - names moved to a submodule in cunumpy 0.5 (see ``_MOVED``), which still - resolve with a ``DeprecationWarning``. + backend's ``__all__`` and for ``numpy_backend``/``cupy_backend``. """ if name == "numpy_backend": return xp.numpy_backend if name == "cupy_backend": return xp.cupy_backend - submodule = _MOVED.get(name) - if submodule is not None: - _warnings.warn( - f"cunumpy.{name} moved to cunumpy.{submodule}.{name}; the top-level " - "name is deprecated and will be removed in cunumpy 0.6", - DeprecationWarning, - stacklevel=2, - ) - return getattr(globals()[submodule], name) return getattr(xp.xp, name) # `xp.zeros` must be as fast as `numpy.zeros`. A module-level __getattr__ runs # only after the normal lookup failed, which costs about 3 us per access, so the # public names of the active backend module are copied into this namespace, and -# replaced whenever the backend changes. cunumpy's own names and the deprecated -# names of _MOVED (e.g. `fuse`, which CuPy also has) are never overwritten. +# replaced whenever the backend changes. cunumpy's own names are never overwritten. _OWN_NAMES = frozenset(globals()) _backend_names: dict[int, dict[str, object]] = {} # id(module) -> names to copy _switches: dict[tuple[int, int], tuple[tuple[str, ...], dict[str, object]]] = {} @@ -186,7 +136,6 @@ def _names_of(module) -> dict[str, object]: for name in getattr(module, "__all__", ()) if not name.startswith("_") and name not in _OWN_NAMES - and name not in _MOVED and hasattr(module, name) } _backend_names[id(module)] = names diff --git a/src/cunumpy/_deprecated.py b/src/cunumpy/_deprecated.py deleted file mode 100644 index f27d7a7..0000000 --- a/src/cunumpy/_deprecated.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Module aliases kept for one release after a module was renamed.""" - -import importlib -import sys -import warnings - - -def alias_module(old: str, new: str, use: str) -> None: - """Make the module `old` an alias of `new`, with a ``DeprecationWarning``. - - Called from the module `old` itself, which is then replaced in - ``sys.modules`` by `new`; `use` is the module to import instead. - """ - warnings.warn( - f"{old} is deprecated and will be removed in cunumpy 0.6; import {use} instead", - DeprecationWarning, - stacklevel=3, - ) - sys.modules[old] = importlib.import_module(new) diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 95a0667..5a026fc 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -25,9 +25,6 @@ use ``import cupy; cupy.cuda`` for CuPy's module. """ -import importlib as _importlib -import warnings as _warnings - from cunumpy._cuda_kernel import ( DEBUG_OPTIONS, CudaParameter, @@ -91,28 +88,3 @@ "stream", "wait_event", ] - -# Names that were in this module before they moved to cunumpy.kernels and -# cunumpy.arguments. They still resolve (with a DeprecationWarning) until cunumpy 0.6. -_MOVED = { - "CudaKernel": "kernels", - "CudaKernelVariants": "kernels", - "CudaArguments": "arguments", - "CudaStruct": "arguments", - "CudaStructArguments": "arguments", - "CudaStructValue": "arguments", - "write_cuda_header": "arguments", -} - - -def __getattr__(name: str): - submodule = _MOVED.get(name) - if submodule is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - _warnings.warn( - f"cunumpy.cuda.{name} moved to cunumpy.{submodule}.{name}; the old name " - "is deprecated and will be removed in cunumpy 0.6", - DeprecationWarning, - stacklevel=2, - ) - return getattr(_importlib.import_module(f"cunumpy.{submodule}"), name) diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py deleted file mode 100644 index 02cd61d..0000000 --- a/src/cunumpy/cuda_kernel.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Deprecated alias of :mod:`cunumpy._cuda_kernel`, removed in cunumpy 0.6. - -The public names are in :mod:`cunumpy.kernels` and :mod:`cunumpy.arguments`. -""" - -from cunumpy._deprecated import alias_module - -alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.kernels") diff --git a/src/cunumpy/dispatch.py b/src/cunumpy/dispatch.py deleted file mode 100644 index 8781b5b..0000000 --- a/src/cunumpy/dispatch.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Deprecated alias of :mod:`cunumpy._dispatch`, removed in cunumpy 0.6. - -The public names are in :mod:`cunumpy.kernels`. -""" - -from cunumpy._deprecated import alias_module - -alias_module(__name__, "cunumpy._dispatch", "cunumpy.kernels") diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py deleted file mode 100644 index d0b8717..0000000 --- a/src/cunumpy/kernel.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Deprecated alias of :mod:`cunumpy._kernel`, removed in cunumpy 0.6. - -The public names are in :mod:`cunumpy.kernels`. -""" - -from cunumpy._deprecated import alias_module - -alias_module(__name__, "cunumpy._kernel", "cunumpy.kernels") diff --git a/src/cunumpy/main.py b/src/cunumpy/main.py deleted file mode 100644 index 063d2ee..0000000 --- a/src/cunumpy/main.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -Main module of the python package. -""" - - -def main(): - """Main method called from from the command line.""" - print("Hello, world") diff --git a/src/cunumpy/testing.py b/src/cunumpy/testing.py deleted file mode 100644 index 9f3593a..0000000 --- a/src/cunumpy/testing.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Deprecated alias of :mod:`cunumpy.kernel_testing`, removed in cunumpy 0.6. - -A submodule named ``testing`` replaced NumPy's ``xp.testing`` once imported, so -``xp.testing.assert_allclose`` failed after any ``import cunumpy.testing``. -""" - -from cunumpy._deprecated import alias_module - -alias_module(__name__, "cunumpy.kernel_testing", "cunumpy.kernel_testing") diff --git a/tests/unit/test_mpi_serial.py b/tests/unit/test_mpi_serial.py index 87484a9..75a5ddd 100644 --- a/tests/unit/test_mpi_serial.py +++ b/tests/unit/test_mpi_serial.py @@ -29,7 +29,6 @@ def test_reexported_from_maybempi(): assert getattr(xp.mpi, name) is getattr(maybempi, name), name assert xp.mpi.OVERRIDE_VARIABLE == "MAYBEMPI" assert xp.mpi.is_serial(MPI) and xp.mpi.is_serial(comm) - assert "get_mpi" not in xp._MOVED # never was at the top level class _DeviceArray: diff --git a/tests/unit/test_namespaces.py b/tests/unit/test_namespaces.py index 2f679aa..df25611 100644 --- a/tests/unit/test_namespaces.py +++ b/tests/unit/test_namespaces.py @@ -1,4 +1,4 @@ -"""Tests for the submodule layout of cunumpy (0.5) and the deprecated top-level names.""" +"""Tests for the submodule layout of cunumpy and the backend names of the top level.""" import importlib import os @@ -39,40 +39,19 @@ def test_submodule_exports_resolve(name): def test_top_level_is_backend_and_numpy_only(): - moved = set(xp._MOVED) - assert not moved & set(xp.__all__) for name in xp.__all__: assert hasattr(xp, name) -@pytest.mark.parametrize("name", sorted(xp._MOVED)) -def test_moved_names_warn_and_resolve(name): - submodule = xp._MOVED[name] - with pytest.warns(DeprecationWarning, match=f"cunumpy.{submodule}.{name}"): - value = getattr(xp, name) - assert value is getattr(getattr(xp, submodule), name) - - def test_kernel_classes_are_in_kernels_and_arguments(): import cunumpy._cuda_kernel as impl assert xp.kernels.CudaKernel is impl.CudaKernel assert xp.kernels.PyccelKernel is not None assert xp.arguments.CudaStructArguments is impl.CudaStructArguments - assert not set(xp.cuda._MOVED) & set(xp.cuda.__all__) - # the pre-0.5 top-level names point to the new homes - assert xp._MOVED["CudaKernel"] == "kernels" - assert xp._MOVED["CudaStruct"] == "arguments" - - -@pytest.mark.parametrize("name", sorted(xp.cuda._MOVED)) -def test_names_moved_out_of_cuda_warn_and_resolve(name): - submodule = xp.cuda._MOVED[name] - with pytest.warns(DeprecationWarning, match=f"cunumpy.{submodule}.{name}"): - value = getattr(xp.cuda, name) - assert value is getattr(getattr(xp, submodule), name) - with pytest.raises(AttributeError): - _ = xp.cuda.no_such_name + for name in ("CudaKernel", "CudaStruct"): + assert not hasattr(xp.cuda, name) + assert name not in vars(xp) def test_numpy_names_do_not_warn(): @@ -86,6 +65,8 @@ def test_numpy_names_do_not_warn(): def test_unknown_name_raises(): with pytest.raises(AttributeError): _ = xp.no_such_function_in_cunumpy + with pytest.raises(AttributeError): + _ = xp.cuda.no_such_name def test_kernel_testing_keeps_numpy_testing(): @@ -95,25 +76,6 @@ def test_kernel_testing_keeps_numpy_testing(): assert xp.testing is np.testing -def test_testing_alias_is_deprecated(): - # a fresh process: importing cunumpy.testing rebinds xp.testing for the rest of it - code = ( - "import warnings\n" - "warnings.simplefilter('error')\n" - "try:\n" - " import cunumpy.testing\n" - "except DeprecationWarning as w:\n" - " assert 'cunumpy.kernel_testing' in str(w)\n" - "else:\n" - " raise SystemExit('no warning')\n" - "warnings.simplefilter('ignore')\n" - "import cunumpy.testing, cunumpy.kernel_testing\n" - "assert cunumpy.testing.assert_kernels_agree is " - "cunumpy.kernel_testing.assert_kernels_agree\n" - ) - subprocess.run([sys.executable, "-c", code], check=True) - - def test_backend_names_are_plain_attributes(): # copied into the namespace: no module __getattr__ call per access assert vars(xp)["zeros"] is xp.xp.xp.zeros @@ -125,8 +87,6 @@ def test_backend_names_are_plain_attributes(): def test_backend_names_never_hide_cunumpy_names(): assert xp.cuda is importlib.import_module("cunumpy.cuda") assert xp.scipy is importlib.import_module("cunumpy._scipy_backend").scipy - for name in xp._MOVED: - assert name not in vars(xp), name def test_switching_the_backend_replaces_the_names(): From ad60a634037e36f33528b941525a447b5ef70b39 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 08:40:02 +0200 Subject: [PATCH 7/8] Cleanup (#76) --- src/cunumpy/_cuda_kernel.py | 6 +----- tests/unit/test_cuda_kernel.py | 4 ++-- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 1ca92e2..1c691d1 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -326,10 +326,6 @@ def _includes(source: str) -> list[tuple[str, bool]]: ] -def _quoted_includes(source: str) -> list[str]: - return [name for name, quoted in _includes(source) if quoted] - - def resolve_includes( source: str, include_dirs: Iterable[str | Path] = (), @@ -949,7 +945,7 @@ def _header_source( if any(f.view_ndim is not None for s in structs for f in s.fields): includes.insert(0, _ARRAY_VIEW_INCLUDE) lines = [ - "// Generated by cunumpy.cuda.CudaStruct from the Python definition; do not edit.", + "// Generated by cunumpy.arguments.CudaStruct from the Python definition; do not edit.", f"#ifndef {guard}", f"#define {guard}", "", diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 15eedc6..011fcf2 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1111,7 +1111,7 @@ def test_to_header(tmp_path): struct = CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") header = struct.to_header() assert header == ( - "// Generated by cunumpy.cuda.CudaStruct from the Python definition; do not edit.\n" + "// Generated by cunumpy.arguments.CudaStruct from the Python definition; do not edit.\n" "#ifndef MARKERARGS_CUH\n" "#define MARKERARGS_CUH\n" "\n" @@ -1151,7 +1151,7 @@ def test_write_cuda_header(tmp_path): ) assert path.read_text() == header assert header.startswith( - "// Generated by cunumpy.cuda.CudaStruct from the Python definition; do not edit.\n" + "// Generated by cunumpy.arguments.CudaStruct from the Python definition; do not edit.\n" "#ifndef PUSHER_ARGS_CUH\n#define PUSHER_ARGS_CUH\n\n" '#include "cunumpy/array_view.cuh"\n#include \n\n', ) From 144c40f317f6d709957b92157c76019f008922e4 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 10:18:52 +0200 Subject: [PATCH 8/8] Set version number to 0.6.0 (#77) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6f11be9..329ab9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.7.0" +version = "0.6.0" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ]