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/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/ diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 7f7a8e7..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" # @@ -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/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/pyodide.yml b/.github/workflows/pyodide.yml index dc35744..8ba8ea2 100644 --- a/.github/workflows/pyodide.yml +++ b/.github/workflows/pyodide.yml @@ -6,6 +6,8 @@ on: pull_request: branches: [main, devel] +# 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 @@ -22,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 @@ -29,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/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 02b58a7..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 @@ -82,6 +81,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/.pre-commit-config.yaml b/.pre-commit-config.yaml index af67085..9fd7df5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,15 +1,33 @@ +--- 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 + - repo: https://github.com/numpy/numpydoc + rev: v1.11.0 + hooks: + - id: numpydoc-validation # docstring sections, configured in pyproject.toml + files: ^src/cunumpy/ + exclude: _fake_cupy diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b46bef..17305bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.7.2] - 2026-10-27 + +### 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 +- `xp.kernels.KernelPackage(name)`: the `PyccelCache` and the `Kernel.from_folder` options + shared by the kernel folders of one package (`host_suffix="_pyccel"`, `dispatch="arrays"`, + `block_size=256`, `n_threads_from="last_axis"`), with `from_folder(package, **overrides)` + and `metal_kernel(...)` (float32 `MetalKernel`s with zeroed outputs), so packages no longer + keep their own copies of these settings. +- `xp.kernels.source_roots(*packages)`: the source roots of installed packages, for + `include_dirs` of a CUDA kernel that includes another package's header. +- `xp.cuda.shared_memory_fits(num_values, itemsize=8)`: whether a per-block buffer fits + `DEFAULT_SHARED_MEMORY_PER_BLOCK`. +- 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, @@ -440,7 +468,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/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/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/docs/src/content/docs/kernels/dispatch.md b/docs/src/content/docs/kernels/dispatch.md index 3eabdf0..62a5e33 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. @@ -282,6 +286,43 @@ Packages that compile their kernels at install time with the `pyccel` command do not need `compile_host`: the compiled modules are imported like the Python ones. +### One set of options per package + +A package with many kernel folders passes the same options to every +`from_folder()`. A `KernelPackage` holds them, with the package's +`PyccelCache`, so they are declared once: + +```python +# my_sim/kernels/__init__.py +KERNEL_PACKAGE = xp.kernels.KernelPackage("my-sim") + +# my_sim/kernels/push/__init__.py +from my_sim.kernels import KERNEL_PACKAGE + +kernel = KERNEL_PACKAGE.from_folder(__name__) + +# my_sim/kernels/push/push_metal.py +_KERNEL = KERNEL_PACKAGE.metal_kernel("push", SOURCE, ["x", "v"], ["x_new"]) +``` + +The defaults are `host_suffix="_pyccel"`, `dispatch="arrays"` (see below), +`block_size=256` and `n_threads_from="last_axis"` (the marker count of a +component-major `(ncomp, N)` array). A folder overrides any of them in its own +call, e.g. `KERNEL_PACKAGE.from_folder(__name__, n_threads_from=rows)`; +`KERNEL_PACKAGE.options` are the same options for +`KernelCatalog.from_package()`. `metal_kernel()` builds a `MetalKernel` with +`threadgroup=256`, `float64="cast"` and `init_value=0`. + +A CUDA source that includes another package's header (`#include +"other_pkg/kernels/cell.cuh"`) needs that package's source root as an include +directory; `xp.kernels.source_roots()` finds it without importing the package: + +```python +kernel = KERNEL_PACKAGE.from_folder( + __name__, include_dirs=xp.kernels.source_roots("my_sim", "other_pkg") +) +``` + ## Choosing the kernel by where the arrays are `Kernel` picks the CUDA kernel on the CuPy backend. Real codes also hand host diff --git a/docs/src/content/docs/reference/overview.md b/docs/src/content/docs/reference/overview.md index ffe525f..f084b3a 100644 --- a/docs/src/content/docs/reference/overview.md +++ b/docs/src/content/docs/reference/overview.md @@ -47,7 +47,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | Submodule | Backends | Contents | |---|---|---| | [`cunumpy`](/cunumpy/api/cunumpy/) | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `as_device_array`, `to_host_async`, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` | -| [`cunumpy.kernels`](/cunumpy/api/cunumpy/kernels/) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, `PyccelCache`, host implementations, `as_kernel_array`, `kernel_output`, `kernel_outputs`, `fuse` | +| [`cunumpy.kernels`](/cunumpy/api/cunumpy/kernels/) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, `PyccelCache`, `KernelPackage`, `source_roots`, host implementations, `as_kernel_array`, `kernel_output`, `kernel_outputs`, `fuse` | | [`cunumpy.arguments`](/cunumpy/api/cunumpy/arguments/) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | | [`cunumpy.cuda`](/cunumpy/api/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`](/cunumpy/api/cunumpy/rng/) | both | `random_streams`, `get_rng`, `philox_*` | diff --git a/pyproject.toml b/pyproject.toml index f70ac73..f26fe5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.7.1" +version = "0.7.2" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] @@ -25,6 +25,7 @@ dependencies = [ "array-api-compat", "maybempi>=0.1.2", "numpy", + "scope-profiler>=0.7.5", ] optional-dependencies.dev = [ @@ -37,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" @@ -47,6 +51,12 @@ 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" } +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 extend-exclude = [ ".claude" ] 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/requirements-devel.txt b/requirements-devel.txt new file mode 100644 index 0000000..84b537c --- /dev/null +++ b/requirements-devel.txt @@ -0,0 +1,5 @@ +# 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 +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..d240d97 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -15,11 +15,14 @@ import re as _re from importlib.metadata import PackageNotFoundError, version +import scope_profiler as _sp + from cunumpy import ( algorithms, arguments, cuda, integrate, + interop, kernels, memory, mpi, @@ -30,6 +33,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 +58,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.""" @@ -116,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/_algorithms.py b/src/cunumpy/_algorithms.py index d7e1621..15876e7 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,22 @@ 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. @@ -445,21 +464,7 @@ 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 2db6d94..e6ae384 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 @@ -582,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