Add auto parallelization

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

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

src/struphy/topology/domain_decomposition.py

@@ -0,0 +1,168 @@
+"""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))