Linkable runtime: compile a kernel again in chunks when it would not fit

[wsmoses]

Stacked on #3435 (the linkable-runtime stack). Needs EnzymeAD/Enzyme-JAX#3481 for the parallel_while_to_batched_scatter_chunked pattern op. The chunked retry only runs for a module that would otherwise run out of memory, so before an Enzyme-JAX bump that case fails in the pipeline instead of at allocation.

XLA allocates all the values an executable keeps as one temporary. A raised kernel's batched parallel loops keep a value per iteration for many of them, so past some problem size the executable can't run. MFEM ceed at order 4, side 240 (216,000 elements): BP3 needs a 35.79 GiB temporary, BP4 47.29 GiB.

Change:

• After compiling (raw path), the runtime sums the GPU executable's preallocated temp buffers.
• If they exceed the budget, it compiles the module again from the start. The budget is 90% of what the allocator has left; else half the device's memory; REACTANT_EXEC_TEMP_BUDGET overrides.
• The retry uses a pattern list with parallel_while_to_batched_scatter_chunked(N) in place of parallel_while_to_batched_scatter, and without greedy_while_loop_batch_fission, which lifts ops out of a loop over all its iterations at once and was the source of the 47 GiB.
• N starts at ⌈temp/budget⌉ (at least 2) and keeps multiplying while the temporary is over budget, up to 1024; past that, the attempt with the smallest temporary is used. More chunks can also let loops batch that a whole batch couldn't, so the temp isn't monotone: ceed BP4/2/160 needed 22.9 GB with 1 chunk, 34.8 GB with 2, and 9.6 GB with 6.
• The specialization path (specializeRaw) does the same. Its pipeline would otherwise rebuild the oversized executable from the already-optimized text.
• The GPU-executable code is under REACTANT_CUDA/REACTANT_ROCM, and BUILD declares its deps.
• The request is kept per module and argument signature. A module that fits is compiled exactly as before: ceed BP4/4/96 148 ms and nlvc AssembleGradDiagonal3d/2/96 3.3 ms, unchanged.

Measured (RTX 5090; Enzyme-JAX with #3481; native = clang CUDA):

case	before	this branch	with the exec-folder rounds (#3434)	native
ceed BP3/4/240	OOM, 35.79 GiB	1030 ms	642 ms	230 ms
ceed BP4/4/240	OOM, 47.29 GiB	—	2736 ms with REACTANT_XLA_MEM_FRACTION=0.92	881 ms
ceed BP4/2/160	OOM, 34.8 GiB	—	1162 ms (6 chunks)	528 ms
ceed BP6/2/240	OOM	—	2476 ms	1530 ms

For BP4/4/240, the chunked temporary is 9.66 GiB, but the default 75% pool holds 14.5 GiB of arguments.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

[wsmoses]

Fixed the test-localjll build: the GpuExecutable code is now under REACTANT_CUDA/REACTANT_ROCM (tempBytes is 0 elsewhere), and the BUILD deps declare @xla//xla/service/gpu:gpu_executable (cuda and rocm lists) and @xla//xla/stream_executor/integrations:tf_allocator_adapter.