Fuse adaptive GPUTsit5 stage arithmetic

[ChrisRackauckas-Claude]

Ignore this draft until reviewed by @ChrisRackauckas.

#544 (the endpoint-stall fix this PR used to be stacked on) is merged, and this branch now merges master bacef856 (which includes #563, fixing the :missing_docs Documentation failure and the ensemblegpuarray_scalar_batch.jl failures that were red on the previous head). The diff against master is only commit 97c0e86 (Fuse adaptive GPUTsit5 stage arithmetic): the adaptive Tsit5 stage block and AGENTS.md.

What changed and why

This applies @muladd to the adaptive GPUTsit5/GPUTsit5IController stage combinations and embedded error estimate, which the fixed-step stepper already does (#529). The macro is scoped to that block, so the controller and endpoint logic keep their current rounding, and the function keeps @inline. It also adds an AGENTS.md note explaining why @inline must stay on kernel steppers.

Root causes of the three failures that made #529 drop adaptive FMA

I reproduced each one on current master before changing anything.

1. "Adaptive FMA + forced inlining timed out on LTS". An infinite loop at runtime, not a compile blowup. I reproduced it at Fuse fixed-step GPUTsit5 arithmetic and document controller accuracy #529's first commit (8d23077033) on Julia 1.10.12. SIGUSR1 on the stuck process showed it spinning inside step!, and tracing the steps showed t = 0.099999994f0 stuck with dt = dtnew = 3.7252903f-9 forever. The dtnew clamp tf - t - dt does not match the rounded t + dt, so a sub-half-ulp step never advances t. FMA only moved that test case onto a stall; master stalls on other inputs without FMA (31–53 of 1800 Float32 configurations per adaptive algorithm). Fix Float32 endpoint stall in adaptive GPU kernel steppers #544 (merged) fixes this for all steppers by clamping dtnew from the accepted time while keeping the length-based snap. Follow-ups Respect pending tstops in the adaptive next-step clamp #556 (merged) and Interpolate to tstops by step length and handle sub-dtmin stop spacing #561 refine that clamp around tstops; they are independent of this PR.
2. OpenCL "Batch solve failed; output times [0.0f0, 0.0f0]". Caused by removing @inline, not by FMA: function step! without @muladd reproduces it too. When step! is not inlined, the mutable GPUATsit5Integrator escapes into a call, so Julia heap-allocates it. The pocl device IR then contains call i64 @gpu_malloc(...) → gpu_report_oom → gpu_signal_exception. The OpenCL runtime's malloc returns null, so the kernel exits before writing anything. With @inline kept, the device IR has 0 malloc/gc_pool_alloc calls.
3. Tight Float64 Enzyme exponential gradient bound. Also caused by removing @inline, not by FMA; it reproduces without @muladd. It happens on Julia 1.12 only (LTS passes). The uninlined build gives dp[1] = 1.2214058645… (max rel. error 3.8e-6 > 1e-7), while the primal loss is identical. Central finite differences of the same primal give 1.22140275843 (h = 1e-5), which matches the inlined Enzyme result 1.22140275841 (rel. error 4.4e-10). So the uninlined gradient is wrong, which points to Enzyme mishandling the heap-allocated mutable integrator passed to a call. The bound itself is appropriate. With @inline kept and FMA applied, the Float64 adaptive result is 4.43e-10 max relative error, the same as master.

Why not other forms

• Removing @inline: causes failures 2 and 3 above.
• @muladd on the whole function: gives the same contract count, since the controller has only one a*b+c pattern (the tstops check). Scoping the macro keeps controller and endpoint arithmetic byte-for-byte unchanged, which rules out any interaction with step acceptance at tf.

FMA / contract counts: code_llvm / code_native of step!, Lorenz, CPU backend

Julia	T	adaptive master: FP ops / contract / x86 vfmadd	adaptive this PR	fixed-step (unchanged)
1.12.7	Float32	259 / 0 / 0	340 / 195 / 69	241 / 141 / 66
1.12.7	Float64	271 / 0 / 0	271 / 162 / 81	206 / 126 / 63
1.10.12	Float32	293 / 0 / 0	309 / 155 / 54	236 / 143 / 67
1.10.12	Float64	274 / 36 / 22	274 / 155 / 71	206 / 141 / 63

PTX from GPUCompiler.code_native for step! (sm_80, PTX ISA 7.8; this is compiled PTX, nothing was run on a GPU). Cells show fma.rn / unfused add|sub|mul:

GPUCompiler	T	adaptive master	adaptive this PR
1.23.0 (llc, no contraction)	Float32	0 / 328	90 / 155
1.23.0	Float64	0 / 262	81 / 100
2.8.1 (NVPTX fma_contraction=true)	Float32	88 / 152	100 / 135
2.8.1	Float64	87 / 88	87 / 88

So on the current CUDA stack (CUDACore ≥ 6.4 requires GPUCompiler ≥ 2.7, whose NVPTX back-end contracts by default, matching nvcc --fmad=true), master already gets most of the FMA in PTX, and this PR adds 12 more in Float32 and none in Float64. The larger effect is on GPUCompiler 1.x, and on backends that honour only explicit contract flags (OpenCL/SPIR-V, x86 CPU). Therefore FMA does not explain the reported GRADSOLVE Float64 gap on a current CUDA.jl stack.

CPU timing (CPU only, not GPU)

julia +1.12 benchmark/tsit5.jl, KernelAbstractions CPU(), 4096 Lorenz trajectories, reltol 1e-6, abstol 1e-9. Minimum of 20 samples, two alternating rounds, identical manifests. Shared host with load average about 31–35.

	master r1 / r2 (ms)	this PR r1 / r2 (ms)
Float32 adaptive	45.69 / 47.33	45.25 / 45.19
Float64 adaptive	47.06 / 47.72	47.33 / 47.40
Float32 fixed	215.73 / 215.34	215.82 / 215.45
Float64 fixed	215.85 / 226.37	216.80 / 216.41

No meaningful CPU change: the adaptive Lorenz step is not bound by stage FLOPs on x86. Final states agree to rounding (Float64 differs in the 15th digit; Float32 in the 6th–7th).

Verification

On head 2dd2fd4 (the FMA commit merged with master bacef856), Julia 1.12.7, GROUP=<name> julia +1.12 --project -e 'using Pkg; Pkg.test()':

GROUP=CPU     -> Tsit5 accuracy | 144 144; Adaptive endpoint termination | 132 132; Testing DiffEqGPU tests passed (exit 0)
GROUP=JLArrays -> Tsit5 accuracy | 84 84; Testing DiffEqGPU tests passed (exit 0)
GROUP=OpenCL  (pocl) -> Tsit5 accuracy | 84 84; Testing DiffEqGPU tests passed (exit 0)
GROUP=QA      -> Quality Assurance + JET static analysis | 75 75; Testing DiffEqGPU tests passed (exit 0)
GROUP=Enzyme  -> see below: fails on Enzyme v0.13.209, passes on v0.13.208 and earlier

Earlier, on head baf4560 (the FMA commit merged with master 5e8faee), from separate clean worktrees, with a fresh Pkg.test per group:

# JULIA_PKG_PRECOMPILE_AUTO=0, clean worktree per branch/version, 1.12 Enzyme jobs run one at a time
GROUP=CPU    julia +lts  -> Tsit5 accuracy | 144 144; Adaptive endpoint termination | 132 132; Testing DiffEqGPU tests passed
GROUP=CPU    julia +1.12 -> Tsit5 accuracy | 144 144; Adaptive endpoint termination | 132 132; Testing DiffEqGPU tests passed
GROUP=OpenCL julia +lts  (pocl) -> Tsit5 accuracy | 84 84; Testing DiffEqGPU tests passed
GROUP=OpenCL julia +1.12 (pocl) -> Tsit5 accuracy | 84 84; Testing DiffEqGPU tests passed
GROUP=Enzyme julia +lts  -> Enzyme ensemble gradients (Float64, adaptive=true) | 1 1; Lorenz parameter gradients | 2 2; Testing DiffEqGPU tests passed
GROUP=Enzyme julia +1.12 -> Enzyme ensemble gradients (Float64, adaptive=true) | 1 1; Lorenz parameter gradients | 2 2; Testing DiffEqGPU tests passed

Also: an OpenCL adaptive Lorenz probe (8 trajectories, Float32) completes on LTS and 1.12, with max |OpenCL − CPU| = 2.3e-5 / 1.8e-5 at t = 1. Runic --check, typos, and git diff --check exit 0. No assertion or tolerance was changed.

CI on head 2dd2fd4 — every red check is pre-existing or upstream, none PR-caused

• CUDA Tests (Julia 1) fails at test/enzyme/cuda_records.jl:9 (Allocate records Enzyme error), identical to master's bacef856 run — CUDA CI: test/enzyme/cuda_records.jl fails on released CUDACore (Enzyme undef-shadow fill! for non-numeric eltypes) #546 (upstream CUDACore Enzyme-rule bug). The previous head's CUDA failure (ensemblegpuarray_scalar_batch.jl:102, "Scalar indexing is disallowed") is gone — Fix #555 fallout: docs missing_docs and CUDA scalar-indexing test #563's fix is in, that testset now runs 55 55 on CUDA. Tsit5 accuracy runs 144 144 on CUDA.
• Tests - Enzyme (Julia 1) and Documentation both fail on Enzyme v0.13.209 (released 2026-10-01, after master's last green run resolved v0.13.205): MethodError: no method matching augmented_primal(::RevConfigWidth{..., (false, true, true), ...}, ::Const{typeof(DiffEqGPU._kernel_transfer)}, ...) — Enzyme now passes the backend as Duplicated{CPU} where the custom rule in ext/EnzymeExt.jl declares backend::Const. Verified locally on this head: the full test/enzyme suite (imports, transfers, gradients incl. Float64 adaptive, Lorenz, tspan, u0) passes end-to-end on Enzyme v0.13.208 and v0.13.205, and errors on v0.13.209. The failing call is in _kernel_transfer on the fixed-step Float32 path — ext/EnzymeExt.jl and the test files are identical on master, so any branch resolving Enzyme 0.13.209 hits this. Upstream regression, not caused by this diff.
• Buildkite oneAPI v1 and v1.10 fail, same as master (gpu_kernel_de/conversions.jl:82). Metal v1 and v1.10 pass — the adaptive FMA path runs clean on real Metal hardware.
• The AMDGPU jobs never reported a status on build 1773 (no rocm agent picked them up); on the previous head they failed at ensemblegpuarray_scalar_batch.jl:102, the same Fix EnsembleGPUArray scalar parameter packing for batch_size=1 #555 regression fixed by the merged Fix #555 fallout: docs missing_docs and CUDA scalar-indexing test #563. Buildkite logs are not readable from the verification host, so AMDGPU on this head is unverified either way.
• All other checks pass: both benchmarks, Downgrade, CPU (1, 1/x86, lts, pre), JLArrays (1, lts, pre), OpenCL (1, lts, pre), Enzyme lts, QA, runic, typos.

Not verified / findings

• GPU throughput, and AMDGPU/oneAPI execution: the verification host has no GPU. The PTX counts above are compile-only.
• Allowed-to-fail pre jobs and downstream packages.
• Worth pushing back on: FMA changes adaptive rounding, so accepted step sequences and results can differ at the reltol level from master (as happened for the fixed-step path in Fuse fixed-step GPUTsit5 arithmetic and document controller accuracy #529).

Investigation of the earlier attempt: #529 and https://github.com/ChrisRackauckas/InternalJunk/pull/99.

🤖 Generated with Claude Code (model: claude-opus-5-5[1m])

https://claude.ai/code/session_01SJfy3ez7u4CqvAerKrT5np

Body updated for head 2dd2fd4 by Devin CLI 3000.11.3, model swe-2-high; local session transcript at /home/crackauc/sandbox/agent-jobs/DiffEqGPU.jl/jobs/545-devin/log.txt on amdci2.julia.csail.mit.edu

Risk assessment

• Risk: low
• Blast radius: Every user of the adaptive GPUTsit5 kernel stepper (GPUAT5I, which covers the PI and I controllers) on every backend (CUDA, AMDGPU, Metal, oneAPI, OpenCL, CPU/KA). Only the stage and error-estimate arithmetic now uses @muladd. Results and accepted step sequences can change at the reltol level. No public API or signature changes, @inline stays, and it doesn't use internals of any dependency (@muladd is already used on the fixed-step step! in the same file). The diff doesn't touch any test and doesn't loosen any tolerance.
• Evidence: Head baf4560 has 6 failing checks. All 6 fail the same way on the merge base 5e8faee, and the ones that ran on the latest master 46c4a76 fail the same way there too. None is a regression:
  • CUDA Tests (Julia 1): fails on head at test/ensemblegpuarray_scalar_batch.jl:102 ("Scalar indexing is disallowed", a test file added by Fix EnsembleGPUArray scalar parameter packing for batch_size=1 #555). Master 46c4a76 fails with the identical error (54 passed / 1 errored). Pre-existing.
  • Documentation: fails with :missing_docs for DiffEqGPU.ensemble_param and DiffEqGPU.pack_ordinary_parameters, the same failure as master 46c4a76 and 5e8faee. Pre-existing.
  • buildkite amdgpu v1 and v1.10: fail at the same ensemblegpuarray_scalar_batch.jl:102 as master build 1734 (5e8faee). They passed on d6be47f (before Fix EnsembleGPUArray scalar parameter packing for batch_size=1 #555), so Fix EnsembleGPUArray scalar parameter packing for batch_size=1 #555 introduced the failure, not this PR. Pre-existing.
  • buildkite oneapi v1 and v1.10: fail at gpu_kernel_de/conversions.jl:82 ("Initialization preprocessing runs once"), the same as master builds 1734 and 1732. Pre-existing.
  • buildkite aggregate: fails only because of the jobs above.
  • The Tsit5 accuracy testset runs before those failures and passed on head for CUDA (144/144), AMDGPU (144/144), oneAPI (84/84) and Metal (84/84). CPU, OpenCL, Enzyme, JLArrays, QA, Downgrade, Runic and typos all pass.
• Independent review: Devin CLI (Mac) / fusion-claude-opus-5-5-high-sidekick-swe-2-medium rated it low, high confidence: it is a 21-line @muladd wrap of stage arithmetic that keeps @inline and controller rounding unchanged, and every red check fails identically on master.
• Merge: needs human review. Findings:
  1. The PR is still a draft and says "Ignore this draft until reviewed by @ChrisRackauckas."
  2. It changes numerical results for all adaptive GPUTsit5 users with no demonstrated benefit. CPU timing shows no gain, the PR says it adds 0 FMA in Float64 and 12 in Float32 on the current CUDA stack, and no GPU benchmark was run.
  3. The description's list of pre-existing failures is stale. It names the CUDA test/enzyme/cuda_records.jl:9 failure and a downgrade-lane failure. On this head the downgrade lane passes and the CUDA failure is ensemblegpuarray_scalar_batch.jl:102. It doesn't mention the Documentation, AMDGPU or oneAPI failures.
  4. It adds no test for the FMA path. That is acceptable for a perf change covered by the existing accuracy tests, but worth noting.
  5. It adds an AGENTS.md policy note.

---

🤖 Risk assessment posted by an AI agent (fleet master) — harness: Devin CLI (Mac) / fusion-claude-opus-5-5-high-sidekick-swe-2-medium; dispatched by Claude Code head, model claude-opus-5-5[1m]
Conversation: local Claude Code session 3cd6500a-1f81-46b5-ac0b-c466e15b6a53 on Chris's Mac (session ID, no URL)

[ChrisRackauckas-Claude]

Taking this — Devin CLI 3000.11.3, model swe-2-high (local job 545-devin on amdci2.julia.csail.mit.edu).

[ChrisRackauckas]

We should just @muladd all kernel functions, and the full function rather than parts?