From 870f51868b23d4ab9945ee48d7156a744647f95a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 21 Oct 2026 22:25:17 +0200 Subject: [PATCH 1/9] Test PRs to devel against the max-models devel branches and PRs to main against PyPI (#112) --- .github/workflows/gpu_ci_trigger.yml | 6 ++++-- .github/workflows/macos.yml | 2 ++ .github/workflows/pyodide.yml | 2 ++ .github/workflows/static_analysis.yml | 1 + .github/workflows/testing.yml | 19 +++++++++++++++++++ .gitlab-ci.yml | 9 +++++++++ .../docs/getting-started/installation.md | 5 +++++ pyproject.toml | 5 +++++ requirements-devel.txt | 4 ++++ 9 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 requirements-devel.txt diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 7f7a8e7..80d4415 100644 --- a/.github/workflows/gpu_ci_trigger.yml +++ b/.github/workflows/gpu_ci_trigger.yml @@ -63,8 +63,10 @@ jobs: # 3. Add GitLab SSH remote git remote add gitlab git@gitlab.mpcdf.mpg.de:maxlin/cunumpy.git - # 4. Force push (This automatically starts the GitLab Pipeline) - git push -f gitlab "HEAD:refs/heads/$TARGET_BRANCH" + # 4. Force push (This automatically starts the GitLab Pipeline). The pushed branch + # is gpu-test-*, so pass the target branch for the max-models dependencies. + git push -f -o ci.variable="MAX_MODELS_TARGET=${{ github.base_ref || github.ref_name }}" \ + gitlab "HEAD:refs/heads/$TARGET_BRANCH" echo "PUSHED_SHA=$(git rev-parse HEAD)" >> "$GITHUB_ENV" # 5. Provide the direct link diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index b13bab3..388d28f 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -43,6 +43,8 @@ jobs: - name: Install project run: | pip install --upgrade pip + # devel branches of the max-models dependencies, except on runs that target main + if [ "${GITHUB_BASE_REF:-$GITHUB_REF_NAME}" != main ]; then python -m pip install -r requirements-devel.txt; fi pip install ".[test-compiled,metal]" # Hosted macOS runners are virtual machines and usually have no Metal GPU, diff --git a/.github/workflows/pyodide.yml b/.github/workflows/pyodide.yml index dc35744..c431dee 100644 --- a/.github/workflows/pyodide.yml +++ b/.github/workflows/pyodide.yml @@ -6,6 +6,8 @@ on: pull_request: branches: [main, devel] +# PyPI only, also on devel: micropip, which installs the wheel's dependencies in Pyodide, +# cannot install from git. jobs: wheel-test: runs-on: ubuntu-latest diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 02b58a7..ae0eb4d 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -82,6 +82,7 @@ jobs: - name: Docstring sections with numpydoc run: | + if [ "${GITHUB_BASE_REF:-$GITHUB_REF_NAME}" != main ]; then python -m pip install -r requirements-devel.txt; fi pip install . numpydoc python -m numpydoc lint $(git ls-files 'src/cunumpy/*.py' 'src/cunumpy/**/*.py' | grep -v _fake_cupy) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 0f4d562..1676664 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -10,6 +10,8 @@ on: - devel workflow_call: +# Runs that do not target main install the devel branches of the max-models dependencies +# (requirements-devel.txt) before cunumpy; PRs to main and the release use PyPI only. jobs: compiled-fake-cupy: runs-on: ubuntu-latest @@ -22,6 +24,7 @@ jobs: run: | sudo apt-get update sudo apt-get install -y gcc gfortran + if [ "${GITHUB_BASE_REF:-$GITHUB_REF_NAME}" != main ]; then python -m pip install -r requirements-devel.txt; fi python -m pip install ".[test-compiled]" - name: Test compiled host kernels with fake device arrays env: @@ -61,6 +64,7 @@ jobs: - name: Install project run: | pip install --upgrade pip + if [ "${GITHUB_BASE_REF:-$GITHUB_REF_NAME}" != main ]; then python -m pip install -r requirements-devel.txt; fi pip install ".[dev]" - name: Run tests @@ -70,3 +74,18 @@ jobs: - name: Run docstring examples run: | pytest --doctest-modules src/cunumpy + + # PyPI rejects direct references (`name @ git+...`): the wheel's metadata must have none. + no-direct-references: + if: (github.base_ref || github.ref_name) == 'main' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Check the wheel for direct references + run: | + python -m pip install build + python -m build --wheel + ! unzip -p dist/*.whl '*.dist-info/METADATA' | grep -E '^Requires-Dist:.* @ ' diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 1d852e4..cba3e93 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -14,6 +14,13 @@ variables: stages: - test +# Runs that do not target main install the devel branches of the max-models dependencies +# first; main uses PyPI only. The GitHub trigger pushes gpu-test-* branches and passes the +# target branch in MAX_MODELS_TARGET. +.install_devel_dependencies: &install_devel_dependencies + - TARGET="${MAX_MODELS_TARGET:-${CI_MERGE_REQUEST_TARGET_BRANCH_NAME:-$CI_COMMIT_BRANCH}}" + - if [ "$TARGET" != main ]; then python3 -m pip install --user -r requirements-devel.txt; fi + .gpu_job_nvidia: stage: test image: ${CUDA_IMAGE} @@ -27,6 +34,7 @@ stages: - nvidia-smi # Keep NVRTC, CUDA libraries and headers on the same toolkit generation. - python3 -m pip install --user "cupy-cuda13x[ctk]${CUPY_PIN}" "cuda-toolkit==13.2.*" + - *install_devel_dependencies - python3 -m pip install --user -e ".[test]" - export PATH="$HOME/.local/bin:$PATH" - python3 -c "import cupy; cupy.show_config()" @@ -49,6 +57,7 @@ stages: - export ROCM_HOME="${ROCM_PATH:-/opt/rocm}" - export CUPY_INSTALL_USE_HIP=1 - python3 -m pip install --user "cupy${CUPY_PIN}" + - *install_devel_dependencies - python3 -m pip install --user -e ".[test]" - export PATH="$HOME/.local/bin:$PATH" - python3 -c "import cupy; cupy.show_config()" diff --git a/docs/src/content/docs/getting-started/installation.md b/docs/src/content/docs/getting-started/installation.md index 5a04447..7653a25 100644 --- a/docs/src/content/docs/getting-started/installation.md +++ b/docs/src/content/docs/getting-started/installation.md @@ -88,6 +88,11 @@ python -m pytest tests/unit Tests that need a GPU are skipped automatically where CuPy is not functional. +On `devel`, cunumpy is developed against the `devel` branches of the other max-models packages +(currently maybempi): `uv sync` picks them up from `[tool.uv.sources]`, and with pip run +`python -m pip install -r requirements-devel.txt` before the install above. CI does the same for +pull requests to `devel`; pull requests to `main` and the PyPI release use PyPI versions only. + ## Environment variables | Variable | Effect | diff --git a/pyproject.toml b/pyproject.toml index f70ac73..34c15e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,11 @@ where = [ "src" ] [tool.setuptools.package-data] cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] +# PRs to devel test against the devel branches of the max-models packages; PRs to main +# and the release use PyPI only. +[tool.uv.sources] +maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" } + [tool.ruff] # agent worktrees and scratch scripts are local copies, not part of the project extend-exclude = [ ".claude" ] diff --git a/requirements-devel.txt b/requirements-devel.txt new file mode 100644 index 0000000..ab53722 --- /dev/null +++ b/requirements-devel.txt @@ -0,0 +1,4 @@ +# The devel branches of the max-models packages in cunumpy's dependency tree. CI installs +# them before cunumpy on every run that does not target main; main and the release use PyPI. +# pip install -r requirements-devel.txt +maybempi @ git+https://github.com/max-models/maybempi@devel From 5ebc6e39a86ad5b450c69c435b133d2d70552206 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 22 Oct 2026 08:42:17 +0200 Subject: [PATCH 2/9] Opt-in scope-profiler regions, off by default (#113) * Add opt-in scope-profiler regions, off by default * Merge the uv sources tables * Install scope-profiler from devel too * Give the Pyodide test the devel branches as wheels --- .github/workflows/pyodide.yml | 10 +++- CHANGELOG.md | 9 +++ docs/src/content/docs/guides/profiling.md | 22 +++++++ pyproject.toml | 2 + requirements-devel.txt | 1 + src/cunumpy/__init__.py | 7 +++ src/cunumpy/_algorithms.py | 5 ++ src/cunumpy/_cuda_kernel.py | 33 ++++++----- src/cunumpy/_dispatch.py | 71 ++++++++++++----------- src/cunumpy/_linalg.py | 3 + src/cunumpy/_mirror.py | 34 ++++++----- src/cunumpy/_morton.py | 3 + src/cunumpy/_mpi.py | 41 +++++++------ src/cunumpy/_optimize.py | 4 ++ src/cunumpy/_profile_tags.py | 9 +++ src/cunumpy/_pyccel_cache.py | 5 ++ src/cunumpy/_scipy_host.py | 8 +++ tests/pyodide/run.mjs | 20 +++++-- tests/unit/test_opt_in_profiling.py | 37 ++++++++++++ 19 files changed, 234 insertions(+), 90 deletions(-) create mode 100644 src/cunumpy/_profile_tags.py create mode 100644 tests/unit/test_opt_in_profiling.py diff --git a/.github/workflows/pyodide.yml b/.github/workflows/pyodide.yml index c431dee..8ba8ea2 100644 --- a/.github/workflows/pyodide.yml +++ b/.github/workflows/pyodide.yml @@ -6,8 +6,8 @@ on: pull_request: branches: [main, devel] -# PyPI only, also on devel: micropip, which installs the wheel's dependencies in Pyodide, -# cannot install from git. +# Runs that do not target main take the max-models packages from their devel branches +# (requirements-devel.txt); micropip cannot install from git, so they reach Pyodide as wheels. jobs: wheel-test: runs-on: ubuntu-latest @@ -24,6 +24,10 @@ jobs: cache-dependency-path: tests/pyodide/package-lock.json - name: Build wheel and install compiler-free test dependencies run: | + if [ "${GITHUB_BASE_REF:-$GITHUB_REF_NAME}" != main ]; then + python -m pip install -r requirements-devel.txt + python -m pip wheel --no-deps -w devel-wheels -r requirements-devel.txt + fi python -m pip install build '.[test]' python -m build --wheel - name: Run portable tests in native Python @@ -31,4 +35,4 @@ jobs: - name: Install Pyodide run: npm ci --prefix tests/pyodide - name: Run portable tests against the wheel in Pyodide - run: node tests/pyodide/run.mjs dist/cunumpy-*-py3-none-any.whl + run: node tests/pyodide/run.mjs dist/cunumpy-*-py3-none-any.whl $(ls devel-wheels/*.whl 2>/dev/null) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b46bef..8d0fa9b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- scope-profiler regions around the coarse operations (`Kernel` calls as + `cunumpy.kernel.`, CUDA and Pyccel compilation, `sort_by_key`, + `SegmentPlan.sum`, `compact_by_mask`, `morton_keys`, the solvers, `exchange`, + and the device copies of `mpi_buffer` and `DeviceMirror`), behind the opt-in + tag `cunumpy`: off unless the application enables it, e.g. + `sp.setup(enable_tags=["cunumpy"])` or `SCOPE_PROFILER_ENABLE_TAGS=cunumpy`. + New dependency: `scope-profiler>=0.7.5`. + ### Changed - `xp.algorithms.sort_by_key` gathers a 2-D NumPy array along its last axis (component-major markers, `axis=-1`) row by row, about 2x faster than `take` (1e6 markers: 3 rows 12.9 to 6.7 ms, diff --git a/docs/src/content/docs/guides/profiling.md b/docs/src/content/docs/guides/profiling.md index eb7ee15..001fac4 100644 --- a/docs/src/content/docs/guides/profiling.md +++ b/docs/src/content/docs/guides/profiling.md @@ -86,6 +86,28 @@ another library converting internally) show up in Nsight as `memcpy` rows on the timeline; NVTX ranges tell you which part of the step issued them. See [Moving data between host and device](/cunumpy/guides/data-movement/). +## Cunumpy's own regions in scope-profiler + +Cunumpy ships [scope-profiler](https://github.com/max-models/scope-profiler) +regions around its coarse operations: every `Kernel` call +(`cunumpy.kernel.`), CUDA and Pyccel compilation, `sort_by_key`, +`segment_sum` (`cunumpy.SegmentPlan.sum`), `compact_by_mask`, `morton_keys`, the +solvers, `exchange` and the host/device copies of `mpi_buffer` and +`DeviceMirror`. They are off by default, so an application that profiles itself +does not see them, and cost at most about 0.2 µs per call while off. Turn them on with the `cunumpy` tag: + +```python +import scope_profiler as sp + +sp.setup(enable_tags=["cunumpy"]) +``` + +or `SCOPE_PROFILER_ENABLE_TAGS=cunumpy`, or +`scope-profiler run --enable-tag cunumpy simulate.py` (`"*"` enables every +library's tag). The regions time the host side: on CuPy, a kernel region +measures the launch, not the kernel; wrap a synchronizing `timed_region` around +a loop of steps for device times. + ## A profiling checklist 1. Run a few steps with `count_transfers()` and remove every transfer from the diff --git a/pyproject.toml b/pyproject.toml index 34c15e3..f73a9f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "array-api-compat", "maybempi>=0.1.2", "numpy", + "scope-profiler>=0.7.5", ] optional-dependencies.dev = [ @@ -51,6 +52,7 @@ cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] # and the release use PyPI only. [tool.uv.sources] maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" } +scope-profiler = { git = "https://github.com/max-models/scope-profiler", branch = "devel" } [tool.ruff] # agent worktrees and scratch scripts are local copies, not part of the project diff --git a/requirements-devel.txt b/requirements-devel.txt index ab53722..84b537c 100644 --- a/requirements-devel.txt +++ b/requirements-devel.txt @@ -2,3 +2,4 @@ # them before cunumpy on every run that does not target main; main and the release use PyPI. # pip install -r requirements-devel.txt maybempi @ git+https://github.com/max-models/maybempi@devel +scope-profiler @ git+https://github.com/max-models/scope-profiler@devel diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index f61e4ed..d8743db 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -15,6 +15,8 @@ import re as _re from importlib.metadata import PackageNotFoundError, version +import scope_profiler as _sp + from cunumpy import ( algorithms, arguments, @@ -30,6 +32,7 @@ xp, ) from cunumpy._host import evaluate_on_host, host_call, setup_on_host +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy._scipy_backend import scipy from cunumpy._staging import to_host_async from cunumpy.xp import ( @@ -54,6 +57,10 @@ use_backend, ) +# cunumpy's scope-profiler regions are off unless the application enables the +# "cunumpy" tag +_sp.register_opt_in_tag(*_PROFILE_TAGS) + try: __version__ = version("cunumpy") """The installed version of cunumpy; ``"0.0.0+unknown"`` without package metadata.""" diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index d7e1621..76355df 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -7,8 +7,10 @@ from typing import Any import array_api_compat.numpy as np +import scope_profiler as sp from cunumpy import _fake_cupy +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy.xp import assert_same_backend, get_array_backend, get_array_module @@ -202,6 +204,7 @@ def n_segments(self) -> int: """Number of output segments, fixed when the plan is created.""" return self._n_segments + @sp.profile("cunumpy.SegmentPlan.sum", tags=_PROFILE_TAGS) def sum(self, values: Any, *, out: Any = None) -> Any: """Sum the rows of `values` per segment. @@ -371,6 +374,7 @@ def _radix_argsort(keys: np.ndarray) -> np.ndarray: return order.astype(np.int64, copy=False) +@sp.profile("cunumpy.sort_by_key", tags=_PROFILE_TAGS) def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]: """Sort `keys` and reorder every array the same way, in one stable argsort. @@ -460,6 +464,7 @@ def _take_along_axis(xpm: Any, array: Any, order: Any, axis: int) -> Any: return xpm.take(array, order, axis=axis) +@sp.profile("cunumpy.compact_by_mask", tags=_PROFILE_TAGS) def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int: """Move the entries where `mask` is True to the front of every array, in place. diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 2db6d94..3e72f02 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -34,7 +34,9 @@ from typing import Any, NamedTuple import numpy as np +import scope_profiler as sp +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _record_sync @@ -2387,21 +2389,22 @@ def compile(self, *, log_stream: Any = None) -> Any: "without CUNUMPY_BOUNDS_CHECK (and without debug mode, which adds " "it) for a kernel using array views on this build.", ) - if self._template_args is None: - raw = cp.RawKernel( - self._source, - self._name, - options=options, - ) - raw.compile(log_stream=log_stream) - else: - module = cp.RawModule( - code=self._source, - options=options, - name_expressions=[self.expression], - ) - module.compile(log_stream=log_stream) - raw = module.get_function(self.expression) + with sp.region("cunumpy.CudaKernel.compile", tags=_PROFILE_TAGS): + if self._template_args is None: + raw = cp.RawKernel( + self._source, + self._name, + options=options, + ) + raw.compile(log_stream=log_stream) + else: + module = cp.RawModule( + code=self._source, + options=options, + name_expressions=[self.expression], + ) + module.compile(log_stream=log_stream) + raw = module.get_function(self.expression) limits = _device_limits(device) attributes = dict(raw.attributes) static = max(0, attributes.get("shared_size_bytes", 0)) diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 91b5708..b3bc8a2 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -30,6 +30,7 @@ import array_api_compat import numpy as np +import scope_profiler as sp from cunumpy._cuda_kernel import CudaKernel, _compile_in_threads from cunumpy._kernel import ( @@ -39,6 +40,7 @@ get_device_kernel_implementation, outputs_from_annotations, ) +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _record from cunumpy.xp import get_backend @@ -240,6 +242,7 @@ def __init__( self._host_kernel = host_kernel self._cuda_kernel = cuda_kernel self._name = name if name is not None else host_kernel.name + self._region = f"cunumpy.kernel.{self._name}" self._missing_cuda = missing_cuda self._cuda_path = None if cuda_path is None else Path(cuda_path) self._test_args_module = test_args @@ -675,41 +678,42 @@ def __call__( If the CUDA kernel runs without `n_threads`, `grid` or ``n_threads_from``. """ - if self._dispatch == "arrays": - on_device = any(_on_device(arg) for arg in args) - kernel = self._device_kernel() if on_device else self._host_kernel - else: - on_device = get_backend() == "cupy" - kernel = self.get_kernel() - if kernel is self._host_kernel: - if _COUNTERS and self._cuda_kernel is None and on_device: - _record( - "fallback", - f"Kernel {self._name!r} has no CUDA kernel: host kernel " - "called on the CuPy backend", + with sp.region(self._region, tags=_PROFILE_TAGS): + if self._dispatch == "arrays": + on_device = any(_on_device(arg) for arg in args) + kernel = self._device_kernel() if on_device else self._host_kernel + else: + on_device = get_backend() == "cupy" + kernel = self.get_kernel() + if kernel is self._host_kernel: + if _COUNTERS and self._cuda_kernel is None and on_device: + _record( + "fallback", + f"Kernel {self._name!r} has no CUDA kernel: host kernel " + "called on the CuPy backend", + ) + if ( + self._dispatch == "arrays" + and not on_device + and kernel._use_cupy is not True + ): + # only host arguments, by the choice above: nothing to convert, + # even while CuPy is the active backend + return kernel.kernel(*args) + return kernel(*args) + if n_threads is None and grid is None and kernel.n_threads_from is None: + raise ValueError( + f"{self._name}: n_threads is required to launch the CUDA kernel " + "(or pass grid, or set cuda_kernel.n_threads_from)", ) - if ( - self._dispatch == "arrays" - and not on_device - and kernel._use_cupy is not True - ): - # only host arguments, by the choice above: nothing to convert, - # even while CuPy is the active backend - return kernel.kernel(*args) - return kernel(*args) - if n_threads is None and grid is None and kernel.n_threads_from is None: - raise ValueError( - f"{self._name}: n_threads is required to launch the CUDA kernel " - "(or pass grid, or set cuda_kernel.n_threads_from)", + return kernel( + *args, + n_threads=n_threads, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, ) - return kernel( - *args, - n_threads=n_threads, - grid=grid, - block=block, - shared_mem=shared_mem, - stream=stream, - ) class KernelCatalog(Mapping): @@ -978,6 +982,7 @@ def check_signatures(self) -> None: + "\n ".join(problems), ) + @sp.profile("cunumpy.KernelCatalog.compile_all", tags=_PROFILE_TAGS) def compile_all(self, jobs: int | None = 1) -> list[str]: """Compile all CUDA kernels now, e.g. at setup instead of in the first step. diff --git a/src/cunumpy/_linalg.py b/src/cunumpy/_linalg.py index 41e8ffb..00af1be 100644 --- a/src/cunumpy/_linalg.py +++ b/src/cunumpy/_linalg.py @@ -5,7 +5,9 @@ from typing import Any import numpy as np +import scope_profiler as sp +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy.xp import get_array_module @@ -24,6 +26,7 @@ def _as_inexact(xpm: Any, a: Any) -> Any: return a +@sp.profile("cunumpy.solve_circulant", tags=_PROFILE_TAGS) def solve_circulant( c: Any, b: Any, diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py index 44192f6..746ded7 100644 --- a/src/cunumpy/_mirror.py +++ b/src/cunumpy/_mirror.py @@ -12,7 +12,9 @@ from typing import Any import numpy as np +import scope_profiler as sp +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _describe, _nbytes, _record from cunumpy.xp import _cupy_backend, to_cupy @@ -160,17 +162,18 @@ def to_device(self) -> DeviceMirror: if not _cupy_backend(): return self self._check_device() - if self._device is None: - _ = self.device # allocates as a copy of the host, through to_cupy - else: - # host -> device copy into the existing device buffer - self._device.set(self._host) - if _COUNTERS: - _record( - "to_device", - f"DeviceMirror.to_device({_describe(self._host)})", - nbytes=_nbytes(self._host), - ) + with sp.region("cunumpy.DeviceMirror.to_device", tags=_PROFILE_TAGS): + if self._device is None: + _ = self.device # allocates as a copy of the host, through to_cupy + else: + # host -> device copy into the existing device buffer + self._device.set(self._host) + if _COUNTERS: + _record( + "to_device", + f"DeviceMirror.to_device({_describe(self._host)})", + nbytes=_nbytes(self._host), + ) return self def to_host(self, *, stream: Any = None, event: Any = None) -> DeviceMirror: @@ -222,10 +225,11 @@ def to_host(self, *, stream: Any = None, event: Any = None) -> DeviceMirror: if event is not None: cp.cuda.get_current_stream().wait_event(event) # device -> host copy into the existing host buffer - if stream is None: - self._device.get(out=self._host) - else: - self._device.get(out=self._host, stream=stream) + with sp.region("cunumpy.DeviceMirror.to_host", tags=_PROFILE_TAGS): + if stream is None: + self._device.get(out=self._host) + else: + self._device.get(out=self._host, stream=stream) if _COUNTERS: _record( "to_host", diff --git a/src/cunumpy/_morton.py b/src/cunumpy/_morton.py index 49406da..79ade5c 100644 --- a/src/cunumpy/_morton.py +++ b/src/cunumpy/_morton.py @@ -22,8 +22,10 @@ from typing import Any import numpy as np +import scope_profiler as sp from cunumpy._philox import _module +from cunumpy._profile_tags import _PROFILE_TAGS __all__ = [ "MAX_MORTON_LEVELS", @@ -196,6 +198,7 @@ def morton_scales(lower: Sequence[float], upper: Sequence[float], levels: int) - return float(2**levels) / (upper - lower) +@sp.profile("cunumpy.morton_keys", tags=_PROFILE_TAGS) def morton_keys( positions: Any, lower: Sequence[float], diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index c923ddd..30ec632 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -10,8 +10,10 @@ import array_api_compat import array_api_compat.numpy as np +import scope_profiler as sp from maybempi import get_mpi +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _describe, _nbytes, _record, _record_sync from cunumpy.xp import array_backend, cupy_available, to_numpy @@ -340,9 +342,10 @@ def mpi_buffer( synchronize_for_mpi(array, stream=stream, event=event) transfer_array = staging._transfer_array(array) if send: - if transfer_array is not array: - transfer_array[...] = array - transfer_array.get(out=host) + with sp.region("cunumpy.mpi_buffer.to_host", tags=_PROFILE_TAGS): + if transfer_array is not array: + transfer_array[...] = array + transfer_array.get(out=host) if _COUNTERS: _record( "to_host", @@ -351,21 +354,23 @@ def mpi_buffer( ) yield host if recv: - transfer_array.set(host) - if _COUNTERS: - _record( - "to_device", - f"mpi_buffer({_describe(array)}) staging for recv", - nbytes=_nbytes(array), - ) - # Ensure MPI's host buffer can be reused immediately on context exit. - if transfer_array is not array: - array[...] = transfer_array - if _COUNTERS: - _record_sync("mpi_buffer() staging for recv") - cp.cuda.get_current_stream().synchronize() - - + with sp.region("cunumpy.mpi_buffer.to_device", tags=_PROFILE_TAGS): + transfer_array.set(host) + if _COUNTERS: + _record( + "to_device", + f"mpi_buffer({_describe(array)}) staging for recv", + nbytes=_nbytes(array), + ) + # Ensure MPI's host buffer can be reused immediately on context exit. + if transfer_array is not array: + array[...] = transfer_array + if _COUNTERS: + _record_sync("mpi_buffer() staging for recv") + cp.cuda.get_current_stream().synchronize() + + +@sp.profile("cunumpy.exchange", tags=_PROFILE_TAGS) def exchange( comm: Any, *, diff --git a/src/cunumpy/_optimize.py b/src/cunumpy/_optimize.py index 1bde879..f638e33 100644 --- a/src/cunumpy/_optimize.py +++ b/src/cunumpy/_optimize.py @@ -6,6 +6,9 @@ from collections.abc import Callable from typing import Any, NamedTuple +import scope_profiler as sp + +from cunumpy._profile_tags import _PROFILE_TAGS from cunumpy.xp import array_backend, get_array_backend, get_array_module @@ -35,6 +38,7 @@ def _module(x0: Any) -> Any: return array_backend.xp +@sp.profile("cunumpy.newton", tags=_PROFILE_TAGS) def newton( func: Callable[..., Any], x0: Any, diff --git a/src/cunumpy/_profile_tags.py b/src/cunumpy/_profile_tags.py new file mode 100644 index 0000000..7deda28 --- /dev/null +++ b/src/cunumpy/_profile_tags.py @@ -0,0 +1,9 @@ +"""The scope-profiler tag of cunumpy's own profiling regions. + +``cunumpy/__init__.py`` registers it as opt-in, so the regions record nothing +unless the application enables the ``"cunumpy"`` tag (see +:doc:`/guides/profiling`). +""" + +#: Tags of every cunumpy region (``@sp.profile("cunumpy.", tags=_PROFILE_TAGS)``). +_PROFILE_TAGS = ("cunumpy",) diff --git a/src/cunumpy/_pyccel_cache.py b/src/cunumpy/_pyccel_cache.py index 07cec3a..4fded85 100644 --- a/src/cunumpy/_pyccel_cache.py +++ b/src/cunumpy/_pyccel_cache.py @@ -29,6 +29,10 @@ from pathlib import Path from types import ModuleType +import scope_profiler as sp + +from cunumpy._profile_tags import _PROFILE_TAGS + #: Bumped when the layout of a cached build changes (it was an ``epyccel`` #: folder before the ``pyccel compile`` command line was used). Part of the key. _LAYOUT = "pyccel-cli-1" @@ -245,6 +249,7 @@ def _build(self, module: ModuleType, language: str, staging: Path) -> Path: ) return output + @sp.profile("cunumpy.PyccelCache.compile", tags=_PROFILE_TAGS) def compile(self, module: ModuleType, *, language: str | None = None) -> ModuleType: """Return the compiled form of `module`, building it on first use. diff --git a/src/cunumpy/_scipy_host.py b/src/cunumpy/_scipy_host.py index 0a62dff..123a241 100644 --- a/src/cunumpy/_scipy_host.py +++ b/src/cunumpy/_scipy_host.py @@ -8,8 +8,11 @@ from collections.abc import Callable from typing import Any +import scope_profiler as sp + from cunumpy import xp from cunumpy._host import _to_device +from cunumpy._profile_tags import _PROFILE_TAGS def _scipy(name: str, function: str) -> Any: @@ -79,6 +82,7 @@ def constraints(self, constraints: Any) -> Any: return out +@sp.profile("cunumpy.fsolve", tags=_PROFILE_TAGS) def fsolve( func: Callable[..., Any], x0: Any, @@ -134,6 +138,7 @@ def fsolve( return bridge.result(out) +@sp.profile("cunumpy.root", tags=_PROFILE_TAGS) def root( fun: Callable[..., Any], x0: Any, @@ -190,6 +195,7 @@ def root( return bridge.result(out) +@sp.profile("cunumpy.minimize", tags=_PROFILE_TAGS) def minimize( fun: Callable[..., Any], x0: Any, @@ -266,6 +272,7 @@ def minimize( return bridge.result(out) +@sp.profile("cunumpy.quad", tags=_PROFILE_TAGS) def quad( func: Callable[..., Any], a: float, @@ -313,6 +320,7 @@ def quad( ) +@sp.profile("cunumpy.odeint", tags=_PROFILE_TAGS) def odeint( func: Callable[..., Any], y0: Any, diff --git a/tests/pyodide/run.mjs b/tests/pyodide/run.mjs index 27d9c76..dd337c8 100644 --- a/tests/pyodide/run.mjs +++ b/tests/pyodide/run.mjs @@ -11,18 +11,26 @@ process.on("uncaughtException", (error) => { process.exit(1); }); -const [wheel] = process.argv.slice(2); -if (!wheel || process.argv.length !== 3 || !wheel.endsWith("-py3-none-any.whl")) { - throw new Error("Usage: node run.mjs /path/to/cunumpy-VERSION-py3-none-any.whl"); +// Wheels after cunumpy's are dependencies installed alongside it, e.g. the devel branches of +// the max-models packages, which micropip cannot install from git. +const [wheel, ...dependencies] = process.argv.slice(2); +if (!wheel || ![wheel, ...dependencies].every((path) => path.endsWith("-py3-none-any.whl"))) { + throw new Error( + "Usage: node run.mjs /path/to/cunumpy-VERSION-py3-none-any.whl [dependency wheels...]", + ); } const pyodide = await loadPyodide(); await pyodide.loadPackage(["micropip", "pytest"]); -pyodide.FS.writeFile(`/tmp/${basename(wheel)}`, await readFile(wheel)); -pyodide.globals.set("wheel_uri", `emfs:/tmp/${basename(wheel)}`); +const uris = []; +for (const path of [...dependencies, wheel]) { + pyodide.FS.writeFile(`/tmp/${basename(path)}`, await readFile(path)); + uris.push(`emfs:/tmp/${basename(path)}`); +} +pyodide.globals.set("wheel_uris", pyodide.toPy(uris)); await pyodide.runPythonAsync(` import micropip -await micropip.install(wheel_uri) +await micropip.install(list(wheel_uris)) `); pyodide.FS.mkdirTree("/tests"); pyodide.FS.writeFile( diff --git a/tests/unit/test_opt_in_profiling.py b/tests/unit/test_opt_in_profiling.py new file mode 100644 index 0000000..e4b6c50 --- /dev/null +++ b/tests/unit/test_opt_in_profiling.py @@ -0,0 +1,37 @@ +"""cunumpy's scope-profiler regions: off by default, recorded once the +application enables the ``"cunumpy"`` tag.""" + +import numpy as np +import pytest +import scope_profiler as sp + +import cunumpy as xp + + +@pytest.fixture +def profiler(): + """A fresh scope-profiler state, reset again afterwards so that no + configuration leaks into other tests.""" + sp.reset() + yield + sp.reset() + + +def _cunumpy_regions(**setup_options): + sp.setup(output=None, **setup_options) + keys = np.asarray([2, 0, 1, 0]) + with xp.use_backend("numpy"): + xp.algorithms.sort_by_key(keys, np.arange(4.0)) + xp.kernels.Kernel(lambda a: None, name="noop")(keys) + results = sp.finalize(verbose=False, return_results=True) + return [name for name in results.region_names if name.startswith("cunumpy.")] + + +def test_regions_are_off_by_default(profiler): + assert _cunumpy_regions() == [] + + +@pytest.mark.parametrize("tags", [["cunumpy"], ["*"]]) +def test_enabled_tag_records_regions(profiler, tags): + regions = _cunumpy_regions(enable_tags=tags) + assert sorted(regions) == ["cunumpy.kernel.noop", "cunumpy.sort_by_key"] From 5986211654bec895ee226a6d0fdc5c8b7d41a7fd Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 22 Oct 2026 09:33:51 +0200 Subject: [PATCH 3/9] Check that a new pyccel build imports before caching it, and key builds on NumPy (#114) --- CHANGELOG.md | 8 ++ docs/src/content/docs/kernels/dispatch.md | 8 +- src/cunumpy/_pyccel_cache.py | 97 +++++++++++++++++++---- tests/unit/test_pyccel_cache.py | 46 +++++++++++ 4 files changed, 140 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d0fa9b..741056f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- `PyccelCache` imports a new build once in a subprocess before it stores it: a build the + interpreter cannot import, or one that crashes it (an extension linked to `libpython` in a + Python that has it built in, such as uv's CPython 3.12 on macOS), is a `RuntimeError` and is + not stored, instead of a segmentation fault in every process that loads it. +- The `PyccelCache` key includes the NumPy version, so environments that share an interpreter + but not NumPy no longer share builds. + ### Added - scope-profiler regions around the coarse operations (`Kernel` calls as `cunumpy.kernel.`, CUDA and Pyccel compilation, `sort_by_key`, diff --git a/docs/src/content/docs/kernels/dispatch.md b/docs/src/content/docs/kernels/dispatch.md index 3eabdf0..4770592 100644 --- a/docs/src/content/docs/kernels/dispatch.md +++ b/docs/src/content/docs/kernels/dispatch.md @@ -257,8 +257,8 @@ inside `with xp.kernels.use_host_kernel_implementation("numpy"):` (or set `epyccel` compiles again on every call, in every process. `PyccelCache(name)` compiles a module once with the `pyccel compile` command and keeps the build -on disk, keyed on the module source, the OpenMP setting, the Pyccel and Python -versions and the platform, so only the first run after an edit compiles: +on disk, keyed on the module source, the OpenMP setting, the Pyccel, NumPy and +Python versions and the platform, so only the first run after an edit compiles: - builds go to `~/.cache//kernels`, or the directory in `_KERNEL_CACHE`; if that is not writable, to `/-kernels`; @@ -266,6 +266,10 @@ versions and the platform, so only the first run after an edit compiles: modules that call the OpenMP runtime; - builds are published with an atomic rename, so MPI ranks compiling the same module on a cold cache never load a partial build; +- a new build is imported once in a subprocess before it is published; if that + import fails or crashes the interpreter (an extension linked to `libpython` in + a Python that has it built in), `compile` raises instead and nothing is + stored, so the process that uses the cache does not crash; - the key covers the module's own source only: after editing a module the kernel imports, edit the kernel or clear the cache. diff --git a/src/cunumpy/_pyccel_cache.py b/src/cunumpy/_pyccel_cache.py index 4fded85..563c2b4 100644 --- a/src/cunumpy/_pyccel_cache.py +++ b/src/cunumpy/_pyccel_cache.py @@ -2,10 +2,12 @@ :class:`PyccelCache` builds a module with the ``pyccel compile`` command line interface, in a subprocess, once. It keys a build on the module source, the -target language, the OpenMP setting, the Pyccel and Python versions and the -platform, and stores it under a per-package cache directory. Later processes, -and other MPI ranks, import the stored extension directly; only the first run -after an edit compiles. +target language, the OpenMP setting, the Pyccel, NumPy and Python versions and +the platform, and stores it under a per-package cache directory. Later +processes, and other MPI ranks, import the stored extension directly; only the +first run after an edit compiles. A new build is imported once in a subprocess +before it is stored, so a build this interpreter cannot load is an error, not a +crash of the process that uses it. cunumpy needs no Pyccel to import this module: Pyccel is looked up when a module is compiled, and a missing Pyccel is an :class:`ImportError` there, which @@ -64,20 +66,77 @@ def _writable(directory: Path) -> bool: return os.access(directory, os.W_OK | os.X_OK) -def _load_extension(directory: Path, stem: str) -> ModuleType | None: +def _extension_path(directory: Path, stem: str) -> Path | None: for suffix in importlib.machinery.EXTENSION_SUFFIXES: path = directory / f"{stem}{suffix}" - if not path.is_file(): - continue + if path.is_file(): + return path + return None + + +def _load_extension(directory: Path, stem: str) -> ModuleType | None: + path = _extension_path(directory, stem) + if path is not None: spec = importlib.util.spec_from_file_location(stem, path) - if spec is None or spec.loader is None: - continue - extension = importlib.util.module_from_spec(spec) - spec.loader.exec_module(extension) - return extension + if spec is not None and spec.loader is not None: + extension = importlib.util.module_from_spec(spec) + spec.loader.exec_module(extension) + return extension return None +#: Imports the extension at ``argv[1]`` as module ``argv[2]``. +_IMPORT_CHECK = ( + "import importlib.util, sys\n" + "spec = importlib.util.spec_from_file_location(sys.argv[2], sys.argv[1])\n" + "spec.loader.exec_module(importlib.util.module_from_spec(spec))\n" +) + + +def _check_importable(directory: Path, stem: str) -> None: + """Import a new build in a subprocess; raise if this interpreter cannot. + + A build that crashes the interpreter on import (an extension linked to + ``libpython`` loaded into a Python with ``libpython`` built in, as with some + standalone CPython builds, runs two Pythons in one process) would otherwise + crash every process that loads it from the cache. + """ + path = _extension_path(directory, stem) + if path is None: + raise ImportError(f"pyccel compile wrote no extension module for {stem}") + result = subprocess.run( + [sys.executable, "-c", _IMPORT_CHECK, str(path), stem], + capture_output=True, + text=True, + check=False, + cwd=directory, + ) + if result.returncode == 0: + return + how = ( + f"crashed with signal {-result.returncode}" + if result.returncode < 0 + else f"failed ({result.returncode})" + ) + hint = ( + "\nAn extension linked to libpython crashes a Python that has libpython " + "built in; `otool -L` (macOS) or `ldd` shows what it links." + if result.returncode < 0 + else "" + ) + raise RuntimeError( + f"importing the new build of {stem} with {sys.executable} {how}; it is " + f"not stored:\n{result.stdout}{result.stderr}{hint}" + ) + + +def _numpy_version() -> str: + try: + return importlib.metadata.version("numpy") + except importlib.metadata.PackageNotFoundError: + return "none" + + def _pyccel_version() -> str: try: return importlib.metadata.version("pyccel") @@ -183,9 +242,10 @@ def _openmp_for(self, module: ModuleType) -> bool: def build_key(self, module: ModuleType, language: str | None = None) -> str: """Return the key of `module`'s build, which changes with anything that changes it. - The module's own source, the language, the OpenMP setting, the Pyccel - and Python versions and the platform. Modules the kernel imports are not - part of it: edit the kernel (or clear the cache) after changing one. + The module's own source, the language, the OpenMP setting, the Pyccel, + NumPy and Python versions and the platform. Modules the kernel imports + are not part of it: edit the kernel (or clear the cache) after changing + one. Parameters ---------- @@ -211,6 +271,7 @@ def build_key(self, module: ModuleType, language: str | None = None) -> str: _LAYOUT, f"openmp={self._openmp_for(module)}", _pyccel_version(), + f"numpy={_numpy_version()}", sys.version, platform.platform(), ): @@ -270,8 +331,9 @@ def compile(self, module: ModuleType, *, language: str | None = None) -> ModuleT ImportError If Pyccel is not installed. RuntimeError - With the output of ``pyccel compile`` if the build fails; callers - decide on a fallback. + With the output of ``pyccel compile`` if the build fails, or of + the import if this interpreter cannot import the new build (it is + then not stored); callers decide on a fallback. PermissionError If neither the cache nor the fallback directory is writable. """ @@ -293,6 +355,7 @@ def compile(self, module: ModuleType, *, language: str | None = None) -> ModuleT ).absolute() try: build = self._build(module, language, staging) + _check_importable(build, stem) try: # Atomic, so concurrent builders never see a partial directory. os.rename(build, target) diff --git a/tests/unit/test_pyccel_cache.py b/tests/unit/test_pyccel_cache.py index ce74dfb..8719d69 100644 --- a/tests/unit/test_pyccel_cache.py +++ b/tests/unit/test_pyccel_cache.py @@ -1,5 +1,6 @@ """PyccelCache: compile once, keep the build on disk, reuse it in later processes.""" +import importlib.machinery import importlib.metadata import importlib.util import os @@ -99,6 +100,51 @@ def test_key_changes_with_the_source_and_the_openmp_setting( assert cache.build_key(edited) != key +def test_key_changes_with_the_numpy_version(make_module, cache, monkeypatch): + # a build is compiled against one NumPy's C API + monkeypatch.setattr(_pyccel_cache, "_pyccel_version", lambda: "test") + module = make_module() + monkeypatch.setattr(_pyccel_cache, "_numpy_version", lambda: "2.4.0") + key = cache.build_key(module) + monkeypatch.setattr(_pyccel_cache, "_numpy_version", lambda: "2.5.0") + assert cache.build_key(module) != key + + +def _fake_build(content): + """A ``PyccelCache._build`` that writes `content` as the extension module.""" + + def build(self, module, language, staging): + output = staging / "build" + output.mkdir() + stem = module.__name__.rsplit(".", 1)[-1] + suffix = importlib.machinery.EXTENSION_SUFFIXES[0] + (output / f"{stem}{suffix}").write_bytes(content) + return output + + return build + + +def test_a_build_that_cannot_be_imported_is_not_stored(make_module, cache, monkeypatch): + # no compiler needed: the "build" is a file no interpreter can import + monkeypatch.setattr(_pyccel_cache, "_pyccel_version", lambda: "test") + monkeypatch.setattr(xp.kernels.PyccelCache, "_build", _fake_build(b"not a library")) + with pytest.raises(RuntimeError, match="not stored"): + cache.compile(make_module()) + assert not list(cache.cache_root().iterdir()) + + +def test_a_build_that_crashes_the_interpreter_names_the_signal( + make_module, cache, monkeypatch +): + monkeypatch.setattr(_pyccel_cache, "_pyccel_version", lambda: "test") + monkeypatch.setattr(xp.kernels.PyccelCache, "_build", _fake_build(b"")) + crashed = SimpleNamespace(returncode=-11, stdout="", stderr="Segmentation fault") + monkeypatch.setattr(_pyccel_cache.subprocess, "run", lambda *a, **k: crashed) + with pytest.raises(RuntimeError, match="signal 11(.|\n)*libpython"): + cache.compile(make_module()) + assert not list(cache.cache_root().iterdir()) + + def test_a_module_calling_the_openmp_runtime_keeps_openmp( make_module, cache, monkeypatch ): From bc852eade7e3417fc5c6c7d6b92d4aa8d6eeec59 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 10:10:17 +0200 Subject: [PATCH 4/9] Build the docs with starlight-theme v0.5.0 (#115) --- .github/workflows/docs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 21277ec..8f4a5ad 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -22,7 +22,7 @@ concurrency: jobs: docs: - uses: max-models/starlight-theme/.github/workflows/docs.yml@v0.4.0 + uses: max-models/starlight-theme/.github/workflows/docs.yml@v0.5.0 with: notebooks: false # stable-branch: main # deploy the latest release at // and devel at //dev/ From 18c7c2fbef68a73a32029981662213b73f1d2a21 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 24 Oct 2026 09:09:07 +0200 Subject: [PATCH 5/9] Implemented optional interop with jax, torch mlx (#116) * Implemented optional interop with jax, torch mlx * ruff format * numpydoc --- pyproject.toml | 3 ++ src/cunumpy/__init__.py | 2 + src/cunumpy/__init__.pyi | 1 + src/cunumpy/interop/__init__.py | 9 ++++ src/cunumpy/interop/_common.py | 96 +++++++++++++++++++++++++++++++++ src/cunumpy/interop/dlpack.py | 54 +++++++++++++++++++ src/cunumpy/interop/jax.py | 76 ++++++++++++++++++++++++++ src/cunumpy/interop/mlx.py | 71 ++++++++++++++++++++++++ src/cunumpy/interop/torch.py | 72 +++++++++++++++++++++++++ tests/unit/test_interop.py | 39 ++++++++++++++ 10 files changed, 423 insertions(+) create mode 100644 src/cunumpy/interop/__init__.py create mode 100644 src/cunumpy/interop/_common.py create mode 100644 src/cunumpy/interop/dlpack.py create mode 100644 src/cunumpy/interop/jax.py create mode 100644 src/cunumpy/interop/mlx.py create mode 100644 src/cunumpy/interop/torch.py create mode 100644 tests/unit/test_interop.py diff --git a/pyproject.toml b/pyproject.toml index f73a9f1..dda36b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,9 @@ optional-dependencies.dev = [ # the shared docs workflow registers an ipykernel even without notebooks optional-dependencies.docs = [ "griffe", "ipykernel" ] optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] +optional-dependencies.jax = [ "jax" ] +optional-dependencies.mlx = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] +optional-dependencies.torch = [ "torch" ] optional-dependencies.test = [ "coverage", "pytest", "scipy" ] optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ] urls."Source" = "https://github.com/max-models/cunumpy" diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index d8743db..d240d97 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -22,6 +22,7 @@ arguments, cuda, integrate, + interop, kernels, memory, mpi, @@ -123,6 +124,7 @@ def require_version(minimum: str) -> None: "get_backend", "host_call", "integrate", + "interop", "is_cpu", "is_gpu", "kernels", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index e239928..e89ee4c 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -12,6 +12,7 @@ from cunumpy import algorithms as algorithms from cunumpy import arguments as arguments from cunumpy import cuda as cuda from cunumpy import integrate as integrate +from cunumpy import interop as interop from cunumpy import kernels as kernels from cunumpy import memory as memory from cunumpy import mpi as mpi diff --git a/src/cunumpy/interop/__init__.py b/src/cunumpy/interop/__init__.py new file mode 100644 index 0000000..e7deedb --- /dev/null +++ b/src/cunumpy/interop/__init__.py @@ -0,0 +1,9 @@ +"""Explicit interoperability with scientific machine-learning libraries. + +The adapters are optional and import their framework lazily. They do not +participate in CuNumpy backend selection. +""" + +from cunumpy.interop import dlpack, jax, mlx, torch + +__all__ = ["dlpack", "jax", "mlx", "torch"] diff --git a/src/cunumpy/interop/_common.py b/src/cunumpy/interop/_common.py new file mode 100644 index 0000000..b16baec --- /dev/null +++ b/src/cunumpy/interop/_common.py @@ -0,0 +1,96 @@ +"""Shared implementation for optional framework adapters.""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + +import numpy as np + +from cunumpy._transfers import _ACTIVE, _describe, _nbytes, _record + + +def require(module: str, *, extra: str | None = None) -> Any: + """Import an optional framework and give a consistent installation error. + + Parameters + ---------- + module : str + Import name of the optional framework. + extra : str or None, optional + CuNumpy extra name, inferred from `module` when omitted. + + Returns + ------- + module + Imported framework module. + """ + try: + return import_module(module) + except ImportError as error: + extra = module.split(".", 1)[0] if extra is None else extra + raise ImportError( + f"{module} interoperability needs {module.split('.', 1)[0]!r}; " + f"install cunumpy[{extra}]" + ) from error + + +def record(source: Any, target: Any) -> None: + """Record an adapter conversion when transfer profiling is active. + + Parameters + ---------- + source, target : object + Source and converted arrays. + """ + if _ACTIVE: + _record( + "device_copy", + f"interop {type(source).__name__} -> {type(target).__name__} " + f"({_describe(target)})", + nbytes=_nbytes(target), + ) + + +def target_backend(backend: str) -> Any: + """Load a NumPy or CuPy target module. + + Parameters + ---------- + backend : {"numpy", "cupy"} + Name of the target array library. + + Returns + ------- + module + The requested array module. + """ + if backend == "numpy": + return np + if backend == "cupy": + try: + import cupy + except ImportError as error: + raise ImportError( + "CuPy interoperability needs CuPy; install a compatible CuPy release" + ) from error + return cupy + raise ValueError(f"backend must be 'numpy' or 'cupy', got {backend!r}") + + +def cast(array: Any, dtype: Any) -> Any: + """Cast an array only when a dtype was requested. + + Parameters + ---------- + array : object + Array to cast. + dtype : object or None + Requested dtype, or None to preserve the input. + + Returns + ------- + object + The original or cast array. + """ + return array if dtype is None else array.astype(dtype) diff --git a/src/cunumpy/interop/dlpack.py b/src/cunumpy/interop/dlpack.py new file mode 100644 index 0000000..4258974 --- /dev/null +++ b/src/cunumpy/interop/dlpack.py @@ -0,0 +1,54 @@ +"""Small explicit helpers around the DLPack array interchange protocol.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import record, target_backend + + +def to_dlpack(array: Any) -> Any: + """Return a DLPack capsule or protocol object for `array`. + + Parameters + ---------- + array : object + Array implementing the DLPack protocol. + + Returns + ------- + object + The DLPack capsule or protocol object. + """ + try: + return array.__dlpack__() + except AttributeError as error: + raise TypeError( + f"{type(array).__name__} does not implement the DLPack protocol" + ) from error + + +def from_dlpack(source: Any, *, backend: str = "numpy") -> Any: + """Construct a NumPy or CuPy array from a DLPack source. + + Parameters + ---------- + source : object + DLPack capsule or object implementing the protocol. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + + Returns + ------- + object + The constructed array. + """ + target = target_backend(backend) + if not hasattr(target, "from_dlpack"): + raise TypeError(f"{backend} does not provide from_dlpack") + result = target.from_dlpack(source) + record(source, result) + return result + + +__all__ = ["from_dlpack", "to_dlpack"] diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py new file mode 100644 index 0000000..98393f4 --- /dev/null +++ b/src/cunumpy/interop/jax.py @@ -0,0 +1,76 @@ +"""Optional interoperability with JAX arrays.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to a JAX array. + + DLPack is used when available; `device` and `dtype` are explicit optional + conversions. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + JAX device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + JAX array. + """ + jax = require("jax") + import jax.numpy as jnp + + if type(array).__module__.startswith("jax"): + result = array + elif hasattr(array, "__dlpack__"): + result = jax.dlpack.from_dlpack(array) + else: + result = jnp.asarray(array) + result = cast(result, dtype) + if device is not None: + result = jax.device_put(result, device) + record(array, result) + return result + + +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: + """Convert a JAX array to a NumPy or CuPy array. + + Parameters + ---------- + array : object + JAX array to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ + target = target_backend(backend) + if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): + result = target.from_dlpack(array) + else: + result = target.asarray(array) + result = cast(result, dtype) + record(array, result) + return result + + +to_jax = to_framework +from_jax = from_framework + +__all__ = ["from_framework", "from_jax", "to_framework", "to_jax"] diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py new file mode 100644 index 0000000..976978c --- /dev/null +++ b/src/cunumpy/interop/mlx.py @@ -0,0 +1,71 @@ +"""Optional interoperability with Apple's MLX arrays.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to an MLX array. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + MLX device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + MLX array. + """ + mlx = require("mlx.core", extra="mlx") + if type(array).__module__.startswith("mlx"): + result = array + elif hasattr(mlx, "from_dlpack") and hasattr(array, "__dlpack__"): + result = mlx.from_dlpack(array) + else: + result = mlx.array(array) + result = cast(result, dtype) + if device is not None: + result = mlx.array(result, device=device) + record(array, result) + return result + + +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: + """Convert an MLX array to a NumPy or CuPy array. + + Parameters + ---------- + array : object + MLX array to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ + target = target_backend(backend) + if hasattr(array, "__dlpack__") and hasattr(target, "from_dlpack"): + result = target.from_dlpack(array) + else: + result = target.asarray(array) + result = cast(result, dtype) + record(array, result) + return result + + +to_mlx = to_framework +from_mlx = from_framework + +__all__ = ["from_framework", "from_mlx", "to_framework", "to_mlx"] diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py new file mode 100644 index 0000000..6b77123 --- /dev/null +++ b/src/cunumpy/interop/torch.py @@ -0,0 +1,72 @@ +"""Optional interoperability with PyTorch tensors.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to a PyTorch tensor. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + PyTorch device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + PyTorch tensor. + """ + torch = require("torch") + if type(array).__module__.startswith("torch"): + result = array + elif hasattr(torch, "from_dlpack") and hasattr(array, "__dlpack__"): + result = torch.from_dlpack(array) + else: + result = torch.as_tensor(array) + if dtype is not None: + result = result.to(dtype=dtype) + if device is not None: + result = result.to(device=device) + record(array, result) + return result + + +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: + """Convert a PyTorch tensor to a NumPy or CuPy array. + + Parameters + ---------- + array : object + PyTorch tensor to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ + target = target_backend(backend) + if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): + result = target.from_dlpack(array) + else: + result = target.asarray(array.detach().cpu().numpy()) + result = cast(result, dtype) + record(array, result) + return result + + +to_torch = to_framework +from_torch = from_framework + +__all__ = ["from_framework", "from_torch", "to_framework", "to_torch"] diff --git a/tests/unit/test_interop.py b/tests/unit/test_interop.py new file mode 100644 index 0000000..22442ff --- /dev/null +++ b/tests/unit/test_interop.py @@ -0,0 +1,39 @@ +"""Tests for optional framework interoperability.""" + +import types + +import numpy as np +import pytest + +from cunumpy.interop import dlpack +from cunumpy.interop._common import require, target_backend + + +def test_dlpack_helpers_use_array_protocol(monkeypatch): + class Source: + def __dlpack__(self): + return "capsule" + + assert dlpack.to_dlpack(Source()) == "capsule" + + target = types.SimpleNamespace( + from_dlpack=lambda source: np.array([source.__dlpack__()]), + ) + monkeypatch.setattr("cunumpy.interop.dlpack.target_backend", lambda backend: target) + result = dlpack.from_dlpack(Source()) + np.testing.assert_array_equal(result, ["capsule"]) + + +def test_dlpack_rejects_objects_without_protocol(): + with pytest.raises(TypeError, match="DLPack protocol"): + dlpack.to_dlpack(object()) + + +def test_optional_dependency_error_is_actionable(): + with pytest.raises(ImportError, match="cunumpy\\[not-installed-framework\\]"): + require("not-installed-framework") + + +def test_target_backend_rejects_unknown_backend(): + with pytest.raises(ValueError, match="backend must be"): + target_backend("jax") From f456ff4ad2500d398986e9c5a54a4ad10db9dfb7 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 27 Oct 2026 14:47:47 +0100 Subject: [PATCH 6/9] Use ruff instead of black and isort (#117) * Use ruff for formatting and import sorting, with the shared pre-commit config * Remove the mypy check --- .gitattributes | 2 +- .github/workflows/gpu_ci_trigger.yml | 2 +- .github/workflows/publish.yml | 1 - .github/workflows/static_analysis.yml | 13 +- .pre-commit-config.yaml | 30 +- CHANGELOG.md | 2 +- pytest.ini | 2 +- src/cunumpy/_algorithms.py | 30 +- src/cunumpy/_cuda_kernel.py | 634 +++++++++++++------------- src/cunumpy/_dispatch.py | 98 ++-- src/cunumpy/_fake_cupy.py | 20 +- src/cunumpy/_fake_cupy_impl.py | 146 +++--- src/cunumpy/_fake_cupyx.py | 72 +-- src/cunumpy/_kernel.py | 116 ++--- src/cunumpy/_metal_kernel.py | 24 +- src/cunumpy/_mirror.py | 48 +- src/cunumpy/_pyccel_cache.py | 6 +- src/cunumpy/_random_streams.py | 16 +- src/cunumpy/_scipy_backend.py | 50 +- src/cunumpy/_staging.py | 12 +- src/cunumpy/_streams.py | 12 +- src/cunumpy/_transfers.py | 33 +- src/cunumpy/kernel_testing.py | 20 +- src/cunumpy/xp.py | 112 ++--- testing/smoke_wheel.py | 1 + tests/unit/test_cuda_kernel.py | 88 ++-- tests/unit/test_mpi_serial.py | 6 +- tests/unit/test_staging.py | 12 +- tests/unit/test_transfers.py | 14 +- 29 files changed, 818 insertions(+), 804 deletions(-) diff --git a/.gitattributes b/.gitattributes index 2f41707..8425c58 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1 @@ -*.ipynb filter=strip-notebook-output +*.ipynb filter=strip-notebook-output diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 80d4415..2bc5a14 100644 --- a/.github/workflows/gpu_ci_trigger.yml +++ b/.github/workflows/gpu_ci_trigger.yml @@ -1,7 +1,7 @@ # SETUP INSTRUCTIONS: # ------------------ # This workflow synchronizes the code to GitLab via SSH to trigger GPU-enabled CI. -# +# # 1. GENERATE SSH KEY PAIR (on your local machine): # ssh-keygen -t ed25519 -f ~/.ssh/gitlab_sync_key -N "" -C "github-to-gitlab-sync" # diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 2e22dd2..d272aff 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -45,4 +45,3 @@ jobs: TWINE_USERNAME: __token__ # Use API token TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} run: twine upload dist/* - diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index ae0eb4d..b7f3a4c 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -37,19 +37,18 @@ jobs: - name: Code formatting with ruff run: | pip install ruff - ruff format --check src/ tests/ + ruff format --check . - mypy: + ssort: runs-on: ubuntu-latest - continue-on-error: true steps: - name: Checkout the code uses: actions/checkout@v4 - - name: Type checking with mypy + - name: Statement order with ssort run: | - pip install mypy - mypy src/ || true + pip install ssort + ssort --check --diff . prospector: runs-on: ubuntu-latest @@ -72,7 +71,7 @@ jobs: - name: Linting with ruff (including import sorting) run: | pip install ruff - ruff check src/ tests/ + ruff check . numpydoc: runs-on: ubuntu-latest diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index af67085..2effd49 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,15 +1,27 @@ +--- repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 # Use the latest stable version + rev: v6.0.0 hooks: + - id: check-yaml + - id: check-toml + - id: end-of-file-fixer + - id: trailing-whitespace - id: check-added-large-files # Prevent giant files from being committed. - args: ["--maxkb=1000"] - - id: check-merge-conflict # Check for files that contain merge conflict strings. - - id: check-toml # Attempts to load all TOML files to verify syntax. - - id: check-yaml # Attempts to load all yaml files to verify syntax. - args: ["--unsafe"] - + args: [--maxkb=500] + exclude: ^(uv\.lock|docs/package-lock\.json)$ # generated lock files + - id: check-merge-conflict # Check for files that contain merge conflict strings. - repo: https://github.com/kynan/nbstripout - rev: 0.8.1 + rev: 0.9.1 hooks: - - id: nbstripout # remove jupyter notebook cell output + - id: nbstripout # strip jupyter notebook output + - repo: https://github.com/bwhmather/ssort + rev: 0.17.0 + hooks: + - id: ssort # order statements so names are defined before they are used + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.17.3 + hooks: + - id: ruff-check # lint, including import sorting (the I rules) + args: [--fix] + - id: ruff-format diff --git a/CHANGELOG.md b/CHANGELOG.md index 741056f..ab8d227 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -457,7 +457,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - **Dynamic Dispatch Architecture**: Refactored `src/cunumpy/xp.py` to use module-level `__getattr__`. This ensures that `cunumpy.` calls always resolve to the currently active backend module, enabling seamless runtime switching via `set_backend`. - **Type Safety**: Updated `src/cunumpy/__init__.pyi` stubs to provide full IDE autocompletion and type-checking for all new API methods. -- **Documentation**: +- **Documentation**: - Simplified `README.md` and documentation to exclusively focus on PyPI installation (`pip install cunumpy`). - Enhanced `quickstart.md` and `api.md` with usage examples for the new backend control and synchronization features. - **CI/CD**: Restricted GitHub Pages documentation deployment to the `devel` branch only. diff --git a/pytest.ini b/pytest.ini index fa4a21d..fcccae1 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,2 @@ -[pytest] +[pytest] pythonpath = src diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index 76355df..15876e7 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -374,6 +374,21 @@ def _radix_argsort(keys: np.ndarray) -> np.ndarray: return order.astype(np.int64, copy=False) +def _take_along_axis(xpm: Any, array: Any, order: Any, axis: int) -> Any: + """Return `array` taken at `order` along `axis`, as a new C-contiguous array. + + A component-major ``(ncomp, n)`` NumPy array is gathered one contiguous row at a + time, which is about twice as fast as ``take`` along its last axis; everything + else, and CuPy, goes through ``take``. + """ + if xpm is np and array.ndim == 2 and axis == 1: + gathered = np.empty((array.shape[0], order.shape[0]), dtype=array.dtype) + for row in range(array.shape[0]): + gathered[row] = array[row][order] + return gathered + return xpm.take(array, order, axis=axis) + + @sp.profile("cunumpy.sort_by_key", tags=_PROFILE_TAGS) def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]: """Sort `keys` and reorder every array the same way, in one stable argsort. @@ -449,21 +464,6 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]: ) -def _take_along_axis(xpm: Any, array: Any, order: Any, axis: int) -> Any: - """Return `array` taken at `order` along `axis`, as a new C-contiguous array. - - A component-major ``(ncomp, n)`` NumPy array is gathered one contiguous row at a - time, which is about twice as fast as ``take`` along its last axis; everything - else, and CuPy, goes through ``take``. - """ - if xpm is np and array.ndim == 2 and axis == 1: - gathered = np.empty((array.shape[0], order.shape[0]), dtype=array.dtype) - for row in range(array.shape[0]): - gathered[row] = array[row][order] - return gathered - return xpm.take(array, order, axis=axis) - - @sp.profile("cunumpy.compact_by_mask", tags=_PROFILE_TAGS) def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int: """Move the entries where `mask` is True to the front of every array, in place. diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 3e72f02..e6ae384 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -584,98 +584,6 @@ def _template_arg(value: Any) -> str: return ctype_of(value) -def parse_cuda_signature( - source: str, - name: str, - *, - structs: Iterable[CudaStruct] = (), - template_args: Sequence[Any] | None = None, -) -> tuple[CudaParameter, ...]: - """Parse the parameters of a ``__global__`` function. - - C types map to dtypes as on Linux (LP64): ``int`` is int32, ``long`` and - ``long long`` are int64. - - Parameters - ---------- - source : str - CUDA C source code. - name : str - Name of the ``__global__`` function. - structs : iterable of CudaStruct, optional - Struct types that may appear as parameters (by value). A definition - of the struct in `source` must match. - template_args : sequence, optional - Template arguments if `name` is a function template: C types or - dtypes (see :func:`~cunumpy.cuda.ctype_of`) for type parameters, - ints or bools for non-type parameters. - - Returns - ------- - tuple of CudaParameter - The parameters, in order. - - Raises - ------ - ValueError - If there is no such function, `template_args` do not fit, a struct - definition does not match its :class:`~cunumpy.arguments.CudaStruct`, - or a parameter type cannot be checked (e.g. a macro or ``double**``). - - Examples - -------- - >>> source = ''' - ... template - ... __global__ void scale(T* x, T factor, int n) {}''' - >>> [(p.name, p.ctype, p.pointer) for p in xp.cuda.parse_cuda_signature( - ... source, "scale", template_args=[np.float32])] - [('x', 'float', True), ('factor', 'float', False), ('n', 'int', False)] - """ - code = _strip_comments(source) - structs = {s.name: s for s in structs} - for struct in structs.values(): - struct.check_source(code) - - pattern = r"(?:template\s*<(?P