EnzymeHLOOpt: a gather of a broadcast along its batching dimensions reads the source

[wsmoses]

gather(broadcast_in_dim(x), idx) where every dimension the broadcast adds is an operand batching dimension of the gather reads the same tensor in every batch. It equals a gather from x with those batching dimensions dropped. The matching start-indices batching dimensions become plain batch dimensions, so the result type is unchanged. Batching dimensions that belong to x itself are kept.

Conditions:

• the broadcast keeps its source dimensions in order without expanding them;
• every dimension it adds is an operand batching dimension.

Otherwise, e.g. when an added dimension is indexed through start_index_map, the gather stays.

New pattern: gather_of_broadcast_batching_dims, registered in the enzyme-hlo-opt set, the transform ops, and the exec-time pattern list in EnzymeXLA.cpp.

Motivation: batching a loop's gather from a buffer the loop doesn't change creates such a broadcast. greedy_while_loop_batch_fission broadcasts every operand defined outside the loop (DEFINED_OUTSIDE_WHILE), and XLA materializes it: one copy of the buffer per iteration. MFEM's bench_assembly_levels BP5ELEMENT made tensor<39304x943296xf64> (296 GB) and tensor<39304x3x3x6367248xf64>, and ran out of GPU memory compiling (68 GB of executable arguments). The minimal set of patterns that reproduces the broadcast is greedy_while_loop_batch_fission alone. With this pattern, the largest tensor in that module is 76 MB.

Test: new gather_of_broadcast_batching_dims.mlir with full-line goldens:

• one batching dimension;
• several broadcast dimensions plus a window and a source batching dimension, which stays;
• a broadcast dimension that is indexed, where the gather stays.

The first two match the original on JAX CPU exactly. Lit: 1503/1503 pass.

Merging with #3444 conflicts textually on the adjacent registration lines (TransformOps.td, the EnzymeXLA.cpp list, the pattern set); keep both.

🤖 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 (fb49aed) to head (574d027).
⚠️ Report is 28 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3479      +/-   ##
==========================================
- Coverage   29.92%   29.73%   -0.19%     
==========================================
  Files         241      241              
  Lines       48628    48655      +27     
==========================================
- Hits        14551    14468      -83     
- Misses      34077    34187     +110     

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