Repository navigation
Conversation
PJRT exposes no raw buffer-to-buffer copy, so cudaMemcpyDeviceToDevice staged every copy through the host: a read back, a host buffer and a write, which MFEM's conjugate gradient does several times an iteration. A copy on the raw device pointers would not see the uses PJRT still has in flight, so the copy is an executable through the exec path, which orders it with the rest and addresses views and shared buffers as for any kernel: one per size, cached under a stable module string. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
…to it A kernel's scalar argument reached the device as an allocation, a synchronous copy of four or eight bytes (a full stream synchronization, ~50 us) and a free, several times per kernel. A small allocation now waits for its first use: the host-to-device copy of the whole of it makes the buffer in one call that does not wait (BufferFromHostBuffer, with the copy as the buffer's definition event, so what runs after it waits), and any other first use makes it uninitialized as before. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
…small host writes as bytes Each kernel launch of the exec path ran its executable at once and took a future for it, which PJRT fills through a callback on the compute stream; the solver's dozen launches per CG iteration each paid that, and the host's scalar writes (a kernel's coefficient) each became a device buffer at once. The launches now queue (up to 64, or until the host reads or writes a buffer, which runs what is pending first), and run through ExecuteSharded with fill_future=false: what runs on the device after an executable is ordered by its buffers' events, and a copy to the host waits for the buffer it reads, so nothing on the host waits for an executable. An allocation freed while an executable using it is pending is kept, whole, until that has run. A small allocation (64 bytes or less) the host fills whole keeps its bytes instead of making a buffer: the buffer is made when a queued executable is run with it, and a read of it back comes from the bytes with no device round trip (mfem's Vector::operator=(real_t) on a scalar, read back by the host as a norm or a coefficient). An in-place write into part of an allocation stays a staged buffer plus the copy executable. ex1 (star, p=3, PA, 400 PCG iterations) warm solve: 0.37-0.41 s with REACTANT_EXEC_UNROLL=8, from 0.45-0.5 s; the residuals are identical to the synchronous runtime's on star and fichera. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
…tream executor client On the GPU the exec path's kernels went through PJRT buffers: each launch an ExecuteSharded with its argument holds, events and host callbacks on the stream, each host write a buffer made from host memory, each read a literal copy, each free a deferred release. A PJRT execute of a trivial executable costs 17 us of host time against 5 us for the LocalExecutable it wraps run directly (and the solver's kernels, with their arguments, cost 40 us each), and the GPU ran 54 us of kernels per CG iteration of a millisecond. On a stream executor client the runtime now keeps raw device memory: an allocation is the device allocator's memory, an executable is compiled with the LocalClient and run with RunAsync on the compute stream over that memory (a written argument donated, so its result is in place; should XLA not keep it there, the allocation takes the result's memory), a host-to-device copy goes through a pinned ring on the stream, a device-to-host copy waits for the stream, a device-to-device copy is a stream memcpy. The stream orders everything, so freed memory returns to the allocator at once. No PJRT buffer, event or host callback per call. Other clients (the CPU one) keep the PJRT buffer path; REACTANT_EXEC_PJRT_BUFFERS=1 keeps it on the GPU too. ex1 (star, p=3, PA, 400 PCG iterations) warm solve: 0.14 s, from 0.39 s (native CUDA on the same device: 0.037 s); fichera likewise; the residuals are identical to the PJRT path's. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
…utes A while loop whose bound the kernel computes from its inputs (the longest row of a CSR structure, for mfem's ElementRestriction::MultTranspose) has a trip count XLA does not know, so XLA reads the loop's predicate back to the host every step: six synchronizations per CG iteration, and six runs of the loop's body, three kernels each. The generic executable now returns such bounds after its results (a scalar integer the condition compares the induction variable against, neither constant nor loop-carried, computed at the function's top level). Once read, a specialization is compiled with the bounds as constants, the exec path's optimization run again on it (REACTANT_EXEC_UNROLL unrolls the loop, and the unrolled gathers fuse), and it runs instead for as long as the inputs the bounds derive from are the allocations observed, unwritten since (each allocation counts its writes); another observation compiles another specialization, four at most. REACTANT_EXEC_NO_SPECIALIZE=1 turns it off. Also, a small raw allocation freed (256 KB or less) is kept, by size, for the next of that size: a solver allocates and frees the same scalars and workspaces every iteration. ex1 (star, p=3, PA, 400 PCG iterations) warm solve: 0.096 s, from 0.14 s (native CUDA on the same device: 0.037 s); 11 kernel launches per iteration, from 28, and no stream synchronization but the two dot products'; the residuals are identical. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
Each allocation's handle was an address range mapped inaccessible for it and unmapped at its free: a system call each, and the unmapping a TLB shootdown across the compiler's threads -- the solver's dozen allocations an iteration cost it a fifth of its time. One range (1 TB, inaccessible, reserved at the start) is carved by size class instead, a power of two of 4 KB at least, each class keeping the handles freed for the next allocations of its size; a range of its own only once the reservation is spent. A stray dereference still faults, and pointer arithmetic on a handle still stays within its own range. ex1 (star, p=3, PA, 400 PCG iterations) warm solve: 0.077 s, from 0.096 s (native CUDA on the same device: 0.037 s); the residuals are identical. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
…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. The specialization of loop bounds runs the same. Also, REACTANT_EXEC_DUMP prints each module after the pipeline as well. Needs Enzyme-JAX with gather_scatter_constant_forward (#3444) and gather_of_masked_scatter (#3445). 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 at 0.076 s; 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 Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
This was referenced Oct 9, 2026
Member
Author
|
Closing: no separate pattern list in the runtime. The folding rounds this ran before the main list (constant propagation of lane index arithmetic, then gather_scatter_constant_forward / gather_of_masked_scatter / scatter_of_scatter_simplify) need to happen within enzymexlaGetTransformPassesList instead. Anything the main list lacks for that (e.g. folding add/and/or/pad of constants beyond the zero/one simplifications, or forwarding before batching) is a follow-up there. |
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
1 participant
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.
Stacked on #3429. Needs Enzyme-JAX with
gather_scatter_constant_forward(EnzymeAD/Enzyme-JAX#3444) andgather_of_masked_scatter(EnzymeAD/Enzyme-JAX#3445) — the pattern list names them, so this waits for those and the bump.A raised kernel's lane index arithmetic is over constant tensors (
or/add/pad/selectof constants, broadcasts), which the exec pattern list does not fold (its simplifications only drop zeros and ones) and StableHLO's folder only partly (nopad, 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
stripMarksOnConstants),stablehlo-aggressive-folder, which also rungather_scatter_constant_forward,gather_of_masked_scatterandscatter_of_scatter_simplify: with the indices constant, the scratch is forwarded while the lanes are still a loop.The loop-bound specialization (#3428) re-runs the same.
REACTANT_EXEC_DUMPalso prints each module after the pipeline now.ex1 (p=3, PA) warm solve on an RTX 5090: fichera 0.46 s, from 0.60 s (native CUDA 0.13 s; the apply module 2.0 → 1.5 ms per call by NVTX); star unchanged. Residuals identical on star; on fichera they differ only past iteration ~230, at residuals ~1e-13, the same round-off profile the XLA build already has against native (fusion order). MFEM GPU suite with this (+ #3432 and the two Enzyme-JAX PRs): 74/74 (sweep final64, 2026-10-09).
🤖 Generated with Claude Code
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD