From a7572bfc0b58c32dc24f26123930344205f1ec80 Mon Sep 17 00:00:00 2001 From: Lars Pastewka Date: Fri, 25 Sep 2026 19:47:54 +0200 Subject: [PATCH 1/4] FIX: Put cupy on the rank's GPU _resolve_device binds each MPI rank to GPU local_rank % nb_gpus, but cupy was never told: it places arrays and launches kernels on its *current* device, device 0 unless switched. On two GPUs every cupy temporary of rank 1 (to_device(), the reductions in _reference_lame, ...) therefore landed on GPU 0 while its fields lived on GPU 1 -- slow peer traffic at best, and an illegal-address fault for code that mixes the two in one kernel, as muGrid's hybrid preconditioner did. Homogenization now makes the resolved GPU cupy's current device right after resolving it. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 3 +++ muTopOpt/homogenization.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 630e1fa..d24c3a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ Change log for muTopOpt unreleased ---------- +- FIX: `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 - FIX: 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 diff --git a/muTopOpt/homogenization.py b/muTopOpt/homogenization.py index 543f10f..fccda7d 100644 --- a/muTopOpt/homogenization.py +++ b/muTopOpt/homogenization.py @@ -228,6 +228,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 From 1ef57e1a63790e346aefb1384a4eb1ea95c7655c Mon Sep 17 00:00:00 2001 From: Lars Pastewka Date: Fri, 25 Sep 2026 19:48:16 +0200 Subject: [PATCH 2/4] ENH: Offer muGrid's hybrid Fourier/tridiagonal preconditioner --preconditioner hybrid applies the same reference stiffness as 'green', but inverts it with muGrid's HybridFourierTridiagonalPreconditioner: FFT in the rank-local axes and a block-tridiagonal (PCR) solve in the distributed one, so no all-to-all. hybrid-jacobi wraps it in the same J-FFT Jacobi scaling as green-jacobi, with the same refresh() on a material update. The FFT engine's real-space split is already the slab ([1, ..., P]) the hybrid needs, so it runs on the engine directly. It is exact: on a fixed design it matches 'green' to ~1e-14 per apply and takes identical CG iterations. In a full trust-region run (fp64, fixed rtol 1e-8, 128^3) the optimizer takes the same path to 7 digits, with ~1% more CG iterations from round-off. Time per CG iteration, solve_bench.py, fp32, Q1, 2x GTX TITAN X: n^3 green-jacobi 1 GPU / 2 GPU hybrid-jacobi 1 GPU / 2 GPU 128 11.4 / 8.78 ms 13.6 / 8.50 ms 160 23.3 / 17.7 ms 32.2 / 15.8 ms 192 40.4 / 31.1 ms 62.6 / 30.2 ms So it is slower on one GPU, where it only adds work, and scales ~2x from one GPU to two against 1.3x for the distributed FFT. solve_bench.py gains the new choices and now runs under mpirun: it used to build a serial muGrid.Communicator on every rank, so each rank solved the whole problem alone. Rank 0 reports; the CSV gains nb_ranks, so it refuses to append to a file written with the old columns. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 8 ++++++++ benchmarks/solve_bench.py | 29 +++++++++++++++++++++++------ muTopOpt/homogenization.py | 32 +++++++++++++++++++++++++++++++- simulate.py | 11 ++++++++--- 4 files changed, 70 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d24c3a0..4dc27a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,14 @@ Change log for muTopOpt unreleased ---------- +- 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 - FIX: `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 diff --git a/benchmarks/solve_bench.py b/benchmarks/solve_bench.py index eec311d..3d09ced 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, ) @@ -176,7 +188,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 +228,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 +250,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 +270,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 fccda7d..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, ) @@ -413,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/simulate.py b/simulate.py index 135f96f..cf21864 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", From af36cc2fd284e1787a9d9cd4da535657b7556585 Mon Sep 17 00:00:00 2001 From: Lars Pastewka Date: Fri, 25 Sep 2026 19:48:44 +0200 Subject: [PATCH 3/4] FIX: Start every rank count from the same initial design simulate.py and solve_bench.py called initial_density(homog.nb_pixels, ...) -- the rank-local shape -- with the same seed on every rank. Each rank thus drew the *same* noise and filtered it periodically on its own subdomain, so a 2-rank run started from two copies of a half-size design instead of the serial one, and the problem being solved changed with the number of ranks (checksum of the gathered 64^3 start: 0188fc8aca1b on one rank, 8acb9d447da1 on two). It showed up as iteration counts that moved with the rank count: 161 vs 143 CG iterations per solve at 128^3 in solve_bench.py. initial_density now takes the global shape and, optionally, the rank's subdomain_locations / nb_subdomain_grid_pts, and returns that rank's slice of the field generated on the whole grid. simulate_conduction.py already did this by hand; it now uses the same arguments. Serial callers are unchanged. With the fix the start is 0188fc8aca1b on one and two ranks, and the iterations per solve agree (160.8 vs 161.0). The new test checks that 2-, 3- and 4-slab subdomains reassemble bit-for-bit into the global field. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 7 +++++++ benchmarks/solve_bench.py | 4 +++- muTopOpt/optimize.py | 26 +++++++++++++++++++++++++- simulate.py | 4 +++- simulate_conduction.py | 13 +++++-------- test/test_initial_density.py | 33 +++++++++++++++++++++++++++++++++ 6 files changed, 76 insertions(+), 11 deletions(-) create mode 100644 test/test_initial_density.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4dc27a2..869c77b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ Change log for muTopOpt unreleased ---------- +- FIX: `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 diff --git a/benchmarks/solve_bench.py b/benchmarks/solve_bench.py index 3d09ced..066a600 100644 --- a/benchmarks/solve_bench.py +++ b/benchmarks/solve_bench.py @@ -114,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 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 cf21864..290b8a4 100644 --- a/simulate.py +++ b/simulate.py @@ -544,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)) From 4b4f23e6436f1dd6f4d461cef33a9275b737a188 Mon Sep 17 00:00:00 2001 From: Lars Pastewka Date: Fri, 25 Sep 2026 19:57:06 +0200 Subject: [PATCH 4/4] DOC: Updated CHANGELOG.md --- CHANGELOG.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 869c77b..9859ad3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Change log for muTopOpt unreleased ---------- -- FIX: `simulate.py` and `benchmarks/solve_bench.py` start every rank count +- 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 @@ -19,10 +19,10 @@ unreleased `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 -- FIX: `Homogenization` makes the rank's GPU cupy's current device. cupy +- 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 -- FIX: The inner CG's norms come from muGrid's `linalg.norm_sq` / `vecdot` / +- 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 @@ -34,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