Raising: dead lanes must not scatter at all

[wsmoses]

A store under a lane mask raises as gather-select-scatter, writing the original value back on masked-out lanes. But a dead lane's index expression is unconstrained — strip-mine bound clamps (max(extent, tid + bd)) collapse dead lanes onto live lanes' slots — and scatter applies duplicate indices in unspecified order, so the stale write-back can land after the live lane's update and clobber it.

Found via mfem's QuadratureInterpolator under the XLA backend (#2968): the dynamic 2D tensor-product eval kernel produced partially-zero, element-varying results; a PJRT harness on the raised module reproduced it deterministically, and redirecting dead lanes out of bounds fixes it bit-exactly.

The fix selects dead lanes' indices to -1; out-of-bounds scatter updates are dropped, so masked-out lanes write nothing. This also removes the gather of the original values.

This PR now carries only the raising change. The enzyme-hlo-opt half that resolves the masked index ahead of the single-point scatter-to-DUS rewrite (ScatterMaskedIndexSimplify) is split out into its own PR, and the bounds guard on that rewrite is #3022. Ordering matters: with this change a masked-out lane's correctness rests on the scatter dropping its out-of-bounds index, so the DUS rewrite must not clamp it. This should land with, or after, one of those two.

The previous form let enzyme-hlo-opt fold a full-grid masked store into a pure select; with the index-select form the scatter survives in affine_to_stablehlo_masked_scatter.mlir. Teaching the dense-scatter fold to look through the index select would recover that.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

[wsmoses]

@avik-pal @Pangoraw can I get a proper review/in depth audit of this [unreviewed by me so far to be clear] vibed up PR.

In particular one bad consequence is that this breaks the scatter removal so we may need to update things

[wsmoses]

Pushed a required companion fix: the single-point scatter→dynamic_update_slice rewrite in enzyme-hlo-opt was unsound against the new out-of-bounds dead-lane indices — scatter drops OOB updates but DUS clamps the index back in range, so a masked-off store corrupted slot zero (caught by mfem's PAIdentityInterp, isolated with a buffer-level runtime diff + PJRT harness oracle). A new always-on pattern resolves the mask before those rewrites: scatter at the live index and write the original value back when masked off (a no-op store), which then lowers to a sound DUS. Lit test included.

[Pangoraw]

i think this can just be a reshape (since there should be a single additional trailing 1 element dim)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 30.49%. Comparing base (a1d5338) to head (e496b38).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2993      +/-   ##
==========================================
+ Coverage   29.44%   30.49%   +1.05%     
==========================================
  Files         238      238              
  Lines       48117    48576     +459     
==========================================
+ Hits        14166    14814     +648     
+ Misses      33951    33762     -189     

☔ 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.