Opt-in scope-profiler regions, off by default

.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,11 +24,15 @@ 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
         run: python -m pytest -q tests/portable
       - 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)

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.<name>`, 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,

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.<name>`), 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

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

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

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."""

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.
 

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))

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.
 

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,

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",