Speed up connect_the_chunks for large numbers of chunks

[jin-castle]

fields::connect_the_chunks() tests every not-owned boundary point of every
chunk against every other chunk, on every process. The point scan is the
expensive half and it runs num_chunks times per chunk, so the cost grows as
O(num_chunks² × points-per-boundary). PML splitting pushes the chunk count well
past the process count, and any workflow that rebuilds fields repeatedly —
reset_meep() in an optimization loop — pays it every iteration.

On a 2.2M-voxel 3D cell on a dual-EPYC-9554 node (128 physical cores), the
first step() after init_sim(), which is where the connection tables get
built, costs 12.3 s at 64 processes and 40.9 s at 128, against roughly 1 ms per
timestep of actual stepping. It grows superlinearly, so it gets worse exactly
where you were adding nodes to make things faster. (Those two figures are from
an earlier run on that machine; everything below is measured on this tree.)

This prunes the candidate chunk pairs before the point scan. The connection
tables come out identical.

What changed

Chunk-pair pruning. A chunk_candidates[i] list is precomputed once per
call from a bounding-box test, and both point-scan loops iterate that instead
of for (int j = 0; j < num_chunks; j++).

Skipping chunks with no local pair. chunk_needed[i] is false when no
candidate pair (i, j) involves a process-local chunk. j == i is always a
candidate, since a box intersects itself, so this is true whenever chunk i is
mine; it only drops chunks whose every candidate pair is remote-to-remote,
which the loop body's (chunks[i]->is_mine() || chunks[j]->is_mine()) guard
rejects anyway. Those chunks previously had their boundary points scanned and
filtered one at a time.

Sparse comm-buffer allocation. Buffers were allocated for all
num_chunks² × num_field_types pairs, most of them zero-length. They are now
allocated only for pairs present in comm_sizes.

That third one is the only change that alters state rather than the order
things are computed in, so to be explicit about why leaving the rest NULL is
safe: NULL is already the value the constructor writes (src/fields.cpp:66-68,
:120-122), the destructor handles it (:136), and every read of
comm_blocks[ft][pair_idx] goes through comms_sequence_for_field[ft], whose
operations are only created for pairs with comm_size_tot(f, pair) != 0
(src/boundaries.cpp:688-689). A pair without a buffer is never reached.

Why the box test cannot drop a real pair

The pruning is only sound if it over-admits. The inner loop's sole acceptance
test is chunks[j]->gv.owns(here), so it is enough to bound what here can
be:

1. It starts as a point p from LOOP_OVER_VOL_NOTOWNED(vi, ...), so
   p ∈ [little_corner(vᵢ), big_corner(vᵢ)] — call that box Bᵢ.
2. locate_component_point() → locate_point_in_user_volume() maps
   p ↦ S.transform(p + s, sn), with s over 0 and ±1 lattice vector per
   periodic direction and sn ∈ [0, S.multiplicity()).
3. owns(here) implies here ∈ [little_corner(v_j), big_corner(v_j)] = B_j.

So a nonzero contribution from (i, j) requires
S.transform(Bᵢ + s, sn) ∩ B_j ≠ ∅ for some (s, sn), which is what the
precomputation evaluates.

Two details keep that from being merely plausible. S.transform is a signed
permutation of the coordinate axes about a fixed center, so it carries a box to
a box — transforming the two corners and taking a per-direction min/max
recovers the image exactly, with no inflation or clipping. And both boxes are
dilated outward by 2 ivec units (one pixel) before intersecting, so the test
admits pairs the point scan then rejects; the per-point owns() checks are
untouched, so a false candidate costs a few wasted comparisons while a true
pair cannot be lost.

The complexity claim is narrow, though. What goes away is the
O(num_chunks² × points-per-boundary) term, replaced by
O(num_chunks² × 3^(#periodic dirs) × multiplicity) box tests plus
O(candidates × points). The pair enumeration is still quadratic — this is a
much smaller constant on the term that was dominating, not a change of order.

Measurements

Single process with forced num_chunks, so only the pruning is exercised and
the chunk_needed skip contributes nothing here. 8×5×4 µm cell, resolution 16,
timed as the first step() after init_sim():

chunks	before	after
16	2.87 s	1.23 s	2.3×
48	7.08 s	1.72 s	4.1×
96	11.66 s	2.19 s	5.3×
192	21.23 s	2.85 s	7.4×

Baseline is master at e7d46f2, patched is that same tree with only
src/boundaries.cpp changed, so nothing else moves between the columns.

Equivalence

Since the claim is that the tables are unchanged, the check is whether any
field anywhere differs. Whole-grid checksums — sum and max of |Ez|, |Ey|, |Hx|
over the entire cell after 700 steps — are bit-identical between the baseline
and patched builds across all 16 combinations of

• all-PML, periodic in y, mirror in y, periodic + mirror
• num_chunks = 3, 7, 16, 48

The periodic and mirror cases are the point of that list: the benchmark above
is all-PML, which reaches neither the lattice-shift enumeration nor the
S.transform() branch, and those are the two places the argument is doing real
work. The run length matters for the same reason — a short run leaves most
chunk interfaces exchanging zeros, where a wrong connection cannot show up at
all, so the comparison aborts if the field is still trivial when it is read.

Regression suite on the patched build: test_mode_coeffs, test_special_kz,
test_n2f_periodic, test_bend_flux and test_dispersive_eigenmode pass
serially; test_mode_coeffs and test_bend_flux pass under mpirun -np 4.

All of that is downstream evidence — identical fields imply identical tables
but do not check the tables directly. Building comm_sizes with the pruning
disabled and diffing it against the pruned build at each chunk division would,
and it is cheap; I am glad to add it as a regression test if you would like it
in here.

Limitations

The pad of 2 is a hardcoded constant. It is one pixel, sized for the not-owned
halo that LOOP_OVER_VOL_NOTOWNED walks, and it would be better derived from
that halo than written down, so it follows if the halo ever changes.

There is no spatial index, so the pair enumeration stays O(num_chunks²) box
tests. That is cheap at the chunk counts I measured, but it is the term that
survives, and at some larger count it would want a grid or tree over the chunk
boxes.

The 16-combination field comparison is single-process. Under MPI only the
regression suite ran, at -np 4 — and chunk_needed is the MPI-specific part,
so it has the least direct evidence behind it.

The cases cover one mirror plane and one periodic axis. Several simultaneous
mirrors raise S.multiplicity() and several periodic directions make the shift
enumeration combinatorial; both follow from the same argument, but neither was
run.

---

Reopened and rebased onto current master; the diff is unchanged from the
original.

🤖 Generated with Claude Code

[stevengj]

This is a reasonable sort of optimization at first glance, but I'm a little worried about anything that increases the complexity of the code that computes the boundary connections, which is already extremely complicated.

How much speedup do real applications get from optimizing this code, which usually only executes once at the beginning of the simulation?