[wsmoses]

Stacked on #3473 (the exec pipeline with its folding rounds); needs Enzyme-JAX #3624 (transposed_view_elementwise), #3634 (reread_barrier, a round of its own after the layout round) and #3627 (reduce_minor_dims, also in the layout round: ceed BP1/2 31.0 → 18.6 ms, BP6/2 159 → 115 ms, assembly_levels BP1PARTIAL/2 12.4 → 8.3 ms end to end, with #3626).

An elementwise op reading a transposed view of memory — an argument the list hoisted a transpose onto, through slices — computes in the view's layout, with the transpose on its result, where the transpose-propagating patterns take it into the broadcasts of the next contraction (transposed_view_elementwise, with the pushes that keep it there). Since Enzyme-JAX #3602 turned the reshape-sandwiched transposes of MFEM's diffusion apply kernel into plain ones, the list hoisted them onto the D tensor and XLA materialized a transposed copy of all of D (340 MB): ceed BP3/2 apply 3.0 ms, twice the Oct 13 time.

A round of its own after the list rather than a pattern in it: in the list it fires before the batching patterns have merged the components of a vector problem, and on the unmerged components it leaves transposes the merged form would not have (ceed BP4 apply 6.9 → 14 ms). On the converged modules of the ceed suite: BP3/2 apply 3.02 → 1.64 ms, BP3/4 1.64 → 0.97 ms, BP6/2 3.9 → 2.9 ms, BP1/2/4/5 unchanged, results bit-identical; the 400 exec modules of the MFEM benchmark suites converge through it. REACTANT_EXEC_NOLAYOUT=1 leaves the round out.

End to end (MFEM benchmark binaries, RTX 5090, 3 repetitions, dots off): ceed BP3/2 80.8 → 40.2 ms, BP3/4 48.3 → 24.4 ms, BP6/4 93.8 → 89.1 ms; assembly_levels BP3PARTIAL 29.7 → 15.9 / 30.8 → 17.2 ms, BK3PARTIAL 0.64 → 0.38 / 0.66 → 0.41 ms; every other case of the 107-case paper subset (assembly_levels, ceed, dg_amr, elasticity, nlvc, tmop, vector, virtuals) within noise.

Second round, reread_barrier (Enzyme-JAX #3634; REACTANT_EXEC_NOBARRIER=1 leaves it out): XLA's fusion nests a producer into every consumer and recomputes it per element of the reader; MFEM's order-2 contractions (reductions over 3, which XLA emits elementally) nest three deep, and its cost model prices the f64 recomputation at f32 rates, so on a GeForce the nested kernel is compute bound (103 µs for one fusion of BK5PARTIAL/2's apply; the whole native kernel is 120 µs). The round materializes a computed value read through a broadcast whose expansion is small and does not divide the warp size. Standalone: BK5PARTIAL/2's apply 469 → 308 µs per call; end to end (full 107-case paper subset, with and without the round): BK5PARTIAL/2 0.33 → 0.22 ms, BP5PARTIAL/2 13.8 → 9.9, ceed BP5/2 27.6 → 21.5, BK3PARTIAL/2 −7%, tmop AssembleGradDiagonalPA/1/2 1.85 → 0.67 ms; BK2PARTIAL/2 +8% (documented in #3634); everything else within noise.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD