[jinyouzhi]

Rebased on master (PRs #9010 and #9011 have merged; this branch is rebased on master at 5dc44e8, so the diff contains only this PR's changes).

ParamAffineMap._validate_destinations (added in #8922) checks the box-shaped pieces of a column split by comparing every box of a rank with every other one. validate() runs whenever a ZeRO, AutoEP or block-gather map is built and before every conversion, and this pass has no budget, so validation grows quadratically with the number of strided pieces per rank.

This PR sweeps the boxes with boxes_disjoint from #9011 first. Only a rank that fails the sweep takes the pairwise walk, which names the shared shard address, so error messages are unchanged.

The second commit moves the remaining hand-written geometry onto affine_geometry: _dest_box returns a LogicalBox built with coordinate(), the shared address comes from intersection() and dot(), and address_range() / overlaps_itself() move there from affine.py. The transfer executor reuses address_range() for its buffer-reach check and dot() for chunk offsets. No behavior or message changes.

validate() of 8 ranks with interleaved columns (64 rows; ms, single CPU process):

pieces per rank	before	after
128	29	4
512	441	17
2048	6654	68

Tests:

• Randomized check over 200 seeds: column blocks of row bands tiling a shard, half with one block shifted onto another. validate() refuses exactly the maps whose pieces write some shard address twice (enumerated independently in the test), and every refusal names an address that is written twice. The test checks this contract only, so it does not depend on which helper does the sweep.
• Mutation checks: making the sweep always report "disjoint" fails 48 of the 200 cases. Forcing every rank through the pairwise walk still passes all 200.
• All tests/unit/checkpoint/test_*affine* files: 839 passed, 1 skipped on 4× NVIDIA L20 at f6e1f6c. At ac3a329 on CPU (DS_ACCELERATOR=cpu, 4 ranks): 842 passed, 5 skipped.

AI assistance

Implementation, tests and this description were drafted with GitHub Copilot CLI and reviewed/validated by the author.

[Achyuthan-S]

Thanks @jinyouzhi, this is a nice speedup for the check I added in #8922, and falling back to the pairwise walk keeps the error messages exactly as they were.

I checked it against master (b9100c5) by running both on the same 20,000 random single-rank maps. Each map tiles a 1-D to 3-D shard grid with random cuts, then sometimes changes one piece: shifting it, swapping two pieces' destinations, moving a destination by 1–3 elements, starting it at a random address so it wraps past a row or leaves the shard, or giving a size-1 axis an arbitrary stride. validate() gave the same outcome on every map, including identical error text for all 6,732 refusals that name a shared address. Making the sweep never report an overlap fails 88 tests, and all affine test files pass on CPU with 4 ranks (840 passed).

One small suggestion on the new test, not blocking: its second assertion monkeypatches affine.boxes_disjoint to force the pairwise path, so it would break under another correct implementation that doesn't call that name. Asserting the contract instead, that the address named in the error is one that written contains twice, would keep the check without depending on the helper.

LGTM.

[jinyouzhi]

Thanks @Achyuthan-S for the differential run against master. Good point about the test: 06edb58 drops the boxes_disjoint monkeypatch and asserts the contract instead. Each refusal must name an address that written contains twice. Making the sweep always report disjoint still fails 48 of the 200 seeds, and forcing every rank through the pairwise walk passes all of them. All affine test files pass on CPU with 4 ranks (835 passed, 5 skipped).

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-13T16:06:10.188247Z	b2ddefe	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Achyuthan-S]

Re-checked at b2ddefe, merged with current master (8702b42). The new test asserts the contract now, and both of its mutation checks reproduce here: a sweep that always reports "disjoint" fails 48 of the 200 seeds, and forcing the pairwise walk passes all 200.

On 20,000 random single-rank maps (1-D to 3-D grids with shifted, swapped, nudged or wrapping destinations, arbitrary strides on size-1 axes, and non-box views), validate() gives master's outcome on every map, error text included. 5,449 of them reach the sweep, and 1,722 of those overlap and go through the pairwise walk. All eight affine test files pass on CPU with 4 ranks (906 passed).

This and #9052 touch the same import line and __all__ in affine_geometry.py; if this lands first, I'll rebase #9052 on top.

LGTM.

[jinyouzhi]

Thanks! @Achyuthan-S
Could you review it when available? Thanks @delock