Make the adjoint design-region gradient independent of the MPI chunk division by jin-castle · Pull Request #3264 · NanoComp/meep

jin-castle · 2026-08-09T02:26:18Z

Fixes #2578 (adjoint gradients inconsistent with finite differences in 3D).

Symptom. The adjoint gradient of a MaterialGrid design region depends on
the number of MPI processes while the objective value does not. Measured on a
20×20×6 design grid (3D waveguide, both with and without subpixel smoothing):

f0 identical to 8+ digits at every process count;
gradient max-norm difference between -np 1 and -np 8: 57% of
|g|_max (smoothing off) / 18% (smoothing on), spread over ~80% of the
design nodes in chunk-sized patches;
directional finite-difference checks that pass at -np 1 (ratio 1.0002)
fail under MPI (ratios 0.86–1.22 at -np 16).

Root causes — all in material_grids_addgradient(), which computed the
result from per-chunk DFT views:

chunks of different components were paired by list index
(forward_dft_chunks[ci_forward][cur_chunk]); the per-component chunk sets
can legitimately differ near the monitor edges, after which every
subsequent pair associates the wrong regions, and cross terms were dropped
entirely when the counts differed (cur_chunk >= num_f_chunks → continue);
neighboring-point lookups bounds-checked only the chunk-local linear
index, so an out-of-range point in one dimension could alias to a valid
index of a different point instead of returning zero;
the ±1-pixel padding of persist dft chunks is clamped to the owning
fields chunk and cannot cross process boundaries, so the data those
lookups need does not exist at rank boundaries;
the same-component term indexed the forward chunk with the adjoint loop
index, relying on the chunk-alignment assumption of (1).

Fix. Gather each monitor component onto its full (padded) design-region
grid first — monitor-owned points authoritatively (the is_old..ie_old
decomposition is disjoint), the padding ring as a multiplicity-average of the
ghost copies so its coverage does not depend on the chunk division — then run
the same per-point math on this chunk-independent view, with the work strided
across processes and per-dimension bounds checks on neighbor lookups. Memory
cost is the gathered design-region DFTs (6 complex arrays over the monitor
grid per frequency), negligible for realistic design regions.

Validation.

check	before	after
`-np 8` vs `-np 1`, 20×20×6, smoothing off	57%	7×10⁻¹²
`-np 8/16/32/128` vs `-np 1` (both smoothing cases)	up to 57%	≤7×10⁻¹²
FD-vs-adjoint at `-np 16` and `-np 128` (six grid configurations incl. incommensurate/offset/coarse/fine grids)	ratios 0.86–1.22	identical to serial (1.000 ± FD noise)
serial behavior	—	preserved (`test_subpixel_3d`, adjoint E2E unchanged)

🤖 Generated with Claude Code

material_grids_addgradient computed its result from per-chunk DFT views: chunks of different components were paired by list index, cross terms were dropped when per-component chunk counts differed, neighboring-point lookups used chunk-local linear indices that could alias across dimensions or silently return zero at chunk boundaries, and the +-1-pixel padding of persist dft chunks cannot cross rank boundaries. Together these made the adjoint gradient depend on the number of MPI processes (measured: up to 57% of |g|max on a 20x20x6 3D design region between -np 1 and -np 8, with the objective value identical; see NanoComp#2578). The computation now first gathers each monitor component onto its full (padded) design-region grid -- owned points authoritatively, the padding ring as a multiplicity-average of the ghost copies -- and then evaluates the same per-point math on this chunk-independent view, with the work strided across processes and per-dimension bounds checks on neighbor lookups. Results are now identical (up to summation roundoff) for any number of processes, and serial semantics are preserved.

The rewrite in the previous commit has no coverage: nothing in the suite fails if material_grids_addgradient goes back to reading the per-chunk DFT views. Add a test that computes the design-region gradient twice for the same problem under two different chunk divisions and requires the two vectors to agree. Two things make this cheap to run. The defects are triggered by a chunk boundary crossing the design region, not by MPI as such, so forcing the split with Simulation(num_chunks=...) reproduces them in a serial run -- the test therefore belongs in TESTS and guards the serial jobs too. And comparing the gradient vectors is far more sensitive than the usual directional finite-difference check, which probes along g/|g| and so largely cancels the error: measured against the pre-fix code, the gradient was off by 44% in L2 at four chunks while fd/adjoint still read 1.0008. Measured with this test: 4.7e-1 relative difference before the rewrite, 3e-16 after (serial), 45 s runtime.

stevengj · 2026-08-14T20:06:56Z

Failing single precision tests?

stevengj · 2026-08-14T20:07:45Z

@lxvm, you are running a lot of Meep MPI adjoint jobs, can you look at this?

The MPI single-precision CI job failed on the objective check, at 1.0000001199705797 vs a places=9 bound -- about one float ulp of relative error. The gradient check would have failed too, at 1.5e-7 against 1e-9. Neither bound is reachable in single precision, and the deviation is not something this PR introduces: in a *serial* single-precision run the two chunk divisions give a bitwise identical gradient (rel = 0.0), and the objective -- which no gradient code touches -- moves by the same 1.2e-7 as the gradient once the split is spread across ranks. What varies is summation order in the collective reductions, which the gradient inherits. Measured deviations: single 0.0 (serial) / 1.5e-7 (np 2); double 3.1e-16 (serial) to 2.2e-12 (np 8). The 1e-5 single-precision bound keeps a ~60x margin over the former while staying four orders of magnitude below the 4.7e-1 this test read before the rewrite. Also drops a claim from the docstring about how "the usual directional finite-difference check" works: upstream's checks perturb along a random direction, not g/|g|. The 1.0008 reading quoted there came from a g/|g| probe, so it is stated that way now.

lxvm · 2026-08-18T21:47:44Z

I tested this PR on my Meep adjoint jobs and can confirm that it dramatically improves the accuracy of the adjoint gradient. I also noticed that changing the number of MPI ranks changed the value of the objective function I calculated and therefore I cannot comment on whether the PR fixes the design-region gradient's dependence on the chunk division because there could be another bug present.

stevengj · 2026-08-20T17:19:44Z

@smartalecH, can you take a look at the PR?

stevengj · 2026-08-20T17:22:05Z

+namespace {
+/* DFT data for one field component of the design-region monitors, gathered
+   from every chunk on every process onto the full (padded) monitor grid.
+   The gradient computation below then runs on this chunk-independent view,


I'm worried that doing it this way hurts the scaling… it would be nice to keep the DFT data parallelized…

smartalecH · 2026-08-21T04:01:42Z

Thanks @jin-castle for pulling this together. You managed to find a few different bugs here. But I'm not confident this is the right approach to resolving them. Here are some thoughts, @stevengj :

We definitely have some indexing bugs, and some loose assumptions (e.g. whether component chunks track each other in their respective linked lists). We should absolutely refactor/clean up the code here (although some of the proposed refactoring is more extensive than it needs to be in a single PR).
I don't think we need to materialize all the DFT component data to the same process. We added behavior with the persist flag that essentially stores an extra array element for each component (and extra "halo") which can be used to access these complicated boundaries. I think there are some bugs with how we access those (particularly in 3D) but we should be able to resolve those easily.

Let me take a stab at making some minimal changes but using your same test. Then I'll report back here and iterate as needed.

Fixes NanoComp#2578. The adjoint gradient of a MaterialGrid design region depended on how the cell was split into chunks, while the objective did not. Two load-bearing changes, both local; no communication is added or needed. 1. src/dft.cpp -- dft_chunk's `persist` pad clamped to fc->gv.little_corner()/big_corner(), which are on the centered grid. Component c runs from little_corner()+iyee_shift(c) to big_corner()+iyee_shift(c) (cf. LOOP_OVER_VOL), so in a yee-shifted direction the clamp truncated the topmost node -- exactly the ghost node the adjoint restriction stencil reads, and one that step_boundaries() already keeps current. The padding was never unable to reach the halo; it simply did not record it. Measured at 2 chunks under MPI with smoothing: 244 stencil lookups fell outside the recorded box, and all 244 were inside the owning fields chunk. The same clamp also left `is` off the component lattice, which desynchronized the LOOP_OVER_IVECS counter from grid_volume::index() (index() subtracts iyee_shift, vec.cpp:476; idx0 does not), so the adjoint DFT was read from the wrong element -- Ex with a 3-way split, ip=(-13,2,-4), is=(-14,0,-6): loop index 64 vs correct index 8. 1365 of 2496 points were misindexed at 3 chunks, 0 at 1 chunk. 2. src/meepgeom.cpp -- material_grids_addgradient() paired forward and adjoint chunks by position in the per-component next_in_dft list. A fields chunk contributes a dft_chunk for a component only where the monitor overlaps that component's owned grid, so the lists can differ in length and order. Pair on (fc, sn, shift), the tuple loop_in_chunks() iterates over, and drop the abort on unequal counts, which is a legitimate configuration. This only shows up at higher process counts: with (1) alone, a 1.7um design region still regressed by 3.3e-02 at -np 8. Also included, each a latent defect rather than something the tested configurations require: index dft[] by position rather than by the loop counter in material_grids_addgradient() and dft_chunk::norm2() (the latter feeds dft_norm() -> stop_when_dft_decayed(), so the solver's own termination criterion was affected); and bounds-check neighbor lookups per dimension rather than on the flat index, since in 3d a point one pixel off in y still yields an index inside [0,N) and silently aliases to an unrelated voxel. Validation. Gradient at num_chunks=N vs 3N over ten configurations (baseline 3d, smoothing off, offset design region, incommensurate resolution 13, coarse resolution 8, thin side padding, wide 1.7um design region, three frequencies, 11x11 grid, 2d), each at 1, 2 and 8 processes: all 30 combinations agree to <= 8e-12, most to ~1e-13 or better. On master nine of the ten fail in serial by 20-62%. 2d passes on master, which is why this went unnoticed -- every adjoint test upstream is 2d or cylindrical. Correctness as well as consistency: test_subpixel_3d (3d adjoint-vs-finite- difference) 5/5, test_adjoint_solver 13/13, plus test_adjoint_cyl, test_adjoint_utils, test_dft_fields, test_dft_energy. The regression test is taken unmodified from NanoComp#3264 (jin-castle). It forces the bug in serial via num_chunks, so it runs in every CI job rather than only the MPI ones. Refs NanoComp#2578, NanoComp#3264.

jin-castle force-pushed the pr/np-invariant-adjoint-gradient branch from 299afb7 to 6d39640 Compare August 9, 2026 02:38

inversedesignlab added 2 commits August 9, 2026 15:27

Apply clang-format to adjoint gradient changes

9abaecd

jin-castle mentioned this pull request Aug 9, 2026

Fix 3D MaterialGrid subpixel smoothing: kernel normalization (#3258) and gradient consistency (#3259) #3263

Merged

stevengj reviewed Aug 20, 2026

View reviewed changes

smartalecH mentioned this pull request Aug 21, 2026

[1/2] Fix adjoint gradient chunk-dependence without gathering the DFT monitors #3274

Merged

Navigation Menu

Uh oh!

Make the adjoint design-region gradient independent of the MPI chunk division - #3264

Make the adjoint design-region gradient independent of the MPI chunk division#3264
jin-castle wants to merge 4 commits into
NanoComp:masterfrom
jin-castle:pr/np-invariant-adjoint-gradient

jin-castle commented Aug 9, 2026

Uh oh!

stevengj commented Aug 14, 2026

Uh oh!

stevengj commented Aug 14, 2026

Uh oh!

lxvm commented Aug 18, 2026 •

edited

Loading

Uh oh!

stevengj commented Aug 20, 2026

Uh oh!

stevengj Aug 20, 2026

Uh oh!

smartalecH commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

jin-castle commented Aug 9, 2026

Uh oh!

stevengj commented Aug 14, 2026

Uh oh!

stevengj commented Aug 14, 2026

Uh oh!

lxvm commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

stevengj commented Aug 20, 2026

Uh oh!

stevengj Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

smartalecH commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

lxvm commented Aug 18, 2026 •

edited

Loading