diff --git a/CHANGELOG.md b/CHANGELOG.md index 630e1fa..9859ad3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,25 @@ Change log for muTopOpt unreleased ---------- -- FIX: The inner CG's norms come from muGrid's `linalg.norm_sq` / `vecdot` / +- BUG: `simulate.py` and `benchmarks/solve_bench.py` start every rank count + from the same design. `initial_density` was handed the rank-local shape, so + each rank drew the same noise and filtered it periodically on its own + subdomain: a 2-rank run started from two copies of a half-size design, not + the serial one. `initial_density` now takes the global shape plus + `subdomain_locations` / `nb_subdomain_grid_pts` and returns this rank's + slice (which `simulate_conduction.py` already did by hand) +- ENH: `--preconditioner hybrid` and `hybrid-jacobi` invert the reference + stiffness with muGrid's `HybridFourierTridiagonalPreconditioner` (FFT in the + rank-local axes, a tridiagonal solve in the distributed one) instead of a + distributed FFT, alone or inside the same J-FFT Jacobi scaling. Same + operator, same CG iterations; on two GPUs 3-11% faster per CG iteration than + `green-jacobi` from 128^3 in float32, slower on one GPU +- ENH: `benchmarks/solve_bench.py` runs under `mpirun` (it used to build a + serial communicator on every rank) and records `nb_ranks` in its CSV +- BUG: `Homogenization` makes the rank's GPU cupy's current device. cupy + otherwise defaults to device 0, so with one GPU per rank every cupy + temporary of rank 1 landed on GPU 0 next to fields on GPU 1 +- BUG: The inner CG's norms come from muGrid's `linalg.norm_sq` / `vecdot` / `axpy_norm_sq` instead of `xp.dot` on the raw field buffer. BLAS `sdot` accumulates a float32 field in float32, so its error grows *linearly* in the number of entries -- measured 4.2e-9 at 32^3 but 5.5e-6 at 128^3 and ~4.6e-4 @@ -16,11 +34,11 @@ unreleased additionally makes `b_norm` bit-identical to the `||b||` the CG itself converges against, instead of dividing a double-accumulated residual by a float32-accumulated norm -- FIX: The consistent-objective correction `-λᵀr` likewise uses +- BUG: The consistent-objective correction `-λᵀr` likewise uses `linalg.vecdot` rather than summing a full-size float32 product array. This value is the reported objective *and* feeds the trust region's accuracy control, so it carries the tightest error budget in the package -- FIX: Inner-CG tolerances are clamped to what the field precision can reach +- BUG: Inner-CG tolerances are clamped to what the field precision can reach (1e-6 in float32, where eps is 1.19e-7 and the true residual `b - Kx` stagnates while only the recursive residual keeps shrinking). The floor is applied at every point a tolerance is *consumed* -- `solve_rhs`, the adaptive diff --git a/benchmarks/solve_bench.py b/benchmarks/solve_bench.py index eec311d..066a600 100644 --- a/benchmarks/solve_bench.py +++ b/benchmarks/solve_bench.py @@ -39,6 +39,7 @@ import os import platform import re +import sys import time import muGrid @@ -50,7 +51,7 @@ #: Fields written by ``--csv``, in order. CSV_COLUMNS = ( "timestamp", "host", "gpu", "mugrid_version", "mugrid_commit", - "mugrid_dirty", "device", "precision", "dim", "n", "nb_pixels", + "mugrid_dirty", "device", "nb_ranks", "precision", "dim", "n", "nb_pixels", "preconditioner", "element", "rtol", "solves", "iters_per_solve", "ms_per_solve", "ms_per_cg_iter", ) @@ -86,6 +87,17 @@ def provenance(homog): } +def _communicator(): + """MPI communicator under a parallel launch, the serial one otherwise.""" + try: + from mpi4py import MPI + except ImportError: + return muGrid.Communicator() + if MPI.COMM_WORLD.size == 1: + return muGrid.Communicator() + return muGrid.Communicator(MPI.COMM_WORLD) + + def build(args, timer=None): """Construct the homogenization problem and apply a fixed design.""" dtype = {"single": np.float32, "double": np.float64}[args.precision] @@ -94,7 +106,7 @@ def build(args, timer=None): penalty=args.penalty, void_ratio=args.void_ratio, ) homog = Homogenization( - (args.n,) * args.dim, material, + (args.n,) * args.dim, material, comm=_communicator(), element=args.element, preconditioner=args.preconditioner, device=args.device, dtype=dtype, timer=timer, ) @@ -102,9 +114,11 @@ def build(args, timer=None): # length 3*eta, eta = one grid spacing), so the material contrast the # preconditioner sees is representative rather than uniform. rho = initial_density( - homog.nb_pixels, kind="filtered_random", seed=args.seed, + homog.nb_grid_pts, kind="filtered_random", seed=args.seed, length=3.0 * float(homog.grid_spacing[0]), grid_spacing=homog.grid_spacing, + subdomain_locations=homog.engine.subdomain_locations, + nb_subdomain_grid_pts=homog.nb_pixels, ) homog.set_density(rho) return homog @@ -176,7 +190,7 @@ def main(): choices=("single", "double"), help="solver precision (default: single)") p.add_argument("--preconditioner", default="green-jacobi", - choices=("green-jacobi", "green")) + choices=("green-jacobi", "green", "hybrid-jacobi", "hybrid")) p.add_argument("--element", default="p1", choices=("p1", "q1")) p.add_argument("--rtol", type=float, default=1e-2, help="inner CG relative tolerance. The default matches the " @@ -216,8 +230,12 @@ def main(): timer = Timer() homog = build(args, timer=timer) strains = unit_strains(args.dim) + root = homog.comm.rank == 0 + if not root: + # Every rank takes part in the solves; only rank 0 reports. + sys.stdout = open(os.devnull, "w") - print(f"n={args.n}^{args.dim} device={args.device} " + print(f"n={args.n}^{args.dim} ranks={homog.comm.size} device={args.device} " f"precision={args.precision} preconditioner={args.preconditioner} " f"element={args.element} rtol={args.rtol:g}", flush=True) @@ -234,6 +252,7 @@ def main(): "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), **provenance(homog), "device": args.device, + "nb_ranks": homog.comm.size, "precision": args.precision, "dim": args.dim, "n": args.n, @@ -253,10 +272,10 @@ def main(): f"min {min(vals):.3f} median {np.median(vals):.3f} " f"max {max(vals):.3f}") - if timer is not None: + if timer is not None and root: timer.print_summary(title="wall time per CG phase") - if args.csv: + if args.csv and root: new = not os.path.exists(args.csv) if not new: with open(args.csv, newline="") as fh: diff --git a/muTopOpt/homogenization.py b/muTopOpt/homogenization.py index 543f10f..9ee5ea5 100644 --- a/muTopOpt/homogenization.py +++ b/muTopOpt/homogenization.py @@ -32,6 +32,8 @@ import muGrid from muGrid import linalg from muGrid.Preconditioners import ( + GreenJacobiPreconditioner, + HybridFourierTridiagonalPreconditioner, make_green_jacobi_preconditioner, make_reference_stiffness_preconditioner, ) @@ -228,6 +230,13 @@ def __init__( self.domain_volume = float(np.prod(self.domain_lengths)) self.device = _resolve_device(device, self.comm) + # cupy places arrays (and launches kernels) on its *current* device, + # device 0 unless told otherwise. With one GPU per rank that would put + # every cupy temporary of rank 1 on GPU 0, next to fields on GPU 1. + if self.device is not None and self.device.is_device: + import cupy + + cupy.cuda.Device(self.device.device_id).use() # On a unified-memory APU, default to the managed allocator so device # fields can use the full HBM rather than the smaller coarse-grained # window (see _enable_managed_device_allocator). On a *discrete* GPU @@ -406,11 +415,39 @@ def apply_ref(u, f): self.engine, apply_ref, self.dim, dtype=self.dtype, timer=self.timer, ) + elif self.preconditioner_kind in ("hybrid", "hybrid-jacobi"): + # The same reference operator as 'green', inverted exactly by + # FFT in the rank-local axes and a tridiagonal solve in the + # distributed one -- no all-to-all. The engine's real-space + # split is the slab ([1, ..., P]) it needs. + green = HybridFourierTridiagonalPreconditioner( + self.engine, self.grid_spacing, lam_ref, mu_ref, + communicator=self.comm, element=self.element, + timer=self.timer, dtype=self.dtype, + ) + if self.preconditioner_kind == "hybrid": + self._prec = green + else: + diagonal = self.engine.real_space_field( + "to_hybrid_jacobi_diagonal", components=(self.dim,), + dtype=self.dtype) + self.op.assemble_diagonal(self.lam, self.mu, diagonal) + prec = GreenJacobiPreconditioner( + green, diagonal, timer=self.timer, + name="hybrid-jacobi", communicator=self.comm, + ) + + def refresh(): + self.op.assemble_diagonal(self.lam, self.mu, diagonal) + prec.update_diagonal(diagonal) + + prec.refresh = refresh + self._prec = prec else: raise ValueError( f"unknown preconditioner '{self.preconditioner_kind}'" ) - elif self.preconditioner_kind == "green-jacobi": + elif self.preconditioner_kind in ("green-jacobi", "hybrid-jacobi"): # Reuse the (reference) Green part; refresh only the Jacobi diagonal # from the updated material. self._prec.refresh() diff --git a/muTopOpt/optimize.py b/muTopOpt/optimize.py index 56929d9..bc02d3c 100644 --- a/muTopOpt/optimize.py +++ b/muTopOpt/optimize.py @@ -205,9 +205,19 @@ def _make_inner_tolerance(problem, cg_tol_start, cg_tol_min, cg_forcing_c, def initial_density(shape, kind="uniform", volume_fraction=0.5, seed=0, - smoothing=2, length=None, grid_spacing=None, contrast=0.5): + smoothing=2, length=None, grid_spacing=None, contrast=0.5, + subdomain_locations=None, nb_subdomain_grid_pts=None): """Build an initial element-wise density. + ``shape`` is the *global* grid. Under MPI pass the rank's + ``subdomain_locations`` and ``nb_subdomain_grid_pts`` (e.g. + ``homog.engine.subdomain_locations`` and ``homog.nb_pixels``): the field is + then generated on the whole grid, identically on every rank, and only this + rank's slice is returned -- so a parallel run starts from exactly the design + a serial run does. Passing the *local* shape as ``shape`` instead would draw + and filter an independent field per subdomain, making the starting design + depend on the number of ranks. + ``kind='uniform'`` fills with ``volume_fraction``. ``kind='random'`` draws a box-smoothed random field (``smoothing`` sweeps), @@ -229,6 +239,20 @@ def initial_density(shape, kind="uniform", volume_fraction=0.5, seed=0, eta`` -- the regularization then *sharpens the blob boundaries* rather than dissolving the blobs. """ + if (subdomain_locations is None) != (nb_subdomain_grid_pts is None): + raise ValueError("pass both subdomain_locations and " + "nb_subdomain_grid_pts, or neither") + rho = _global_initial_density(shape, kind, volume_fraction, seed, + smoothing, length, grid_spacing, contrast) + if subdomain_locations is None: + return rho + return rho[tuple(slice(int(lo), int(lo) + int(n)) for lo, n in + zip(subdomain_locations, nb_subdomain_grid_pts))].copy() + + +def _global_initial_density(shape, kind, volume_fraction, seed, smoothing, + length, grid_spacing, contrast): + """The initial density on the full grid ``shape``; see initial_density.""" if kind == "uniform": return np.full(shape, float(volume_fraction)) if kind == "random": diff --git a/simulate.py b/simulate.py index 135f96f..290b8a4 100644 --- a/simulate.py +++ b/simulate.py @@ -297,11 +297,16 @@ def main(): ) p.add_argument( "--preconditioner", - choices=["green-jacobi", "green"], + choices=["green-jacobi", "green", "hybrid-jacobi", "hybrid"], default="green-jacobi", help="inner-solve preconditioner: 'green-jacobi' (J-FFT, " - "reference stiffness times a per-pixel Jacobi scaling) " - "or 'green' (plain reference-stiffness Green operator)", + "reference stiffness times a per-pixel Jacobi scaling), " + "'green' (plain reference-stiffness Green operator), or " + "'hybrid-jacobi' / 'hybrid', the same two with the reference " + "stiffness inverted by muGrid's HybridFourierTridiagonal" + "Preconditioner: FFT in the rank-local axes and a tridiagonal " + "solve in the distributed one, so no all-to-all. Same iterations; " + "slower on one GPU, faster from ~128^3 on two", ) p.add_argument( "--element", @@ -539,12 +544,14 @@ def E_nu_from_K_G(K, G): if args.init == "filtered_random" and length is None: length = 3.0 * reg.eta rho0 = initial_density( - homog.nb_pixels, + tuple(args.nb_grid_pts), kind=args.init, volume_fraction=args.init_volume_fraction, seed=args.seed, length=length, grid_spacing=homog.grid_spacing, + subdomain_locations=homog.engine.subdomain_locations, + nb_subdomain_grid_pts=homog.nb_pixels, ) else: # Restart from the last frame of a previous run (every rank reads the diff --git a/simulate_conduction.py b/simulate_conduction.py index 95ba7d3..f793e2f 100755 --- a/simulate_conduction.py +++ b/simulate_conduction.py @@ -411,21 +411,18 @@ def effective_conductivity(fluxes): length = args.init_length if args.init == "filtered_random" and length is None: length = 3.0 * reg.eta - # Generate the initial density on the full global grid so that MPI-parallel - # runs start from exactly the same field as a serial run (the local - # subdomain for each rank is just a slice of that global field). - rho0_global = initial_density( + # Generated on the full global grid, so MPI-parallel runs start from + # exactly the same field as a serial run; each rank keeps its slice. + rho0 = initial_density( tuple(args.nb_grid_pts), kind=args.init, volume_fraction=args.init_volume_fraction, seed=args.seed, length=length, grid_spacing=homog.grid_spacing, + subdomain_locations=homog.engine.subdomain_locations, + nb_subdomain_grid_pts=homog.nb_pixels, ) - rho0 = rho0_global[ - tuple(slice(lo, lo + n) for lo, n in - zip(homog.engine.subdomain_locations, homog.nb_pixels)) - ].copy() else: # Restart from a previous run rho0_global, restart_meta = restart_density( diff --git a/test/test_initial_density.py b/test/test_initial_density.py new file mode 100644 index 0000000..d864c81 --- /dev/null +++ b/test/test_initial_density.py @@ -0,0 +1,33 @@ +# +# Copyright 2026 Lars Pastewka +# +# MIT License (see LICENSE) +# +"""The initial design must not depend on the domain decomposition.""" + +import numpy as np +import pytest + +from muTopOpt.optimize import initial_density + + +@pytest.mark.parametrize("kind", ["uniform", "random", "filtered_random"]) +@pytest.mark.parametrize("nb_ranks", [2, 3, 4]) +def test_subdomains_tile_the_global_field(kind, nb_ranks): + """Slab subdomains, as the FFT engine splits the grid, reassemble to + exactly the field generated on the whole grid.""" + shape = (16, 12, 24) + kwargs = dict(kind=kind, seed=5, length=0.1, volume_fraction=0.4) + glob = initial_density(shape, **kwargs) + nz = shape[-1] // nb_ranks + pieces = [ + initial_density(shape, subdomain_locations=(0, 0, r * nz), + nb_subdomain_grid_pts=shape[:-1] + (nz,), **kwargs) + for r in range(nb_ranks) + ] + np.testing.assert_array_equal(np.concatenate(pieces, axis=-1), glob) + + +def test_subdomain_needs_both_arguments(): + with pytest.raises(ValueError): + initial_density((8, 8), subdomain_locations=(0, 0))