Opt-in scope-profiler regions, off by default

CHANGELOG.md

@@ -80,6 +80,12 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   the frames of an animation. The tutorial publisher turns GIF outputs into animated images.
 - `Patch.local_cell_centers(hierarchy, name)` takes the field whose cells are wanted; without
   one, the first cell-centered field (it used the first field of any centering).
+- Profiling regions, off by default: halo exchanges, restriction, deposit, migration, regridding,
+  tagging and both Poisson solvers record `patchwork.<operation>` scope-profiler regions tagged
+  `patchwork`, an opt-in tag that an application enables with `enable_tags=["patchwork"]`,
+  `SCOPE_PROFILER_ENABLE_TAGS=patchwork` or `scope-profiler run --enable-tag patchwork`.
+  `scope-profiler>=0.7.5` is a new dependency, which uv takes from its `devel` branch until
+  0.7.5 is released.
 
 ### Changed
 

README.md

@@ -151,6 +151,27 @@ needs 0.9 s to factorize; the two potentials agree to 0.13%.
 ![The potential of a moving charge, with the patches that follow
 it](docs/public/figures/regrid_moving_charge.png)
 
+## Profiling
+
+patchwork ships
+[scope-profiler](https://github.com/max-models/scope-profiler) regions
+around its expensive operations (halo exchanges, restriction, deposit
+and migration, tagging and regridding, building and solving the Poisson
+problems), named `patchwork.<operation>`. They are off by default, so an
+application that profiles itself records only its own regions. Turn them
+on with the `patchwork` tag:
+
+``` python
+import scope_profiler as sp
+
+sp.setup(enable_tags=["patchwork"])  # or sp.session(enable_tags=["patchwork"])
+```
+
+or, without editing the code, `SCOPE_PROFILER_ENABLE_TAGS=patchwork`,
+`scope-profiler run --enable-tag patchwork`, or `enable_tags` in the
+TOML configuration. Switched off, a region costs less than 0.2 µs per
+call.
+
 ## Install
 
 With [uv](https://docs.astral.sh/uv/):

README.qmd

@@ -147,6 +147,25 @@ uniform mesh, which also needs 0.9 s to factorize; the two potentials agree to
 
 ![The potential of a moving charge, with the patches that follow it](docs/public/figures/regrid_moving_charge.png)
 
+## Profiling
+
+patchwork ships [scope-profiler](https://github.com/max-models/scope-profiler)
+regions around its expensive operations (halo exchanges, restriction, deposit and
+migration, tagging and regridding, building and solving the Poisson problems), named
+`patchwork.<operation>`. They are off by default, so an application that
+profiles itself records only its own regions. Turn them on with the `patchwork`
+tag:
+
+```python
+import scope_profiler as sp
+
+sp.setup(enable_tags=["patchwork"])  # or sp.session(enable_tags=["patchwork"])
+```
+
+or, without editing the code, `SCOPE_PROFILER_ENABLE_TAGS=patchwork`,
+`scope-profiler run --enable-tag patchwork`, or `enable_tags` in the TOML
+configuration. Switched off, a region costs less than 0.2 µs per call.
+
 ## Install
 
 With [uv](https://docs.astral.sh/uv/):

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
   "mpiarray>=0.2.1",
   "numpy",
   "scipy",
+  "scope-profiler>=0.7.5",
 ]
 
 [project.optional-dependencies]
@@ -95,3 +96,4 @@ show_missing = true
 [tool.uv.sources]
 maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" }
 mpiarray = { git = "https://github.com/max-models/mpiarray", branch = "devel" }
+scope-profiler = { git = "https://github.com/max-models/scope-profiler", branch = "devel" }

src/patchwork/__init__.py

@@ -1,5 +1,8 @@
 """Patchwork: block-structured adaptive meshes backed by mpiarray."""
 
+import scope_profiler as sp
+
+from patchwork._profiling import PROFILE_TAGS
 from patchwork.mesh import AMRHierarchy, IndexBox, Patch
 from patchwork.nodal import (
     CompositeNodes,
@@ -13,6 +16,9 @@
 from patchwork.poisson import solve_poisson
 from patchwork.regrid import cluster, function_tags, gradient_tags
 
+# patchwork's profiling regions are off unless an application enables the tag
+sp.register_opt_in_tag(*PROFILE_TAGS)
+
 __version__ = "0.1.0"
 
 __all__ = [

src/patchwork/_profiling.py

@@ -0,0 +1,10 @@
+"""The tag of patchwork's own profiling regions, which are off unless enabled.
+
+Every region is named ``"patchwork.<operation>"`` and tagged with
+:data:`PROFILE_TAGS`. The package registers the tag as opt-in with
+scope-profiler, so an application that profiles itself records these regions
+only when it asks for them (``enable_tags=["patchwork"]``; see the README).
+"""
+
+#: The scope-profiler tags of every patchwork region.
+PROFILE_TAGS = ("patchwork",)

src/patchwork/mesh.py

@@ -6,8 +6,11 @@
 from typing import TYPE_CHECKING, Any
 
 import mpiarray as mpa
+import scope_profiler as sp
 from maybempi import MPI
 
+from patchwork._profiling import PROFILE_TAGS
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterator, Mapping, Sequence
 
@@ -545,6 +548,7 @@ def allocate_field(
             if patch.is_local:
                 patch._fields[name] = arrays[id(patch)]
 
+    @sp.profile("patchwork.exchange_halos", tags=PROFILE_TAGS)
     def exchange_halos(self, name: str) -> None:
         """Fill the halo of every patch of field ``name`` from the rest of the hierarchy.
 
@@ -559,6 +563,7 @@ def exchange_halos(self, name: str) -> None:
 
         self._schedule("halo", name, HaloExchange).execute()
 
+    @sp.profile("patchwork.restrict", tags=PROFILE_TAGS)
     def restrict(self, name: str) -> None:
         """Give the points of field ``name`` under finer patches the fine values.
 
@@ -571,6 +576,7 @@ def restrict(self, name: str) -> None:
 
         self._schedule("restrict", name, Restriction).execute()
 
+    @sp.profile("patchwork.regrid", tags=PROFILE_TAGS)
     def regrid(
         self,
         tags: Mapping[int, Any],
@@ -835,6 +841,7 @@ def owner(self, x: Any, y: Any, name: str) -> Any:
         found, cell_i, cell_j, _ = self._cells(x, y)
         return self._owners(found, cell_i, cell_j, name).reshape(shape)
 
+    @sp.profile("patchwork.migrate", tags=PROFILE_TAGS)
     def migrate(
         self, x: Any, y: Any, *arrays: Any, name: str, outside: str = "error"
     ) -> tuple[Any, ...]:
@@ -880,6 +887,7 @@ def migrate(
         )
         return tuple(np.asarray(array) for array in received)
 
+    @sp.profile("patchwork.deposit", tags=PROFILE_TAGS)
     def deposit(
         self, x: Any, y: Any, charge: Any, name: str, *, reset: bool = True
     ) -> None:
@@ -949,6 +957,7 @@ def deposit(
             )
             storage += added.reshape(storage.shape)
 
+    @sp.profile("patchwork.accumulate_halos", tags=PROFILE_TAGS)
     def accumulate_halos(self, name: str) -> None:
         """Bring deposited charge to the composite nodes: the scatter-add after :meth:`deposit`.
 
@@ -985,7 +994,8 @@ def _schedule(self, kind: str, name: str, build: Any) -> Any:
             raise KeyError(f"hierarchy has no field {name!r}")
         key = (kind, name)
         if key not in self._schedules:
-            self._schedules[key] = build(self, name)
+            with sp.region("patchwork.build_schedule", tags=PROFILE_TAGS):
+                self._schedules[key] = build(self, name)
         return self._schedules[key]
 
     def __iter__(self) -> Iterator[Patch]:

src/patchwork/nodal.py

@@ -25,6 +25,9 @@
 import numpy as np
 import scipy.sparse as sparse
 import scipy.sparse.linalg
+import scope_profiler as sp
+
+from patchwork._profiling import PROFILE_TAGS
 
 if TYPE_CHECKING:
     from patchwork.mesh import AMRHierarchy, Patch
@@ -106,6 +109,7 @@ class CompositeNodes:
             ``u`` from the others (which it leaves unchanged).
     """
 
+    @sp.profile("patchwork.composite_nodes", tags=PROFILE_TAGS)
     def __init__(self, hierarchy: AMRHierarchy) -> None:
         """Number the nodes of ``hierarchy`` and find its leaf cells and hanging nodes."""
         self.hierarchy = hierarchy
@@ -702,6 +706,7 @@ class NodalPoisson:
         fixed: Whether a node has a prescribed value (a face or a body).
     """
 
+    @sp.profile("patchwork.nodal_poisson_setup", tags=PROFILE_TAGS)
     def __init__(
         self,
         hierarchy: AMRHierarchy,
@@ -1056,6 +1061,7 @@ def constant(self) -> np.ndarray:
         """The share of the right-hand side that the prescribed values give."""
         return self._constant
 
+    @sp.profile("patchwork.nodal_poisson_solve_reduced", tags=PROFILE_TAGS)
     def solve_reduced(self, right_hand_side: np.ndarray) -> np.ndarray:
         """Solve for a right-hand side of the unknowns; the potential on all nodes.
 
@@ -1071,6 +1077,7 @@ def solve_reduced(self, right_hand_side: np.ndarray) -> np.ndarray:
             potential[nodes.hanging] = self._reduction_rows @ potential
         return potential
 
+    @sp.profile("patchwork.nodal_poisson_solve", tags=PROFILE_TAGS)
     def solve(
         self,
         source: np.ndarray | str | None = None,

src/patchwork/poisson.py

@@ -9,6 +9,9 @@
 import numpy as np
 import scipy.sparse as sparse
 import scipy.sparse.linalg
+import scope_profiler as sp
+
+from patchwork._profiling import PROFILE_TAGS
 
 if TYPE_CHECKING:
     from patchwork.mesh import AMRHierarchy, Patch
@@ -51,6 +54,7 @@ def _parents(
     return parents
 
 
+@sp.profile("patchwork.solve_poisson", tags=PROFILE_TAGS)
 def solve_poisson(
     hierarchy: AMRHierarchy,
     rhs: str,

src/patchwork/regrid.py

@@ -26,8 +26,10 @@
 
 import mpiarray as mpa
 import numpy as np
+import scope_profiler as sp
 from maybempi import MPI
 
+from patchwork._profiling import PROFILE_TAGS
 from patchwork.halos import (
     _overlap,
     _Points,
@@ -70,6 +72,7 @@ def function_tags(
     return tags
 
 
+@sp.profile("patchwork.gradient_tags", tags=PROFILE_TAGS)
 def gradient_tags(
     hierarchy: AMRHierarchy, name: str, threshold: float
 ) -> dict[int, np.ndarray]:
@@ -494,6 +497,7 @@ def new_boxes(
     return boxes
 
 
+@sp.profile("patchwork.regrid_transfer", tags=PROFILE_TAGS)
 def transfer(old: AMRHierarchy, new: AMRHierarchy, name: str) -> None:
     """Fill field ``name`` of ``new`` from ``old``, level by level (collective).
 

src/patchwork/tests/unit/test_profiling.py

@@ -0,0 +1,45 @@
+"""Patchwork's profiling regions: off by default, on with ``enable_tags``."""
+
+from __future__ import annotations
+
+import pytest
+import scope_profiler as sp
+
+from patchwork import AMRHierarchy
+
+
+@pytest.fixture(autouse=True)
+def _fresh_profiler():
+    """No profiling state from other tests, and none left for them."""
+    sp.reset()
+    yield
+    sp.reset()
+
+
+def regions_of_a_halo_exchange(**setup_kwargs) -> set[str]:
+    """The names of the regions recorded around one halo exchange."""
+    sp.ProfileManager.setup(output=None, **setup_kwargs)
+    mesh = AMRHierarchy(base_shape=(8, 8))
+    mesh.add_patch(level=1, box=((4, 12), (4, 12)))
+    mesh.allocate_field("u", halo=1)
+    with sp.region("user.step"):
+        mesh.exchange_halos("u")
+    results = sp.finalize(verbose=False, return_results=True)
+    return set(results.region_names)
+
+
+def test_patchwork_regions_are_off_by_default():
+    names = regions_of_a_halo_exchange()
+
+    assert "user.step" in names
+    assert not {name for name in names if name.startswith("patchwork.")}
+
+
+def test_enable_tags_records_the_patchwork_regions():
+    names = regions_of_a_halo_exchange(enable_tags=["patchwork"])
+
+    assert {
+        "user.step",
+        "patchwork.exchange_halos",
+        "patchwork.build_schedule",
+    } <= names