EnzymeHLOOpt: the live box of a masked scatter through broadcasts and constant masks

[wsmoses]

ScatterMaskedIndexSlice (#3191) finds the live box of a masked store from the mask's compare(iota, c) terms. After exec-time specialization the extent those compares test against is a constant, and the runtime's const-prop folds each compare into a constant prefix along its lane dimension (dense<[true, true, true, true, false, ...]>), broadcast onto the lane grid and and-ed — a form the pattern did not read, so the scatter stayed at the padded size.

Now each term is read through the broadcast_in_dims that put it on the grid (a dimension map from the term's dims to the mask's), and a 1-D constant term gives its interval of true values (a splat false empties the box).

The case: MFEM's PADiffusionSetup kernels padded to MAX_Q1D (32×32 lanes per element) run with Q1D=4 — 16 live lanes of 1024. In ex1's exec-time module for it the scatters now slice to 20480x4x4, and (with slice_elementwise, #3440's slice_gather, and #3439's fold of the fully dead scatters) the arithmetic feeding them shrinks with them.

A dimension of extent one that a broadcast expands holds one value for the whole grid dimension (a flag the kernel was specialized on, broadcast_in_dim(dense<[true]> : tensor<1xi1>)): it gives no interval — a 1-element constant or a 1-extent iota compare is read as holding everywhere or nowhere. A term false everywhere is carried as an explicit "nowhere" flag, since a scatter of one point has a zero-dimensional lane grid with no box to carry it in. (The first version of this PR wrote the empty box into the grid extents, so a single masked-off store — the batched form of an scf.if on a specialized-false flag — was read as live and its mask dropped: PA Convection, PA DG Diffusion and L2 Assembly Levels failed in the GPU suite; found by bisecting the runtime and an oracle comparison of the test's exec modules.)

Tests: @constant_mask in scatter_masked_slice.mlir (the folded form of @box: rows below 4, columns below 3 → one 4×3 dynamic_update_slice) @expanded_term (a 1-extent flag term broadcast over the columns: the box is the rows' alone), and @point_off (one point under a false mask: nothing written).

🤖 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.96%. Comparing base (2d3ff23) to head (c4358b7).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3441      +/-   ##
==========================================
+ Coverage   29.63%   29.96%   +0.33%     
==========================================
  Files         240      240              
  Lines       48496    48546      +50     
==========================================
+ Hits        14371    14548     +177     
+ Misses      34125    33998     -127     

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

MFEM GPU unit suite (74 tests) with the runtime built from Enzyme-JAX main + #3437 #3439 #3440 #3441 #3442 and Reactant.jl #3423 + #3425, over objects from main + the open affine-cfg PRs (#3436 #3438 among them): 74/74 (sweep final55, 2026-10-08 17:38).

ex1 (Poisson, order 3, PA, PCG to 1e-12), timed solves in one process (JIT excluded), native CUDA for reference:

• star (185k dofs, 400 iterations): 0.64 s (native 0.047 s) — the morning's state ran ~70 iterations per 600 s
• fichera (802k dofs, 270 iterations): 1.08 s (native 0.14 s)