Repository navigation
Conversation
…ern list A raised kernel's lane index arithmetic is over constant tensors (or/add/pad/select of constants, broadcast), which the exec pattern list does not fold (its simplifications only drop zeros and ones) and StableHLO's folder only partly (no pad, splat broadcasts only) -- and neither touches a constant or an op the raising marked (enzymexla.non_negative and the like). So a lane's scratch, stored and loaded through those indices, stayed a scatter and a gather on a tensor the lane owns, and once the list batched the lanes each access was a pass over memory: in mfem's 3D PA diffusion apply, 18 scatters and 63 gathers per element, most of fichera's GPU time. Before the list, the exec pipeline now strips the raising's marks from constants and ops over constants only, and runs three rounds of the list's constant propagation alternating with StableHLO's folder; with the indices constant, gather_scatter_constant_forward, gather_of_masked_scatter and scatter_of_scatter_simplify forward the scratch while the lanes are still a loop. Also, REACTANT_EXEC_DUMP prints each module after the pipeline as well. ex1 (p=3, PA) warm solve: fichera 0.46 s, from 0.60 s (native CUDA: 0.13 s; the apply module 2.0 -> 1.5 ms per call); star unchanged; the residuals are identical on star, and on fichera differ only past iteration ~230, at residuals ~1e-13, as the XLA build's already do from native's. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
wsmoses
force-pushed
the
pb/xla-exec-folder
branch
from
October 15, 2026 04:06
8a09f31 to
09330d7
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ported onto main from the runtime stack the MFEM numbers were measured with (there it sits on #3429; here it needs none of the raw-path PRs).
A raised kernel's lane index arithmetic is over constant tensors (or/add/pad/select of constants, broadcast), which the exec pattern list does not fold (its simplifications only drop zeros and ones) and StableHLO's folder only partly (no pad, splat broadcasts only) — and neither touches a constant or an op the raising marked (
enzymexla.non_negativeand the like). So a lane's scratch, stored and loaded through those indices, stayed a scatter and a gather on a tensor the lane owns, and once the list batched the lanes each access was a pass over memory: in MFEM's 3D PA diffusion apply, 18 scatters and 63 gathers per element, most of fichera's GPU time.Before the list, the exec pipeline now strips the raising's marks from constants and ops over constants only, and runs three rounds of the list's constant propagation alternating with StableHLO's folder; with the indices constant,
gather_scatter_constant_forward,gather_of_masked_scatterandscatter_of_scatter_simplifyforward the scratch while the lanes are still a loop. Also,REACTANT_EXEC_DUMPprints each module after the pipeline as well.Measured when written (Oct 9, with the raw path): ex1 p=3 PA warm solve, fichera 0.60 → 0.46 s (apply module 2.0 → 1.5 ms per call), star unchanged; residuals identical on star. The patterns it relies on (Enzyme-JAX #3444, #3445) are in the current pin.
Draft: I am measuring whether the same result can be had without this (the folding done inside the pattern list instead), per the question of how much of the runtime stack the numbers depend on.
🤖 Generated with Claude Code
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD