Hybrid Fourier/tridiagonal preconditioner, multi-GPU fixes, rank-independent initial design

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

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,17 +106,19 @@ 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,
     )
     # The same smooth random design simulate.py starts from (correlation
     # 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:

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

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":

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

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(

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