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

[pastewka]

Stacked on #16. The base branch is fix/float32-reductions-and-tolerance, so this diff shows only the three commits below. Retarget to main once #16 is merged.

Three commits, from testing muGrid's HybridFourierTridiagonalPreconditioner on two GPUs (2× GTX TITAN X, one MPI rank per GPU). The companion muGrid fix is muSpectre/muGrid#214.

1. FIX: Put cupy on the rank's GPU

_resolve_device gives each rank the GPU local_rank % nb_gpus, but cupy was never switched to it. cupy works on its current device (0 by default), so rank 1's cupy temporaries landed on GPU 0 next to fields on GPU 1. Homogenization now makes the resolved GPU cupy's current device.

2. ENH: --preconditioner hybrid / hybrid-jacobi

These apply the same reference stiffness as green / green-jacobi, but invert it with FFT in the rank-local axes and a tridiagonal solve along the distributed axis, so there is no all-to-all. The FFT engine already splits the grid into the slabs the hybrid needs, so no new decomposition is required. benchmarks/solve_bench.py now also runs under mpirun (before, every rank built a serial communicator and solved the whole problem alone) and records nb_ranks in its CSV.

3. FIX: Rank-independent initial design

initial_density(homog.nb_pixels, seed=…) generated the random field on each rank's local subdomain with the same seed. A 2-rank run therefore started from two copies of a half-size design, so the problem changed with the number of ranks. initial_density now takes the global shape plus optional subdomain_locations / nb_subdomain_grid_pts and returns this rank's slice. simulate.py, simulate_conduction.py (which already did this by hand) and solve_bench.py use it. A new test checks that 2/3/4-slab subdomains reassemble exactly into the global field.

Correctness

• The hybrid matches green to ~1e-14 per apply and needs identical CG iterations on a fixed design.
• Full trust-region runs at 128³ (fp64, fixed rtol 1e-8, 3 outer iterations, identical start) take the same path with all four setups (FFT/hybrid × 1/2 GPUs): the objective agrees to 7 digits at every step (6.086835 → 5.356614 → 4.718631), with 156 solves each. The hybrid needs about 1.4 % more CG iterations from round-off.
• With the loose adaptive tolerance in fp32, the default setting, runs do diverge between preconditioners. Round-off changes when CG stops (±1 iteration), which moves the gradient by about 0.3 % and flips trust-region accept/reject decisions. For timing, compare ms per CG iteration, not total wall time.
• 182 tests pass.

Performance (time per CG iteration)

Measured with solve_bench.py, Q1, rtol 1e-2, identical design on 1 and 2 GPUs:

	FFT 1 GPU	hybrid 1 GPU	FFT 2 GPU	hybrid 2 GPU
fp32 128³	11.4 ms	13.6 ms	8.78 ms	8.50 ms
fp32 160³	23.3 ms	32.2 ms	17.7 ms	15.8 ms
fp32 192³	40.4 ms	62.6 ms	31.1 ms	30.2 ms
fp64 128³	56.9 ms	65.8 ms	36.3 ms	39.0 ms
fp64 192³	195 ms	232 ms	129.8 ms	130.3 ms

End to end in fp64 on 2 GPUs the two tie: 39.2 vs 39.0 ms/iteration at 128³ and 130.0 vs 130.8 at 192³. The hybrid is always slower on one GPU, where it only adds work. Its advantage is scaling: about 2× from 1 to 2 GPUs versus 1.3× for the distributed FFT, whose all-to-all goes through the host here because the two GPUs have no P2P link. The default therefore stays green-jacobi.

🤖 Generated with Claude Code