Make the adjoint design-region gradient independent of the MPI chunk division - #3264
Open
jin-castle wants to merge 3 commits into
Open
Make the adjoint design-region gradient independent of the MPI chunk division#3264jin-castle wants to merge 3 commits into
jin-castle wants to merge 3 commits into
Conversation
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.
jin-castle
force-pushed
the
pr/np-invariant-adjoint-gradient
branch
from
August 9, 2026 02:38
299afb7 to
6d39640
Compare
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #2578 (adjoint gradients inconsistent with finite differences in 3D).
Symptom. The adjoint gradient of a
MaterialGriddesign region depends onthe 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):
f0identical to 8+ digits at every process count;-np 1and-np 8: 57% of|g|_max(smoothing off) / 18% (smoothing on), spread over ~80% of thedesign nodes in chunk-sized patches;
-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 theresult from per-chunk DFT views:
(
forward_dft_chunks[ci_forward][cur_chunk]); the per-component chunk setscan 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);index, so an out-of-range point in one dimension could alias to a valid
index of a different point instead of returning zero;
persistdft chunks is clamped to the owningfields chunk and cannot cross process boundaries, so the data those
lookups need does not exist at rank boundaries;
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_olddecomposition 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.
-np 8vs-np 1, 20×20×6, smoothing off-np 8/16/32/128vs-np 1(both smoothing cases)-np 16and-np 128(six grid configurations incl. incommensurate/offset/coarse/fine grids)test_subpixel_3d, adjoint E2E unchanged)🤖 Generated with Claude Code