EnzymeHLOOpt: forward reads of a scratch written through scatters at constant indices

[wsmoses]

A kernel's per-lane scratch (mfem's shared-memory sm[3][125] in the 3D PA diffusion apply, PADiffusionApply3D) is stored through by index tensors the raising knows and loaded back the same way; once the lanes are batched, both become a stablehlo.scatter and a stablehlo.gather on a tensor the lane owns — 18 scatters and 63 gathers per element in that kernel's exec-time module, each a pass over the scratch, where native code keeps the values in registers and shared memory. On fichera (p=3, 802k dofs) the XLA iteration runs 144 kernels against native's 13, and that kernel is most of it (four input_scatter_fusions ~130 µs each, a 198 µs transpose).

Where every index is constant and the scatters overwrite one element per index (the setindex combiner), what each read reads is known: the update of the last scatter in the chain writing it, or the chain's input.

• GatherScatterConstantForward: a gather at constant indices (one element per index tuple) of such a chain becomes gathers of those updates directly, the unwritten positions read from the chain's input by the same gather, selected together by constant masks where there are several sources.
• ScatterChainConstantToGather: the top of a chain read any other way (whole, through a pad) becomes the same over the whole operand, the input where unwritten. A lone scatter into a value that is not a constant is left alone: it updates in place, where the select would read its whole input (e.g. 4 writes into a 100-element argument).
• A scatter writing a position twice ends a chain (which update wins is not defined), as does one with a non-constant index tensor.

On that kernel's body (EXEC_DUMP of the fichera run, replayed with the runtime's list), the scatters go from 18 to 5, the five left being the stores to the output E-vector — once the index arithmetic is folded: the lanes' index tensors are or/add/pad/compare/select over constants, which the exec list never folds (add_simplify and friends only strip zeros), so the pattern sees nothing without stablehlo-aggressive-folder run first. The runtime side (adding that folder to the exec pipeline) is a Reactant.jl change, measured next.

MFEM GPU suite with it (and #3445, the folding rounds): 74/74 (sweep final63, 2026-10-09).

Test gather_scatter_constant_forward.mlir: a permutation read, a read of unwritten positions, a two-scatter chain with an overwrite, a chain read whole, a scatter writing twice (unchanged). Registered as gather_scatter_constant_forward in the scatter group and the runtime list.

Rebased onto main with #3442/#3440/#3445. Main's tests that this changes are regolded: divscatter, autobatching/elementwise_loop_multiindex (a lone scatter into a constant zero becomes select(mask, gather, 0)), and raising/raise_lane_private_scratch (@lane_memref's scratch is forwarded; @lane_atomic's zero-initialized first store folds away). gather_scatter_constant_forward.mlir: @whole now writes into a constant; the new @whole_in_place checks that a scatter into an argument stays a scatter.

Rebased on main 73603ba (2026-10-12).

• Compile-time blowup, now resolved: on MFEM bilininteg_vecdiffusion_pa.cpp, this PR's ScatterChainConstantToGather produced IR on which two main patterns, DUSDUSSubsuming and SliceOfDynamicUpdate, traded slices forever. One function's enzyme-hlo-opt never finished (> 10 min). EnzymeHLOOpt: DUSDUSSubsuming moves slices only when it replaces the DUS #3470, now merged, fixes that cycle; with it, the same input takes 15 s.
• Local lit: 1509/1510; //test:neuralgcm_test runs out of memory on our 32 GB GPUs, as it does on main.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 29.73%. Comparing base (6775ac9) to head (c8e00a6).
⚠️ Report is 24 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3444      +/-   ##
==========================================
- Coverage   29.73%   29.73%   -0.01%     
==========================================
  Files         241      241              
  Lines       48677    48655      -22     
==========================================
- Hits        14476    14468       -8     
+ Misses      34201    34187      -14     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[wsmoses]

Closing: measured strictly worse with the exec runtime's main list only (no separate folding rounds). BK6PARTIAL/2 takes 403 ms with these patterns against 64.6 ms without them, BK6PARTIAL/4 159 vs 30.3 ms, and ceed BP3/2/96 115 vs 84.0 ms. The general primitive suggested in review, a gather that reads nothing a scatter writes reading the scatter's input, is #3490.