diff --git a/CHANGELOG.md b/CHANGELOG.md index 94dd615..7022e86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,33 @@ 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 + `write_cuda_header` moved to the new `cunumpy.arguments`. `cunumpy.cuda` + keeps the device runtime and the CUDA source tools. The old + `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 + 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. @@ -17,6 +44,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 The former function names are removed without compatibility aliases. ### Added +- Strided and C-contiguous CUDA array views through 16 dimensions, including + `Array5D`/`Array6D` for matrix accumulations. Higher-dimensional views work + as kernel parameters, struct fields, annotations and in CPU emulation. +- C-contiguous array views `CArray1D` to `CArray16D` 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/README.md b/README.md index 92d709d..21e5c34 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,35 +380,21 @@ 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 shapes and strides rather than bare pointers. The shipped header `cunumpy/array_view.cuh` (found by every `CudaKernel`) provides the strided -views `Array1D` to `Array4D`; a parameter or struct field of that type +views `Array1D` to `Array16D`; a parameter or struct field of that type takes a CuPy array, contiguous or not, and indexes `a(i, j)`. The struct can be generated from the annotations of the pyccel argument class, so the Python class is the one definition, and written to a header that a test keeps in sync: @@ -417,11 +404,13 @@ 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..06971ac 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` | @@ -53,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 @@ -365,7 +368,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") @@ -479,7 +482,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: @@ -488,7 +492,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 ``` @@ -853,12 +857,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 +877,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 +904,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 +945,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 +1008,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 +1041,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 +1107,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]) ``` @@ -1113,8 +1117,8 @@ cunumpy ships CUDA headers that every `CudaKernel` finds automatically; (`-I`). `cunumpy/array_view.cuh` defines the strided views `Array1D` to -`Array4D` (4D e.g. for a 3D grid of vector components `(nx, ny, nz, -ncomp)`): `T* data`, `long long shape[ndim]`, `long long +`Array16D` (for example, a 4D view can describe a 3D grid of vector +components `(nx, ny, nz, ncomp)`): `T* data`, `long long shape[ndim]`, `long long strides[ndim]` (in elements, not bytes), `operator()(i, j, ...)` returning a reference to the element, and `size()`. A kernel indexes `a(i, j)` like the pyccel kernel it is ported from indexes `a[i, j]`, without hand-passed sizes. @@ -1149,7 +1153,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 +1162,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 +1182,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 +1218,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 +1235,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 +1251,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) ``` @@ -1257,7 +1261,7 @@ changes one definition instead of every kernel signature. `CudaStruct(name, fields)` takes the fields as `(name, C type)` pairs; scalar fields, pointers to the scalar types above (or `void*`), and array views -`Array1D` to `Array4D` of those scalar types (see "CUDA headers and +`Array1D` to `Array16D` of those scalar types (see "CUDA headers and array views") are supported. * `declaration`: the C definition of the struct, to put in the CUDA source @@ -1296,7 +1300,9 @@ 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 +1339,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 +1355,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 +1364,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 +1378,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 +1408,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 +1423,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 +1455,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 +1476,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, @@ -1803,9 +1717,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` @@ -1986,7 +1899,7 @@ and the kernel's include directories and `-D` options apply. Then the kernel is called once per thread, for every block and thread index of the launch shape. Arguments follow the signature, with NumPy arrays in place of CuPy arrays: -pointer and view parameters (`Array1D` to `Array4D`) take arrays of the +pointer and view parameters (`Array1D` to `Array16D`) take arrays of the declared dtype (and ndim), passed as contiguous copies, so any strides work, and written back into the given arrays; scalars are checked and cast like in a launch. Like NVRTC by default, the compiler may fuse `a * b + c` into an FMA, 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/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/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 5e4e80f..6c53e39 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,14 +73,14 @@ 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) ``` Fields may be scalars, pointers to scalar types (or `void*`), and array views -`Array1D` to `Array4D`. Packing checks every field like a kernel +`Array1D` to `Array16D`. Packing checks every field like a kernel argument: pointers need C-contiguous CuPy arrays of the declared dtype, scalars are range-checked and cast. Adding a field means editing the one Python definition; kernels that use the struct pick it up. @@ -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,9 @@ 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 +219,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" ) ``` @@ -325,13 +237,47 @@ 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 +`CArray16D` 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.arguments.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 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` @@ -343,19 +289,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..e0b0949 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) @@ -54,7 +54,7 @@ every call: | `void*` | C-contiguous CuPy array of any dtype | host arrays, views | | `double`, `float`, `complex` | Python `int`/`float`, NumPy scalars that cast safely | strings, arrays, unsafe casts (`np.float64` into `float`) | | `int`, `long long`, `size_t`, `int64_t`, ... | Python `int` and NumPy integers whose value is in range, `bool` | out-of-range values (`OverflowError`), floats | -| `Array1D` ... `Array4D` | CuPy array of dtype `T` and that ndim, contiguous or not | wrong dtype or ndim | +| `Array1D` ... `Array16D` | CuPy array of dtype `T` and that ndim, contiguous or not | wrong dtype or ndim | | a `CudaStruct` type | a value of that struct | anything else | C types map to NumPy dtypes as on 64-bit Linux: `int` is `int32`, `long` and @@ -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,9 @@ 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 +164,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) ``` @@ -177,7 +179,7 @@ Pass extra include directories with `include_dirs=[...]` and NVRTC flags with | Header | Provides | | --- | --- | | `` | `CUNUMPY_THREAD_1D(i, n)`, `_2D`, `_3D`, `CUNUMPY_GRID_STRIDE_1D(i, n)` | -| `` | strided views `Array1D` to `Array4D` | +| `` | strided views `Array1D` to `Array16D` | | `` | `cunumpy_atomic_add` and indexed 2D/3D variants, see [Accumulation kernels](accumulation.md) | | `` | Morton (Z-order) keys `cunumpy_morton_key2(x, y, ...)`, `_key3`, equal to `xp.algorithms.morton_keys` on the host | | `` | counter-based random numbers `cunumpy_uniform(seed, stream, counter)`, `cunumpy_normal2(...)`, equal to `xp.rng.philox_uniform` on the host | @@ -210,7 +212,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 +249,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 +259,12 @@ 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..d4c331e 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 ``` @@ -28,7 +28,7 @@ In debug mode a `CudaKernel`: * is compiled with `-lineinfo` (source lines for `compute-sanitizer` and profilers) and `-DCUNUMPY_BOUNDS_CHECK`, which turns on bounds checks in - `Array1D` to `Array4D` views (an out-of-bounds index prints the index + `Array1D` to `Array16D` views (an out-of-bounds index prints the index and shape, then traps); * synchronizes after every launch, so a failure raises at the launch that caused it, as a `RuntimeError` naming the kernel and its grid and block, with @@ -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..7dbd7ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.5.0" +version = "0.6.1" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] @@ -23,6 +23,7 @@ classifiers = [ ] dependencies = [ "array-api-compat", + "maybempi>=0.1.2", "numpy", ] diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 630a4e3..65a6a8b 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,12 +54,12 @@ 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 - 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 @@ -71,7 +71,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,14 +89,14 @@ 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 | | reduction inside a CUDA kernel (energy, max velocity) | ``: `cunumpy_block_sum_to(out, v)`, `cunumpy_block_min/max`, `cunumpy_warp_sum` | | kernel writes into a host buffer owned by another library | `xp.memory.DeviceMirror(host_array)` + `` | -| N-D indexing in CUDA, non-contiguous arrays | `Array1D`..`Array4D` params from `` | +| N-D indexing in CUDA, non-contiguous arrays | `Array1D`..`Array16D` params from `` | | one MPI rank per GPU | `bind_local_device()` → `from mpi4py import MPI` → `require_cuda_aware_mpi()` → `synchronize_for_mpi(...)` before each call | | timing GPU code | `with xp.profiling.timed_region("name") as t:` → `t.elapsed` | | profiler markers | `xp.profiling.nvtx_range("name")` (context manager or decorator) | @@ -153,6 +153,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 @@ -263,17 +264,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() @@ -306,7 +307,7 @@ Shipped CUDA headers (always on the include path): ```c #include // CUNUMPY_THREAD_1D(i, n) /_2D/_3D, CUNUMPY_GRID_STRIDE_1D(i, n) {...} -#include // Array1D..Array4D: data, shape[], strides[] (elements), a(i, j), size() +#include // Array1D..Array16D: data, shape[], strides[] (elements), a(i, j), size() #include // cunumpy_atomic_add(double*|float*, v), _2d(data, n1, i, j, v), _3d(...) ``` @@ -339,20 +340,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 +353,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 +368,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 +404,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 +438,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 +488,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..1125dcc 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -1,9 +1,19 @@ # cunumpy/__init__.py import re as _re -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, @@ -25,42 +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( - ( - 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: @@ -97,6 +71,7 @@ def require_version(minimum: str) -> None: __all__ = [ "__version__", "algorithms", + "arguments", "as_device_array", "assert_same_backend", "backend_info", @@ -134,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]]] = {} @@ -173,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/__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 c3f931f..fafd890 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`. @@ -73,7 +75,6 @@ class is the one definition of the arguments. "CudaStruct", "CudaStructArguments", "CudaStructValue", - "PyccelStructArguments", "ctype_of", "cuda_include_dir", "cuda_kernel_names", @@ -122,7 +123,7 @@ def cuda_include_dir() -> str: :class:`CudaKernel` adds it to the include path automatically, so kernels can ``#include "cunumpy/array_view.cuh"`` (strided ``Array1D`` - to ``Array4D`` views passed by value) and + to ``Array16D`` views passed by value) and ``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such as ``CUNUMPY_THREAD_1D(i, n)``), ``#include "cunumpy/atomic.cuh"`` (atomic adds) and ``#include "cunumpy/reduce.cuh"`` (warp and block reductions). @@ -187,7 +188,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. + ``Array16D`` or ``CArray1D`` to ``CArray16D``, 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 +201,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 +265,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 Array16D and CArray1D to CArray16D (cunumpy/array_view.cuh), +# T a scalar type of _CTYPES +_VIEW = re.compile(r"\b(C?)Array([1-9]|1[0-6])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>") _TOKEN = re.compile( - r"Array[1234]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>" + r"C?Array(?:[1-9]|1[0-6])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: @@ -320,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] = (), @@ -515,14 +517,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 +717,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 +731,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 +854,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 @@ -886,11 +905,32 @@ def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str ctype = scalars[scalar] if ndim == 0: return ctype - if ndim > 4: - raise ValueError(f"{what}: arrays have at most 4 dimensions, got {ndim}") + if ndim > 16: + raise ValueError(f"{what}: arrays have at most 16 dimensions, got {ndim}") 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}" @@ -905,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}", "", @@ -1035,7 +1075,7 @@ class CudaStruct: ``(field name, C type)`` pairs, in order, e.g. ``("x", "double*")`` or ``("n", "int")``. Scalar fields, pointers to the scalar types of :func:`ctype_of` (or ``void*``), and array views ``Array1D`` to - ``Array4D`` of those scalar types (from ``cunumpy/array_view.cuh``, + ``Array16D`` of those scalar types (from ``cunumpy/array_view.cuh``, packed as pointer, shape and strides in elements) are supported. Examples @@ -1084,6 +1124,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 +1150,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 +1186,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 +1200,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 +1224,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 +1265,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: @@ -1706,143 +1755,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: @@ -1936,7 +1848,7 @@ class CudaKernel: The headers shipped with cunumpy (:func:`cuda_include_dir`) are always found at compile time (see :meth:`compile_options`): ``#include "cunumpy/array_view.cuh"`` gives the ``Array1D`` to - ``Array4D`` views, ``#include "cunumpy/index.cuh"`` the thread-index + ``Array16D`` views, ``#include "cunumpy/index.cuh"`` the thread-index macros, ``#include "cunumpy/atomic.cuh"`` atomic adds, ``#include "cunumpy/reduce.cuh"`` warp and block reductions. source_dir : str | Path | None 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/_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/_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/_emulation.py b/src/cunumpy/_emulation.py index 499b888..56d47d0 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -16,7 +16,7 @@ np.testing.assert_allclose(y, 2.0 * x) Arguments follow the kernel signature: NumPy arrays for pointer and array view -parameters (``Array1D`` to ``Array4D``; any strides, they are passed as +parameters (``Array1D`` to ``Array16D``; any strides, they are passed as contiguous copies), Python or NumPy scalars for scalar parameters (cast and checked like in a launch). Arrays are written back into the given arrays. @@ -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\d+D<(.*)>", 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/_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/_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/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..5a026fc 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -1,24 +1,25 @@ -"""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. @@ -26,20 +27,13 @@ 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 +62,7 @@ __all__ = [ "DEBUG_OPTIONS", "DEFAULT_SHARED_MEMORY_PER_BLOCK", - "CudaArguments", - "CudaKernel", - "CudaKernelVariants", "CudaParameter", - "CudaStruct", - "CudaStructArguments", - "CudaStructValue", "HostEvent", "HostStream", "bind_local_device", @@ -99,5 +87,4 @@ "set_device_for_rank", "stream", "wait_event", - "write_cuda_header", ] diff --git a/src/cunumpy/cuda/include/cunumpy/array_view.cuh b/src/cunumpy/cuda/include/cunumpy/array_view.cuh index c1ebe6a..2a299ae 100644 --- a/src/cunumpy/cuda/include/cunumpy/array_view.cuh +++ b/src/cunumpy/cuda/include/cunumpy/array_view.cuh @@ -1,6 +1,6 @@ -// 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) +// Array1D to Array16D describe a (possibly non-contiguous) // device array the way NumPy/CuPy do: a data pointer, a shape and strides. // Strides are in ELEMENTS, not bytes, so that `a(i, j)` is // `data[i * strides[0] + j * strides[1]]`. Elements are accessed with @@ -9,7 +9,7 @@ // // The views are created on the Python side by cunumpy (a kernel parameter or a // CudaStruct field of type `Array2D` takes a CuPy array). They are -// passed by value, so the memory layout must be exactly, for ndim = 1 to 4: +// passed by value, so the memory layout must be exactly, for ndim = 1 to 16: // // T* data; // 8 bytes // long long shape[ndim]; // ndim * 8 bytes @@ -20,6 +20,19 @@ // the end of this file check the size. Member functions do not change the // layout. // +// CArray1D to CArray16D 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 16: +// +// 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,255 @@ 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}}; + } +}; + +// Higher-dimensional views share an implementation; aliases keep the named +// types used by kernel signatures and Python's argument packer. +namespace cunumpy_detail { +template +struct ArrayView { + T* data; + long long shape[N]; + long long strides[N]; + + template + __device__ __forceinline__ T& operator()(Indices... indices) const { + static_assert(sizeof...(Indices) == N, "wrong number of array indices"); + const long long index[N] = {static_cast(indices)...}; + long long offset = 0; + #pragma unroll + for (int axis = 0; axis < N; ++axis) { + CUNUMPY_CHECK_INDEX(index[axis], axis, shape[axis]); + offset += index[axis] * strides[axis]; + } + return data[offset]; + } + + __device__ __forceinline__ long long size() const { + long long result = 1; + #pragma unroll + for (int axis = 0; axis < N; ++axis) result *= shape[axis]; + return result; + } +}; + +template +struct CArrayView { + T* data; + long long shape[N]; + + template + __device__ __forceinline__ T& operator()(Indices... indices) const { + static_assert(sizeof...(Indices) == N, "wrong number of array indices"); + const long long index[N] = {static_cast(indices)...}; + long long offset = 0; + #pragma unroll + for (int axis = 0; axis < N; ++axis) { + CUNUMPY_CHECK_INDEX(index[axis], axis, shape[axis]); + offset = offset * shape[axis] + index[axis]; + } + return data[offset]; + } + + __device__ __forceinline__ long long size() const { + long long result = 1; + #pragma unroll + for (int axis = 0; axis < N; ++axis) result *= shape[axis]; + return result; + } + + __host__ __device__ operator ArrayView() const { + ArrayView result{}; + result.data = data; + long long stride = 1; + #pragma unroll + for (int axis = N - 1; axis >= 0; --axis) { + result.shape[axis] = shape[axis]; + result.strides[axis] = stride; + stride *= shape[axis]; + } + return result; + } +}; +} // namespace cunumpy_detail + +template using Array5D = cunumpy_detail::ArrayView; +template using CArray5D = cunumpy_detail::CArrayView; +template using Array6D = cunumpy_detail::ArrayView; +template using CArray6D = cunumpy_detail::CArrayView; +template using Array7D = cunumpy_detail::ArrayView; +template using CArray7D = cunumpy_detail::CArrayView; +template using Array8D = cunumpy_detail::ArrayView; +template using CArray8D = cunumpy_detail::CArrayView; +template using Array9D = cunumpy_detail::ArrayView; +template using CArray9D = cunumpy_detail::CArrayView; +template using Array10D = cunumpy_detail::ArrayView; +template using CArray10D = cunumpy_detail::CArrayView; +template using Array11D = cunumpy_detail::ArrayView; +template using CArray11D = cunumpy_detail::CArrayView; +template using Array12D = cunumpy_detail::ArrayView; +template using CArray12D = cunumpy_detail::CArrayView; +template using Array13D = cunumpy_detail::ArrayView; +template using CArray13D = cunumpy_detail::CArrayView; +template using Array14D = cunumpy_detail::ArrayView; +template using CArray14D = cunumpy_detail::CArrayView; +template using Array15D = cunumpy_detail::ArrayView; +template using CArray15D = cunumpy_detail::CArrayView; +template using Array16D = cunumpy_detail::ArrayView; +template using CArray16D = cunumpy_detail::CArrayView; + +// 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"); + +static_assert(sizeof(Array5D) == 88, "unexpected Array5D layout"); +static_assert(alignof(Array5D) == 8, "unexpected Array5D alignment"); +static_assert(sizeof(CArray5D) == 48, "unexpected CArray5D layout"); +static_assert(alignof(CArray5D) == 8, "unexpected CArray5D alignment"); +static_assert(sizeof(Array6D) == 104, "unexpected Array6D layout"); +static_assert(alignof(Array6D) == 8, "unexpected Array6D alignment"); +static_assert(sizeof(CArray6D) == 56, "unexpected CArray6D layout"); +static_assert(alignof(CArray6D) == 8, "unexpected CArray6D alignment"); +static_assert(sizeof(Array7D) == 120, "unexpected Array7D layout"); +static_assert(alignof(Array7D) == 8, "unexpected Array7D alignment"); +static_assert(sizeof(CArray7D) == 64, "unexpected CArray7D layout"); +static_assert(alignof(CArray7D) == 8, "unexpected CArray7D alignment"); +static_assert(sizeof(Array8D) == 136, "unexpected Array8D layout"); +static_assert(alignof(Array8D) == 8, "unexpected Array8D alignment"); +static_assert(sizeof(CArray8D) == 72, "unexpected CArray8D layout"); +static_assert(alignof(CArray8D) == 8, "unexpected CArray8D alignment"); +static_assert(sizeof(Array9D) == 152, "unexpected Array9D layout"); +static_assert(alignof(Array9D) == 8, "unexpected Array9D alignment"); +static_assert(sizeof(CArray9D) == 80, "unexpected CArray9D layout"); +static_assert(alignof(CArray9D) == 8, "unexpected CArray9D alignment"); +static_assert(sizeof(Array10D) == 168, "unexpected Array10D layout"); +static_assert(alignof(Array10D) == 8, "unexpected Array10D alignment"); +static_assert(sizeof(CArray10D) == 88, "unexpected CArray10D layout"); +static_assert(alignof(CArray10D) == 8, "unexpected CArray10D alignment"); +static_assert(sizeof(Array11D) == 184, "unexpected Array11D layout"); +static_assert(alignof(Array11D) == 8, "unexpected Array11D alignment"); +static_assert(sizeof(CArray11D) == 96, "unexpected CArray11D layout"); +static_assert(alignof(CArray11D) == 8, "unexpected CArray11D alignment"); +static_assert(sizeof(Array12D) == 200, "unexpected Array12D layout"); +static_assert(alignof(Array12D) == 8, "unexpected Array12D alignment"); +static_assert(sizeof(CArray12D) == 104, "unexpected CArray12D layout"); +static_assert(alignof(CArray12D) == 8, "unexpected CArray12D alignment"); +static_assert(sizeof(Array13D) == 216, "unexpected Array13D layout"); +static_assert(alignof(Array13D) == 8, "unexpected Array13D alignment"); +static_assert(sizeof(CArray13D) == 112, "unexpected CArray13D layout"); +static_assert(alignof(CArray13D) == 8, "unexpected CArray13D alignment"); +static_assert(sizeof(Array14D) == 232, "unexpected Array14D layout"); +static_assert(alignof(Array14D) == 8, "unexpected Array14D alignment"); +static_assert(sizeof(CArray14D) == 120, "unexpected CArray14D layout"); +static_assert(alignof(CArray14D) == 8, "unexpected CArray14D alignment"); +static_assert(sizeof(Array15D) == 248, "unexpected Array15D layout"); +static_assert(alignof(Array15D) == 8, "unexpected Array15D alignment"); +static_assert(sizeof(CArray15D) == 128, "unexpected CArray15D layout"); +static_assert(alignof(CArray15D) == 8, "unexpected CArray15D alignment"); +static_assert(sizeof(Array16D) == 264, "unexpected Array16D layout"); +static_assert(alignof(Array16D) == 8, "unexpected Array16D alignment"); +static_assert(sizeof(CArray16D) == 136, "unexpected CArray16D layout"); +static_assert(alignof(CArray16D) == 8, "unexpected CArray16D alignment"); #endif // CUNUMPY_ARRAY_VIEW_CUH diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py deleted file mode 100644 index 516e70d..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.cuda`. -""" - -from cunumpy._deprecated import alias_module - -alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.cuda") 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/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/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/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/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_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 1309007..0ef8c8c 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 @@ -704,7 +705,7 @@ def test_parse_view_parameters(): with pytest.raises(ValueError, match="array views"): parse_cuda_signature("__global__ void f(Array2D a) {}", "f") with pytest.raises(ValueError, match="unsupported type"): - parse_cuda_signature("__global__ void f(Array5D a) {}", "f") + parse_cuda_signature("__global__ void f(Array17D a) {}", "f") def test_view_parameters_pack_pointer_shape_and_strides(): @@ -756,6 +757,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")], @@ -869,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)) @@ -964,7 +1088,7 @@ def missing(x, n: int): def unknown(x: "str[:]"): pass - def too_many(x: "float[:, :, :, :, :]"): + def too_many(x: "float[:, :, :, :, :, :, :, :, :, :, :, :, :, :, :, :, :]"): pass def unparsable(x: "float[:](order=F)"): # noqa: F821 (deliberately unparsable) @@ -977,7 +1101,7 @@ def unparsable(x: "float[:](order=F)"): # noqa: F821 (deliberately unparsable) CudaStruct.from_signature(missing, "A") with pytest.raises(ValueError, match="unsupported scalar type 'str'"): CudaStruct.from_signature(unknown, "A") - with pytest.raises(ValueError, match="at most 4 dimensions"): + with pytest.raises(ValueError, match="at most 16 dimensions"): CudaStruct.from_signature(too_many, "A") with pytest.raises(ValueError, match="cannot parse the annotation"): CudaStruct.from_signature(unparsable, "A") @@ -987,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" @@ -1027,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', ) @@ -1403,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) @@ -1693,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" @@ -1758,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")) @@ -1776,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() @@ -1827,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")) @@ -1890,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*"),) @@ -2139,9 +2263,11 @@ def init(self, e: "float[:, :, :, :]", n: int): ... from_annotations = CudaStruct.from_signature(init, "Grid") assert from_annotations.fields[0].ctype == "Array4D" - def too_many(self, e: "float[:, :, :, :, :]"): ... + def too_many( + self, e: "float[:, :, :, :, :, :, :, :, :, :, :, :, :, :, :, :, :]" + ): ... - with pytest.raises(ValueError, match="at most 4 dimensions"): + with pytest.raises(ValueError, match="at most 16 dimensions"): CudaStruct.from_signature(too_many, "Grid") @@ -2257,3 +2383,78 @@ def test_shared_memory_above_the_default_is_opted_in(recorded): with pytest.raises(ValueError, match="exceeds the 100000 bytes"): kernel(1.0, x, y, 1, n_threads=1, shared_mem=100_001) assert [s for *_, s in raw.launches] == [40_000, 80_000, 60_000] + + +@pytest.mark.parametrize("ndim", range(5, 17)) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_high_dimensional_views_pack_and_generate_structs(ndim, contiguous): + ctype = f"{'C' if contiguous else ''}Array{ndim}D" + source = f"__global__ void f({ctype} a) {{}}" + (param,) = parse_cuda_signature(source, "f") + assert (param.view_ndim, param.contiguous) == (ndim, contiguous) + shape = (2, 3) + (1,) * (ndim - 3) + (4,) + array = FakeDeviceArray(np.float64, shape=shape) + if not contiguous: + array.strides = tuple(-2 * s for s in array.strides) + (packed,) = CudaKernel(source, "f").prepare_args(array) + assert packed["data"] == array.data.ptr + assert packed["shape"].tolist() == list(shape) + assert packed.dtype.itemsize == 8 * (1 + ndim * (1 if contiguous else 2)) + if not contiguous: + assert packed["strides"].tolist() == [s // 8 for s in array.strides] + else: + assert packed.dtype.names == ("data", "shape") + with pytest.raises(TypeError, match="must be C-contiguous"): + CudaKernel(source, "f").prepare_args( + FakeDeviceArray(np.float64, shape=shape, strides=(16,) * ndim) + ) + with pytest.raises(TypeError, match=f"must be a {ndim}D array"): + CudaKernel(source, "f").prepare_args(FakeDeviceArray(np.float64)) + with pytest.raises(TypeError, match="dtype"): + CudaKernel(source, "f").prepare_args(FakeDeviceArray(np.float32, shape=shape)) + + def init(a): ... + + init.__annotations__ = {"a": "float[" + ", ".join([":"] * ndim) + "]"} + struct = CudaStruct.from_signature(init, "HighDim", contiguous=contiguous) + assert struct.fields[0].ctype == ctype + assert struct.dtype.fields["a"][0] == packed.dtype + assert f"{ctype} a;" in struct.declaration + value = struct(a=array) + assert value.packed["a"]["shape"].tolist() == list(shape) + + +@pytest.mark.parametrize("ndim", [5, 6, 10, 16]) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_high_dimensional_struct_layout_and_execution_on_gpu(ndim, contiguous): + _skip_without_cupy() + import cupy as cp + + annotation = "float[" + ", ".join([":"] * ndim) + "]" + struct = CudaStruct.from_pyccel_class( + f'class Grid:\n def __init__(self, a: "{annotation}", n: int):\n' + " self.a = a\n self.n = n\n", + "Grid", + contiguous=contiguous, + ) + struct.verify_layout() + indices = ", ".join(["0"] * (ndim - 1) + ["i"]) + source = ( + struct.to_header() + + f""" + #include + #include + extern "C" __global__ void accumulate(Grid g) {{ + CUNUMPY_THREAD_1D(i, g.n); + cunumpy_atomic_add(&g.a({indices}), 2.0); + }} + """ + ) + base = cp.zeros((1,) * (ndim - 1) + (6,)) + a = base if contiguous else base[..., ::2] + CudaKernel(source, "accumulate", structs=(struct,))( + struct(a=a, n=a.size), n_threads=a.size + ) + expected = np.zeros(base.shape) + expected[..., slice(None) if contiguous else slice(None, None, 2)] = 2 + np.testing.assert_array_equal(cp.asnumpy(base), expected) 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_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_emulation.py b/tests/unit/test_emulation.py index 1da8a7d..c5f28f6 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, @@ -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) { @@ -356,3 +387,72 @@ def test_bounds_checks_from_the_view_header(): n_threads=4, options=("-DCUNUMPY_BOUNDS_CHECK",), ) + + +@pytest.mark.parametrize("ndim", range(5, 17)) +@pytest.mark.parametrize("contiguous", [False, True]) +@pytest.mark.parametrize("backend", ["emulation", "cuda"]) +def test_high_dimensional_view_indexing_and_conversion(ndim, contiguous, backend): + # Exercise every axis with unequal extents, including singleton dimensions. + # Reading via a strided device helper also tests CArray -> Array conversion. + ctype = f"{'C' if contiguous else ''}Array{ndim}D" + indices = ", ".join(f"i[{axis}]" for axis in range(ndim)) + source = f""" + #include + #include + __device__ double read(Array{ndim}D a, const long long* i) {{ + return a({indices}); + }} + extern "C" __global__ void update({ctype} a) {{ + CUNUMPY_THREAD_1D(flat, a.size()); + long long i[{ndim}], remaining = flat; + for (int axis = {ndim} - 1; axis >= 0; --axis) {{ + i[axis] = remaining % a.shape[axis]; + remaining /= a.shape[axis]; + }} + a({indices}) = read(a, i) + flat + 1; + }} + """ + shape = (2, 3) + (1,) * (ndim - 3) + (4,) + base = np.arange(48.0).reshape(shape[:-1] + (8,)) + if contiguous: + a = base[..., :4].copy() + else: + a = base[..., ::-2].swapaxes(0, 1) + expected = a.copy() + np.arange(a.size).reshape(a.shape) + 1 + kernel = CudaKernel(source, "update", options=("-DCUNUMPY_BOUNDS_CHECK",)) + if backend == "emulation": + emulate_cuda_kernel(kernel, a, n_threads=a.size) + np.testing.assert_array_equal(a, expected) + if not contiguous: + np.testing.assert_array_equal( + base[..., 0::2], np.arange(48.0).reshape(base.shape)[..., 0::2] + ) + else: + import cunumpy as xp + + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + device_base = cp.asarray(base) + device = cp.asarray(a) if contiguous else device_base[..., ::-2].swapaxes(0, 1) + kernel(device, n_threads=device.size) + np.testing.assert_array_equal(cp.asnumpy(device), expected) + + +@pytest.mark.parametrize("contiguous", [False, True]) +def test_high_dimensional_bounds_check(contiguous): + ctype = f"{'C' if contiguous else ''}Array16D" + indices = ", ".join(["0"] * 15 + ["1"]) + source = f""" + #include + extern "C" __global__ void invalid({ctype} a) {{ a({indices}) = 1; }} + """ + with pytest.raises(RuntimeError, match="crashed"): + emulate_cuda_kernel( + CudaKernel(source, "invalid"), + np.zeros((1,) * 16), + n_threads=1, + options=("-DCUNUMPY_BOUNDS_CHECK",), + ) 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_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..75a5ddd 100644 --- a/tests/unit/test_mpi_serial.py +++ b/tests/unit/test_mpi_serial.py @@ -1,219 +1,34 @@ -"""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) class _DeviceArray: @@ -222,6 +37,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 +45,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 +82,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_namespaces.py b/tests/unit/test_namespaces.py index 28245af..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 @@ -13,6 +13,7 @@ SUBMODULES = ( "algorithms", + "arguments", "cuda", "kernels", "memory", @@ -38,18 +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 + for name in ("CudaKernel", "CudaStruct"): + assert not hasattr(xp.cuda, name) + assert name not in vars(xp) def test_numpy_names_do_not_warn(): @@ -63,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(): @@ -72,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 @@ -102,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(): 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) 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()