cunumpy 0.2.0__tar.gz → 0.3.0__tar.gz

{cunumpy-0.2.0 → cunumpy-0.3.0}/PKG-INFO

@@ -1,19 +1,18 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.2.0
+Version: 0.3.0
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 Project-URL: Source, https://github.com/max-models/cunumpy
 Keywords: python
 Classifier: Development Status :: 3 - Alpha
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Python: >=3.8
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: array-api-compat
 Requires-Dist: numpy
@@ -199,6 +198,22 @@ active CuPy device and `None` on NumPy. `set_device_for_rank(rank)` is a
 round-robin convenience for MPI layouts where local ranks map contiguously to
 GPUs. If your scheduler uses a different mapping, select the device directly.
 
+For MPI programs with one rank per GPU, `bind_local_device()` selects the GPU
+from the node-local rank that the MPI launcher exports (`local_rank()`), so it
+can run before MPI is initialized, as CUDA-aware MPI requires. Before passing
+device buffers to MPI, call `synchronize_for_mpi(*buffers)`: kernels run
+asynchronously, and MPI would otherwise send a buffer a kernel is still
+writing, without an error.
+
+```python
+xp.set_backend("cupy")
+xp.bind_local_device()  # before MPI_Init
+from mpi4py import MPI
+
+xp.synchronize_for_mpi(send, recv)
+MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source)
+```
+
 CuPy caches released allocations in memory pools. This can make process-level
 GPU memory appear occupied after arrays go out of scope. `free_memory()` asks
 CuPy to release currently free cached blocks; it does not free memory still
@@ -253,6 +268,80 @@ The wrapper can also traverse arrays nested in lists, tuples, dictionaries,
 and selected application objects; see the full [API reference](docs/source/api.md)
 for `object_modules`, `is_array`, aliasing, and output declarations.
 
+## Write CUDA kernels next to host kernels
+
+`CudaKernel` wraps a CUDA C kernel (compiled with NVRTC through
+`cupy.RawKernel`) so that it is called with the same arguments as the host
+kernel it mirrors, plus the number of threads. Arrays are never copied: they
+must be CuPy arrays. The `extern "C" __global__` signature is parsed once and
+every call is checked against it: Python scalars are cast to the declared C
+types, and a wrong argument count, an array of the wrong dtype, or a scalar
+that does not fit its type raises instead of silently producing wrong values.
+
+`Kernel` pairs a host kernel with its CUDA kernel and calls the one matching
+the active backend, so kernels can be ported to CUDA one at a time:
+
+```python
+import cunumpy as xp
+
+AXPY = r"""
+extern "C" __global__
+void axpy(double a, const double* x, double* y, int n) {
+    int i = blockDim.x * blockIdx.x + threadIdx.x;
+    if (i < n) y[i] += a * x[i];
+}
+"""
+
+
+def axpy(a, x, y, n):  # host version, e.g. compiled with Pyccel
+    for i in range(n):
+        y[i] += a * x[i]
+
+
+kernel = xp.Kernel(axpy, xp.CudaKernel(AXPY, "axpy"))
+
+with xp.use_backend("cupy"):
+    x = xp.arange(1000, dtype=xp.float64)
+    y = xp.zeros(1000)
+    kernel(2.0, x, y, 1000, n_threads=1000)  # runs the CUDA kernel
+```
+
+On the CuPy backend, a `Kernel` without CUDA kernel raises
+`NotImplementedError` (or, with `missing_cuda="fallback"`, runs the host kernel
+through `PyccelKernel`, with host copies; `host_options` configure that
+`PyccelKernel`). `KernelCatalog.from_package()` collects kernel pairs from a
+package with one folder per kernel (`name/name_kernels.py` and
+`name/name_cuda.cu`), and `catalog.compile_all()` compiles all CUDA kernels at
+setup.
+
+Groups of arguments can be passed as one: objects implementing
+`__cuda_args__()` (see `CudaArguments`) are flattened into several kernel
+arguments, and `CudaStruct` defines a C struct once (its C `declaration` and
+the matching memory layout) and packs values into it, which the kernel takes
+as one parameter:
+
+```python
+Vec = xp.CudaStruct("Vec", [("data", "double*"), ("n", "int")])
+scale = xp.CudaKernel(
+    Vec.declaration
+    + r"""
+    extern "C" __global__ void scale(Vec v, double a) {
+        int i = blockDim.x * blockIdx.x + threadIdx.x;
+        if (i < v.n) v.data[i] *= a;
+    }""",
+    "scale",
+    structs=[Vec],
+)
+scale(Vec(data=y, n=y.size), 0.5, n_threads=y.size)
+```
+
+Launches can be 1D to 3D (`n_threads=(nx, ny)`, `block_size=(16, 16)`) or use
+an explicit `grid`, with dynamic shared memory (`shared_mem`) and a `stream`.
+C++ function templates are instantiated with `template_args`, and
+`CudaKernelVariants` caches kernels whose source is generated per variant
+(e.g. per dimension and dtype). See the [API reference](docs/source/api.md) for
+details.
+
 ## Pyodide
 
 CuNumpy supports the NumPy backend in Pyodide. It does not provide CuPy/CUDA

{cunumpy-0.2.0 → cunumpy-0.3.0}/README.md

@@ -159,6 +159,22 @@ active CuPy device and `None` on NumPy. `set_device_for_rank(rank)` is a
 round-robin convenience for MPI layouts where local ranks map contiguously to
 GPUs. If your scheduler uses a different mapping, select the device directly.
 
+For MPI programs with one rank per GPU, `bind_local_device()` selects the GPU
+from the node-local rank that the MPI launcher exports (`local_rank()`), so it
+can run before MPI is initialized, as CUDA-aware MPI requires. Before passing
+device buffers to MPI, call `synchronize_for_mpi(*buffers)`: kernels run
+asynchronously, and MPI would otherwise send a buffer a kernel is still
+writing, without an error.
+
+```python
+xp.set_backend("cupy")
+xp.bind_local_device()  # before MPI_Init
+from mpi4py import MPI
+
+xp.synchronize_for_mpi(send, recv)
+MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source)
+```
+
 CuPy caches released allocations in memory pools. This can make process-level
 GPU memory appear occupied after arrays go out of scope. `free_memory()` asks
 CuPy to release currently free cached blocks; it does not free memory still
@@ -213,6 +229,80 @@ The wrapper can also traverse arrays nested in lists, tuples, dictionaries,
 and selected application objects; see the full [API reference](docs/source/api.md)
 for `object_modules`, `is_array`, aliasing, and output declarations.
 
+## Write CUDA kernels next to host kernels
+
+`CudaKernel` wraps a CUDA C kernel (compiled with NVRTC through
+`cupy.RawKernel`) so that it is called with the same arguments as the host
+kernel it mirrors, plus the number of threads. Arrays are never copied: they
+must be CuPy arrays. The `extern "C" __global__` signature is parsed once and
+every call is checked against it: Python scalars are cast to the declared C
+types, and a wrong argument count, an array of the wrong dtype, or a scalar
+that does not fit its type raises instead of silently producing wrong values.
+
+`Kernel` pairs a host kernel with its CUDA kernel and calls the one matching
+the active backend, so kernels can be ported to CUDA one at a time:
+
+```python
+import cunumpy as xp
+
+AXPY = r"""
+extern "C" __global__
+void axpy(double a, const double* x, double* y, int n) {
+    int i = blockDim.x * blockIdx.x + threadIdx.x;
+    if (i < n) y[i] += a * x[i];
+}
+"""
+
+
+def axpy(a, x, y, n):  # host version, e.g. compiled with Pyccel
+    for i in range(n):
+        y[i] += a * x[i]
+
+
+kernel = xp.Kernel(axpy, xp.CudaKernel(AXPY, "axpy"))
+
+with xp.use_backend("cupy"):
+    x = xp.arange(1000, dtype=xp.float64)
+    y = xp.zeros(1000)
+    kernel(2.0, x, y, 1000, n_threads=1000)  # runs the CUDA kernel
+```
+
+On the CuPy backend, a `Kernel` without CUDA kernel raises
+`NotImplementedError` (or, with `missing_cuda="fallback"`, runs the host kernel
+through `PyccelKernel`, with host copies; `host_options` configure that
+`PyccelKernel`). `KernelCatalog.from_package()` collects kernel pairs from a
+package with one folder per kernel (`name/name_kernels.py` and
+`name/name_cuda.cu`), and `catalog.compile_all()` compiles all CUDA kernels at
+setup.
+
+Groups of arguments can be passed as one: objects implementing
+`__cuda_args__()` (see `CudaArguments`) are flattened into several kernel
+arguments, and `CudaStruct` defines a C struct once (its C `declaration` and
+the matching memory layout) and packs values into it, which the kernel takes
+as one parameter:
+
+```python
+Vec = xp.CudaStruct("Vec", [("data", "double*"), ("n", "int")])
+scale = xp.CudaKernel(
+    Vec.declaration
+    + r"""
+    extern "C" __global__ void scale(Vec v, double a) {
+        int i = blockDim.x * blockIdx.x + threadIdx.x;
+        if (i < v.n) v.data[i] *= a;
+    }""",
+    "scale",
+    structs=[Vec],
+)
+scale(Vec(data=y, n=y.size), 0.5, n_threads=y.size)
+```
+
+Launches can be 1D to 3D (`n_threads=(nx, ny)`, `block_size=(16, 16)`) or use
+an explicit `grid`, with dynamic shared memory (`shared_mem`) and a `stream`.
+C++ function templates are instantiated with `template_args`, and
+`CudaKernelVariants` caches kernels whose source is generated per variant
+(e.g. per dimension and dtype). See the [API reference](docs/source/api.md) for
+details.
+
 ## Pyodide
 
 CuNumpy supports the NumPy backend in Pyodide. It does not provide CuPy/CUDA

{cunumpy-0.2.0 → cunumpy-0.3.0}/pyproject.toml

@@ -5,22 +5,21 @@ requires = [ "setuptools", "wheel" ]
 
 [project]
 name = "cunumpy"
-version = "0.2.0"
+version = "0.3.0"
 description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
 readme = "README.md"
 keywords = [ "python" ]
 license = { file = "LICENSE.txt" }
 authors = [ { name = "Max" } ]
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 classifiers = [
   "Development Status :: 3 - Alpha",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.8",
-  "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
 ]
 dependencies = [
   "array-api-compat",

{cunumpy-0.2.0 → cunumpy-0.3.0}/src/cunumpy/__init__.py

@@ -2,9 +2,21 @@
 from importlib.metadata import PackageNotFoundError, version
 
 from . import xp
+from .cuda_kernel import (
+    CudaArguments,
+    CudaKernel,
+    CudaKernelVariants,
+    CudaParameter,
+    CudaStruct,
+    CudaStructValue,
+    ctype_of,
+    parse_cuda_signature,
+)
+from .dispatch import Kernel, KernelCatalog
 from .kernel import PyccelKernel
 from .xp import (
     assert_same_backend,
+    bind_local_device,
     cupy_available,
     default_float_dtype,
     device_count,
@@ -15,6 +27,7 @@ from .xp import (
     get_rng,
     is_cpu,
     is_gpu,
+    local_rank,
     memory_info,
     pin_memory,
     same_backend,
@@ -23,6 +36,7 @@ from .xp import (
     set_device_for_rank,
     stream,
     synchronize,
+    synchronize_for_mpi,
     to_cunumpy,
     to_cupy,
     to_numpy,
@@ -35,9 +49,19 @@ except PackageNotFoundError:
     __version__ = "0.0.0+unknown"
 
 __all__ = [
+    "CudaArguments",
+    "CudaKernel",
+    "CudaKernelVariants",
+    "CudaParameter",
+    "CudaStruct",
+    "CudaStructValue",
+    "Kernel",
+    "KernelCatalog",
     "PyccelKernel",
     "__version__",
     "assert_same_backend",
+    "bind_local_device",
+    "ctype_of",
     "cupy_available",
     "cupy_backend",
     "default_float_dtype",
@@ -49,8 +73,10 @@ __all__ = [
     "get_rng",
     "is_cpu",
     "is_gpu",
+    "local_rank",
     "memory_info",
     "numpy_backend",
+    "parse_cuda_signature",
     "pin_memory",
     "same_backend",
     "set_backend",
@@ -58,6 +84,7 @@ __all__ = [
     "set_device_for_rank",
     "stream",
     "synchronize",
+    "synchronize_for_mpi",
     "to_cunumpy",
     "to_cupy",
     "to_numpy",

{cunumpy-0.2.0 → cunumpy-0.3.0}/src/cunumpy/__init__.pyi

@@ -1,13 +1,24 @@
 # Stub file for Pylance/mypy: exposes all numpy symbols so that
 # `import cunumpy as xp` followed by `xp.<Tab>` shows numpy completions.
 # At runtime the real __init__.py dispatches to numpy or cupy via __getattr__.
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator
+from typing import Any
 
 import numpy as np
 from numpy import *
 
 from . import xp as xp
+from .cuda_kernel import CudaArguments as CudaArguments
+from .cuda_kernel import CudaKernel as CudaKernel
+from .cuda_kernel import CudaKernelVariants as CudaKernelVariants
+from .cuda_kernel import CudaParameter as CudaParameter
+from .cuda_kernel import CudaStruct as CudaStruct
+from .cuda_kernel import CudaStructValue as CudaStructValue
+from .cuda_kernel import ctype_of as ctype_of
+from .cuda_kernel import parse_cuda_signature as parse_cuda_signature
+from .dispatch import Kernel as Kernel
+from .dispatch import KernelCatalog as KernelCatalog
 from .kernel import PyccelKernel as PyccelKernel
 
 def to_numpy(array: Any) -> np.ndarray: ...
@@ -26,6 +37,9 @@ def use_backend(backend: str) -> Generator[None]: ...
 def set_backend(backend: str) -> None: ...
 def set_device(device_id: int) -> None: ...
 def set_device_for_rank(rank: int, devices_per_node: int | None = ...) -> int: ...
+def local_rank() -> int: ...
+def bind_local_device() -> int | None: ...
+def synchronize_for_mpi(*arrays: Any) -> None: ...
 def device_count() -> int: ...
 def memory_info() -> tuple[int, int] | None: ...
 def free_memory() -> None: ...