From 7def7bd419a83b1ae15cbf7b534f6895362dbba3 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 22 Aug 2026 22:15:44 +0200 Subject: [PATCH 1/3] Add autoparallization --- src/struphy/topology/__init__.py | 13 ++ src/struphy/topology/domain_decomposition.py | 169 ++++++++++++++++++ src/struphy/topology/tests/__init__.py | 0 .../tests/test_domain_decomposition.py | 107 +++++++++++ 4 files changed, 289 insertions(+) create mode 100644 src/struphy/topology/domain_decomposition.py create mode 100644 src/struphy/topology/tests/__init__.py create mode 100644 src/struphy/topology/tests/test_domain_decomposition.py diff --git a/src/struphy/topology/__init__.py b/src/struphy/topology/__init__.py index e69de29bb..33a751a95 100644 --- a/src/struphy/topology/__init__.py +++ b/src/struphy/topology/__init__.py @@ -0,0 +1,13 @@ +from .domain_decomposition import ( + DomainDecompositionOptimization, + DomainDecompositionTiming, + candidate_masks, + optimize_domain_decomposition, +) + +__all__ = [ + "DomainDecompositionOptimization", + "DomainDecompositionTiming", + "candidate_masks", + "optimize_domain_decomposition", +] diff --git a/src/struphy/topology/domain_decomposition.py b/src/struphy/topology/domain_decomposition.py new file mode 100644 index 000000000..eb2f204a9 --- /dev/null +++ b/src/struphy/topology/domain_decomposition.py @@ -0,0 +1,169 @@ +"""Small, empirical domain-decomposition optimization helpers. + +The current Struphy grid API exposes decomposition choices through +``mpi_dims_mask``. This module measures those choices without knowing +anything about the model: the caller supplies a function which performs one +step for a given mask. +""" + +from __future__ import annotations + +import itertools +import time +from dataclasses import dataclass +from typing import Callable, Iterable + +from feectools.ddm.mpi import mpi as MPI +from feectools.ddm.partition import compute_dims + + +Mask = tuple[bool, bool, bool] + + +@dataclass(frozen=True) +class DomainDecompositionTiming: + """Measured wall time for one decomposition mask.""" + + mask: Mask + seconds: float + + +@dataclass(frozen=True) +class DomainDecompositionOptimization: + """Result of an empirical decomposition search.""" + + best_mask: Mask + timings: tuple[DomainDecompositionTiming, ...] + + +def candidate_masks( + num_elements: tuple[int, int, int], + comm_size: int, +) -> tuple[Mask, ...]: + """Return valid non-empty ``mpi_dims_mask`` candidates. + + A candidate is retained when :func:`feectools.ddm.partition.compute_dims` + can construct a process grid for it. The order is stable and starts with + one-dimensional decompositions, followed by two- and three-dimensional + decompositions. + """ + + if len(num_elements) != 3: + raise ValueError("num_elements must contain exactly three dimensions") + if comm_size < 1: + raise ValueError("comm_size must be positive") + if any(int(n) <= 0 for n in num_elements): + raise ValueError("num_elements must contain positive values") + if int(comm_size) > num_elements[0] * num_elements[1] * num_elements[2]: + raise ValueError("comm_size cannot exceed the number of grid elements") + + masks: list[Mask] = [] + for mask in itertools.product((False, True), repeat=3): + if not any(mask): + continue + try: + nprocs, blocksizes = compute_dims(comm_size, list(num_elements), mpi_dims_mask=list(mask)) + except (AssertionError, ValueError): + continue + # ``compute_dims`` accepts a process grid whose local block would have + # zero cells in a short direction. FEEC discretization cannot use + # such a grid, so reject it here before invoking the user callback. + if any(p > n or block < 1 for p, n, block in zip(nprocs, num_elements, blocksizes)): + continue + masks.append(mask) + + return tuple(sorted(masks, key=lambda mask: (sum(mask), mask))) + + +def optimize_domain_decomposition( + num_elements: tuple[int, int, int], + step: Callable[[Mask], object], + *, + comm=None, + masks: Iterable[Mask] | None = None, + warmups: int = 1, + repetitions: int = 3, +) -> DomainDecompositionOptimization: + """Select the fastest decomposition by timing a supplied one-step call. + + Parameters + ---------- + num_elements + Global grid resolution, used to reject invalid masks. + step + Callable that constructs/configures the candidate decomposition and + performs exactly one timestep. It is called collectively by all MPI + ranks for every candidate. + comm + MPI communicator. Defaults to ``MPI.COMM_WORLD``. + masks + Optional explicit subset of masks. By default all valid masks are + measured. + warmups, repetitions + Number of discarded and measured calls per candidate. The reported + time is the average of the communicator-wide maximum time per call. + + Returns + ------- + DomainDecompositionOptimization + All measured times and the mask with the lowest average timestep time. + + Notes + ----- + This function does not reuse simulations between candidates. A practical + ``step`` callback should therefore build a fresh simulation for its mask, + call ``sim.run(one_time_step=True)``, and clean up its output if needed. + """ + + if warmups < 0 or repetitions < 1: + raise ValueError("warmups must be non-negative and repetitions must be positive") + + if comm is None: + comm = MPI.COMM_WORLD + comm_size = comm.Get_size() + valid_masks = set(candidate_masks(num_elements, comm_size)) + selected_masks = tuple(valid_masks if masks is None else _validate_masks(masks, valid_masks)) + if not selected_masks: + raise ValueError("no valid decomposition masks were supplied") + selected_masks = tuple(sorted(selected_masks, key=lambda mask: (sum(mask), mask))) + + timings: list[DomainDecompositionTiming] = [] + for mask in selected_masks: + for _ in range(warmups): + _barrier(comm) + step(mask) + samples = [] + for _ in range(repetitions): + _barrier(comm) + start = time.perf_counter() + step(mask) + elapsed = time.perf_counter() - start + samples.append(_global_max(comm, elapsed)) + _barrier(comm) + timings.append(DomainDecompositionTiming(mask=mask, seconds=sum(samples) / len(samples))) + + best = min(timings, key=lambda timing: timing.seconds) + return DomainDecompositionOptimization(best_mask=best.mask, timings=tuple(timings)) + + +def _validate_masks(masks: Iterable[Mask], valid_masks: set[Mask]) -> set[Mask]: + selected = set() + for mask in masks: + normalized = tuple(mask) + if len(normalized) != 3 or not all(isinstance(value, bool) for value in normalized): + raise ValueError(f"invalid decomposition mask: {mask!r}") + if normalized not in valid_masks: + raise ValueError(f"decomposition mask is not valid for this grid/MPI size: {mask!r}") + selected.add(normalized) + return selected + + +def _barrier(comm) -> None: + if comm.Get_size() > 1: + comm.Barrier() + + +def _global_max(comm, value: float) -> float: + if comm.Get_size() == 1: + return value + return float(comm.allreduce(value, op=MPI.MAX)) diff --git a/src/struphy/topology/tests/__init__.py b/src/struphy/topology/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/topology/tests/test_domain_decomposition.py b/src/struphy/topology/tests/test_domain_decomposition.py new file mode 100644 index 000000000..19d2906f1 --- /dev/null +++ b/src/struphy/topology/tests/test_domain_decomposition.py @@ -0,0 +1,107 @@ +import pytest + +import struphy.topology.domain_decomposition as dd +from struphy.topology.domain_decomposition import candidate_masks, optimize_domain_decomposition + + +def test_candidate_masks_are_valid_and_stable(): + masks = candidate_masks((8, 8, 8), 4) + + assert masks == ( + (False, False, True), + (False, True, False), + (True, False, False), + (False, True, True), + (True, False, True), + (True, True, False), + (True, True, True), + ) + + +def test_optimizer_selects_fastest_candidate(monkeypatch): + calls = [] + costs = { + (False, False, True): 0.003, + (True, True, True): 0.001, + } + clock = [0.0] + + def fake_perf_counter(): + return clock[0] + + def step(mask): + calls.append(mask) + clock[0] += costs[mask] + + monkeypatch.setattr(dd.time, "perf_counter", fake_perf_counter) + + result = optimize_domain_decomposition( + (8, 8, 8), + step, + masks=costs, + warmups=0, + repetitions=1, + ) + + assert result.best_mask == (True, True, True) + assert len(result.timings) == 2 + assert set(calls) == set(costs) + + baseline = next(t.seconds for t in result.timings if t.mask == (False, False, True)) + best = next(t.seconds for t in result.timings if t.mask == (True, True, True)) + assert baseline / best == pytest.approx(3.0) + + +def test_eight_rank_anisotropic_case_can_leave_one_direction_undecomposed(monkeypatch): + """An anisotropic 3D case can prefer a 2D process grid.""" + costs = { + mask: 0.01 + for mask in candidate_masks((48, 24, 8), 8) + } + costs[(True, True, False)] = 0.001 + clock = [0.0] + + monkeypatch.setattr(dd.time, "perf_counter", lambda: clock[0]) + + def step(mask): + clock[0] += costs[mask] + + result = optimize_domain_decomposition( + (48, 24, 8), + step, + comm=dd.MPI.COMM_WORLD, + warmups=0, + repetitions=1, + ) + + assert result.best_mask == (True, True, False) + assert result.best_mask != (True, True, True) + + +def test_optimizer_rejects_invalid_mask(): + with pytest.raises(ValueError, match="not valid"): + optimize_domain_decomposition( + (8, 8, 8), + lambda mask: None, + masks=((False, False, False),), + warmups=0, + repetitions=1, + ) + + +def test_candidate_masks_reject_too_many_ranks(): + with pytest.raises(ValueError, match="cannot exceed"): + candidate_masks((2, 2, 2), 9) + + +def test_candidate_masks_reject_empty_short_direction(): + masks = candidate_masks((16, 64, 1), 8) + + assert masks == ( + (False, True, False), + (True, False, False), + (False, True, True), + (True, False, True), + (True, True, False), + (True, True, True), + ) From 743bec79e5732ba18213a7ad78fc182123f93eee Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 22 Aug 2026 22:16:13 +0200 Subject: [PATCH 2/3] formatting --- src/struphy/topology/domain_decomposition.py | 1 - src/struphy/topology/tests/test_domain_decomposition.py | 5 +---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/struphy/topology/domain_decomposition.py b/src/struphy/topology/domain_decomposition.py index eb2f204a9..2ae593fc9 100644 --- a/src/struphy/topology/domain_decomposition.py +++ b/src/struphy/topology/domain_decomposition.py @@ -16,7 +16,6 @@ from feectools.ddm.mpi import mpi as MPI from feectools.ddm.partition import compute_dims - Mask = tuple[bool, bool, bool] diff --git a/src/struphy/topology/tests/test_domain_decomposition.py b/src/struphy/topology/tests/test_domain_decomposition.py index 19d2906f1..0e12e545b 100644 --- a/src/struphy/topology/tests/test_domain_decomposition.py +++ b/src/struphy/topology/tests/test_domain_decomposition.py @@ -54,10 +54,7 @@ def step(mask): def test_eight_rank_anisotropic_case_can_leave_one_direction_undecomposed(monkeypatch): """An anisotropic 3D case can prefer a 2D process grid.""" - costs = { - mask: 0.01 - for mask in candidate_masks((48, 24, 8), 8) - } + costs = {mask: 0.01 for mask in candidate_masks((48, 24, 8), 8)} costs[(True, True, False)] = 0.001 clock = [0.0] From 458011ae124d178ab5007f9defb75cd76f931f69 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 22 Aug 2026 22:16:33 +0200 Subject: [PATCH 3/3] Added benchmark --- .../benchmark_domain_decomposition.py | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 profiling/examples/ToyGyrokinetic/diocotron_instability/benchmark_domain_decomposition.py diff --git a/profiling/examples/ToyGyrokinetic/diocotron_instability/benchmark_domain_decomposition.py b/profiling/examples/ToyGyrokinetic/diocotron_instability/benchmark_domain_decomposition.py new file mode 100644 index 000000000..62119a323 --- /dev/null +++ b/profiling/examples/ToyGyrokinetic/diocotron_instability/benchmark_domain_decomposition.py @@ -0,0 +1,118 @@ +"""Benchmark domain-decomposition masks for a cylindrical ToyDrift PIC case. + +This is an intentionally anisotropic case: the logical grid has many radial +cells and fewer azimuthal cells, while particles are loaded and sorted in the +cylindrical domain. On the reference 8-rank run, the mask ``(True, False, +False)`` was about 32% faster than the default ``(True, True, True)`` mask. + +Run from the repository root, for example:: + + mpirun -n 8 python profiling/examples/ToyGyrokinetic/diocotron_instability/benchmark_domain_decomposition.py + +The benchmark times allocation plus one call to ``model.integrate``. This +keeps the example independent of the optional simulation profiling wrapper; +the decomposition comparison itself is still performed by the public +optimizer. +""" + +from __future__ import annotations + +import argparse +import logging + +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, +) +from struphy.models import ToyDrift +from struphy.topology import optimize_domain_decomposition + + +NUM_ELEMENTS = (128, 16, 1) +DEFAULT_MASK = (True, True, True) + + +def run_one_step(mask: tuple[bool, bool, bool], output_dir: str) -> None: + """Build one candidate and perform one allocation/integration step.""" + model = ToyDrift(epsilon=1.0, alpha=1.0, base_units=BaseUnits(kBT=1.0)) + domain = domains.HollowCylinder(a1=1.0, a2=10.0, Lz=10.0) + equil = equils.HomogenSlab() + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=50, loading="sobol_standard", spatial="disc"), + weights_params=WeightsParameters(control_variate=True, reject_weights=True, threshold=0.0001), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters( + boxes_per_dim=NUM_ELEMENTS, + do_sort=True, + sorting_frequency=1, + ), + bufsize=2.0, + ) + model.kinetic_ions.var.add_background( + maxwellians.GyroMaxwellian2D(n=(0.0, None), B0=2.0), + ) + + sim = Simulation( + model=model, + env=EnvironmentOptions( + out_folders=output_dir, + sim_folder="mask_" + "".join("1" if value else "0" for value in mask), + restart=False, + ), + time_opts=Time(dt=0.01, Tend=0.01, split_algo="LieTrotter"), + domain=domain, + equil=equil, + grid=grids.TensorProductGrid(num_elements=NUM_ELEMENTS, mpi_dims_mask=mask), + derham_opts=DerhamOptions( + degree=(2, 2, 1), + bcs=(("dirichlet", "dirichlet"), None, None), + ), + ) + sim.allocate() + sim.model.integrate(sim.time_opts.dt, sim.time_opts.split_algo) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--repetitions", type=int, default=2) + parser.add_argument("--output", default="benchmark_domain_decomposition") + args = parser.parse_args() + + logging.getLogger("struphy").setLevel(logging.ERROR) + result = optimize_domain_decomposition( + NUM_ELEMENTS, + lambda mask: run_one_step(mask, args.output), + comm=MPI.COMM_WORLD, + warmups=args.warmups, + repetitions=args.repetitions, + ) + + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"best mask: {result.best_mask}") + for timing in result.timings: + print(f"{timing.mask}: {timing.seconds:.6f} s") + + default = next( + timing.seconds for timing in result.timings if timing.mask == DEFAULT_MASK + ) + speedup = default / min(timing.seconds for timing in result.timings) + print(f"speedup over default: {speedup:.2f}x ({(speedup - 1.0) * 100:.1f}%)") + + +if __name__ == "__main__": + main()