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
… the pattern list A kernel whose thread block has a run-time extent (MFEM's generic fallback kernels) is raised with its lanes padded to a bound -- 24, or 10 -- and every store from the padding masked. Once the runtime makes the extent a constant (Q1D = 3), the masks fold to constants and most of each lane dimension is dead, but the kernel still computed and stored all of it: assembly_levels BK6PARTIAL ran at 107x native, and BP5ELEMENT's element assembly ran out of memory. The folding rounds before the list now also cut each masked store to its live lanes (scatter_masked_index_slice) and move the cut up its operands (the slice patterns), through the loops that accumulate them (while_slice_elementwise_carry) and the lane iotas the indices are built from (slice_iota), dropping the stores left with nothing to write (scatter_out_of_bounds_noop). Slices move up past transposes (slice_transpose) here; the list itself moves transposes up past slices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
wsmoses
marked this pull request as draft
October 12, 2026 19:19
Member
Author
|
Draft until measured: every pattern added here except |
Member
Author
|
Closing: no separate pattern list in the runtime. Every pattern this added is already in enzymexlaGetTransformPassesList, except slice_transpose (propagate-down only). If anything is missing when only the main list runs, it belongs in that list (or its ordering), not here. |
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 #3434 (it extends that PR's folding rounds; the diff against main includes the stack below it).
MFEM's generic fallback kernels launch with a thread block whose extent is only known at run time, e.g.
SmemPAVectorDiffusionApply3DwithMAX_T1Dregisters, and the element-assembly kernels. The raiser pads those lanes to a bound (24 or 10) and masks every store from the padding. Once the runtime makes the extent a constant (Q1D = 3), the masks fold to constants and most of each lane dimension is dead. The kernel still computed and stored all of it:Change: the folding rounds before the list also cut the padding, with existing Enzyme-JAX patterns:
scatter_masked_index_slice, extended for this in ScatterMaskedIndexSlice: cut stores into any target, and bound the box with constant masks Enzyme-JAX#3483);while_slice_elementwise_carry, WhileSliceElementwiseCarry: carry only the box of an accumulator the loop's readers slice Enzyme-JAX#3484) and the lane iotas the indices are built from (slice_iota, SliceIota: a slice of an integer iota is an iota of the slice's shape Enzyme-JAX#3485);scatter_out_of_bounds_noop).Here slices move up past transposes (
slice_transpose); the main list moves transposes up past slices (transpose_slice), so the two aren't run together.Results (RTX 5090; ms per call; native is clang CUDA):
On the BK6 and BP5 modules, the cut modules' outputs are bit-identical to the uncut ones on the first 300 and 200 elements.
Tried and left out:
Depends on: EnzymeAD/Enzyme-JAX#3483, #3484 and #3485 (merged), plus what #3434 already needs. CI on this stack also needs the BUILD dependency fix for #3427 and #3432 (
pjrt_stream_executor_client,gpu_collectives), which is prepared locally and waiting on a force-push.🤖 Generated with Claude Code
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD