From b061c621d320f98f002990b1f5b4c61fb23d3d51 Mon Sep 17 00:00:00 2001
From: Ryan Senne <50930199+rsenne@users.noreply.github.com>
Date: Sun, 20 Sep 2026 16:56:25 -0400
Subject: [PATCH 1/4] Reuse Reactant staging buffers and copy results on the
 device

ReactantExt copied every compiled AutoReactant call through the host:
Array(x), a new XLA upload, Array(out), then a copy into a new device
array. With CuArray inputs on a CUDA XLA client that is two PCIe transfers
and three allocations per call.

Add two hooks next to needs_host_staging: _host_staging_buffer (CUDAExt
returns pinned memory) and _device_array_from_pointer (CUDAExt wraps a
"cuda" pointer as a non-owning CuArray). ReactantExt builds one callable
per compiled function that owns reusable host stages and copies results
on the device when the XLA buffer can be wrapped, falling back to the host
stage otherwise. Inputs are still uploaded fresh each call: Reactant has no
device-pointer import, and its copyto! into a ConcreteRArray runs a compiled
device copy on top of the upload, so persistent XLA inputs would be slower.

Results are still fresh arrays every call and keep eltype promotions.

Closes #69
---
 CHANGELOG.md                |   2 +
 docs/src/15-gpu.md          |   2 +-
 ext/CUDAExt.jl              |  17 ++++-
 ext/ReactantExt.jl          | 138 ++++++++++++++++++++++++++++++------
 src/ParallelMCMC.jl         |  25 +++++++
 test/test-CUDA-Extension.jl |  56 +++++++++++++++
 test/test-Reactant-HVP.jl   | 111 +++++++++++++++++++++++++++++
 7 files changed, 329 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 931068b..63d609b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -47,6 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   backend in `grad_logdensity` to derive one.
 - Batched derivative slots require `logdensity_batch`. An `hvp_batch` without
   a batched gradient now raises an error.
+- `AutoReactant` reuses pinned host buffers for device-array inputs and copies
+  results on the device when the XLA client is on the GPU (#69).
 
 ### Removed
 
diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md
index f1f678d..f3e91fa 100644
--- a/docs/src/15-gpu.md
+++ b/docs/src/15-gpu.md
@@ -251,7 +251,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant())
 !!! warning "Reactant constraints"
     **Captured arrays are frozen at compile time.** Mutating them after preparation does not change the compiled derivative. Pass mutable data as an argument.
 
-    **Reactant chooses the execution device independently of the input array.** With Reactant's CPU client, `CuArray` inputs round-trip through the host. The sampler warns about this when it prepares the model. Select a GPU client with `Reactant.set_default_backend` when available.
+    **Reactant chooses the execution device independently of the input array.** With a CUDA client, `CuArray` inputs are staged through pinned host memory and results are copied back on the device. With the CPU client, both directions go through the host and the sampler warns when it prepares the model. Select a GPU client with `Reactant.set_default_backend` when available.
 
 - The log-density must be Reactant-traceable. DynamicPPL-built log-densities are not.
 - Across an HVP's two AD passes, use `AutoReactant()` for both or neither. A hand-written gradient may pair with it. `SecondOrder` cannot contain `AutoReactant()`.
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
index ed0e4a7..d5ead9c 100644
--- a/ext/CUDAExt.jl
+++ b/ext/CUDAExt.jl
@@ -1,8 +1,23 @@
 module CUDAExt
 
 using ParallelMCMC: ParallelMCMC
-using CUDA: CuArray
+using CUDA: CUDA, CuArray, CuPtr
 
 ParallelMCMC.needs_host_staging(::CuArray) = true
 
+# Pinned host memory: the driver can DMA directly instead of staging the copy.
+function ParallelMCMC._host_staging_buffer(::CuArray, ::Type{T}, dims::Dims) where {T}
+    return CUDA.pin(Array{T}(undef, dims))
+end
+
+# Non-owning view of an XLA buffer. XLA and CUDA.jl share the device's primary
+# context, so the pointer is usable as a `CuPtr`. Callers copy out of the view
+# while XLA still holds the buffer.
+function ParallelMCMC._device_array_from_pointer(
+    ::CuArray, ::Type{T}, ptr::Ptr{Cvoid}, dims::Dims, platform::AbstractString
+) where {T}
+    platform == "cuda" || return nothing
+    return unsafe_wrap(CuArray, reinterpret(CuPtr{T}, UInt(ptr)), dims; own=false)
+end
+
 end
diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl
index 32a0858..5922de2 100644
--- a/ext/ReactantExt.jl
+++ b/ext/ReactantExt.jl
@@ -25,44 +25,142 @@ end
 function _warn_reactant_host_roundtrip(x::AbstractArray)
     if ParallelMCMC.needs_host_staging(x) && _reactant_client_platform() == "cpu"
         @warn "AutoReactant: preparing on a $(typeof(x)), but Reactant's default XLA " *
-            "client targets \"cpu\". Every call will round-trip to the host and back " *
-            "instead of running where the array lives. Select a GPU client with " *
-            "`Reactant.set_default_backend(\"gpu\")` if one is available." maxlog = 1
+            "client targets \"cpu\". Every call will copy to the host and back. " *
+            "Select a GPU client with `Reactant.set_default_backend(\"gpu\")` if one " *
+            "is available." maxlog = 1
     end
     return nothing
 end
 
-_host(x::Array) = x
-_host(x::AbstractArray) = Array(x)
+#=
+Inputs go through `to_rarray` every call. Reactant's `copyto!` into an existing
+ConcreteRArray uploads to a new buffer and then runs a compiled device copy on
+top, so keeping XLA input buffers around would cost more, not less. The host
+staging array is safe to reuse immediately: `ArrayFromHostBuffer` reads it
+inside a `GC.@preserve` that ends when the call returns.
+
+`to_rarray` wants an `Array`; views get collected first, as in
+`DEER._materialize_ad_array`.
+=#
+_upload(x::Array, ::Nothing) = Reactant.to_rarray(x)
+_upload(x::AbstractArray, ::Nothing) = Reactant.to_rarray(Array(x))
+function _upload(x::AbstractArray, stage::AbstractArray)
+    copyto!(stage, x)
+    return Reactant.to_rarray(stage)
+end
+
+# Host arrays upload directly and need no stage.
+function _in_stage(template::AbstractArray)
+    ParallelMCMC.needs_host_staging(template) || return nothing
+    return ParallelMCMC._host_staging_buffer(template, eltype(template), size(template))
+end
+
+# The output eltype can differ from the template's, so run the compiled
+# function once on the template to size the stage. The template is x0, a valid
+# point. The platform is fixed per compiled function.
+function _prepare_output(template::AbstractArray, out_thunk)
+    ParallelMCMC.needs_host_staging(template) || return nothing, ""
+    out = out_thunk()
+    platform = try
+        string(Reactant.XLA.platform_name(Reactant.XLA.client(out)))
+    catch
+        _reactant_client_platform()
+    end
+    out_stage = ParallelMCMC._host_staging_buffer(template, eltype(out), size(out))
+    return out_stage, platform
+end
 
-# Preserve promotions performed by the compiled function. The result is a fresh
-# array every call: callers keep gradients (tapes, workspaces), so it must not
-# alias a reused buffer.
-function _from_host(template::AbstractArray, out)
-    out_h = Array(out)
+# Keep promotions performed by the compiled function. Callers hold on to
+# gradients, so the result is a fresh array every call.
+function _promote_like(template::AbstractArray, out_h::Array)
     template isa Array && eltype(out_h) === eltype(template) && return out_h
     res = similar(template, eltype(out_h), size(out_h))
     copyto!(res, out_h)
     return res
 end
 
-#=
-Reactant's `copyto!(::ConcreteRArray, ::Array)` uploads to a new buffer and
-then runs a compiled device-to-device copy into the destination, which is
-strictly more work than the upload alone.
-=#
-_upload(x::AbstractArray) = Reactant.to_rarray(_host(x))
+# Pointer access needs an unsharded PJRT buffer that is not already on the
+# host. `nothing` means fall back to host staging.
+function _device_view(template::AbstractArray, out, platform::AbstractString)
+    out isa Reactant.ConcretePJRTArray || return nothing
+    Reactant.Sharding.is_sharded(out.sharding) && return nothing
+    buf = Reactant.get_buffer(out)
+    Reactant.XLA.buffer_on_cpu(buf) && return nothing
+    ptr = Reactant.XLA.unsafe_buffer_pointer(buf)
+    return ParallelMCMC._device_array_from_pointer(
+        template, eltype(out), ptr, size(out), platform
+    )
+end
+
+# No stage: `Array(out)` already allocates a fresh host array.
+function _download(template::AbstractArray, out, ::Nothing, ::AbstractString)
+    return _promote_like(template, Array(out))
+end
+
+function _download(
+    template::AbstractArray, out, out_stage::AbstractArray, platform::AbstractString
+)
+    wait(out)
+    view = _device_view(template, out, platform)
+    if view !== nothing
+        res = similar(template, eltype(out), size(out))
+        GC.@preserve out copyto!(res, view)
+        return res
+    end
+    # CPU client, IFRT, sharded, or a device family without pointer wrapping.
+    copyto!(out_stage, out)
+    res = similar(template, eltype(out), size(out))
+    copyto!(res, out_stage)
+    return res
+end
+
+# One of these per compiled function, owning its staging buffers. `template` is
+# the first argument; HVP outputs follow `x`, not `v`.
+struct ReactantUnary{F,S1,T,OS}
+    compiled::F
+    in1_stage::S1
+    template::T
+    out_stage::OS
+    platform::String
+end
+
+function (c::ReactantUnary)(x::AbstractArray)
+    xr = _upload(x, c.in1_stage)
+    return _download(c.template, c.compiled(xr), c.out_stage, c.platform)
+end
+
+struct ReactantBinary{F,S1,S2,T,OS}
+    compiled::F
+    in1_stage::S1
+    in2_stage::S2
+    template::T
+    out_stage::OS
+    platform::String
+end
+
+function (c::ReactantBinary)(x::AbstractArray, v::AbstractArray)
+    xr = _upload(x, c.in1_stage)
+    vr = _upload(v, c.in2_stage)
+    return _download(c.template, c.compiled(xr, vr), c.out_stage, c.platform)
+end
 
 function _compiled(core, t1::AbstractArray)
     _warn_reactant_host_roundtrip(t1)
-    compiled = @compile core(_upload(t1))
-    return x -> _from_host(x, compiled(_upload(x)))
+    in1_stage = _in_stage(t1)
+    compiled = @compile core(_upload(t1, in1_stage))
+    out_stage, platform = _prepare_output(t1, () -> compiled(_upload(t1, in1_stage)))
+    return ReactantUnary(compiled, in1_stage, t1, out_stage, platform)
 end
 
 function _compiled(core, t1::AbstractArray, t2::AbstractArray)
     _warn_reactant_host_roundtrip(t1)
-    compiled = @compile core(_upload(t1), _upload(t2))
-    return (x, v) -> _from_host(x, compiled(_upload(x), _upload(v)))
+    in1_stage = _in_stage(t1)
+    in2_stage = _in_stage(t2)
+    compiled = @compile core(_upload(t1, in1_stage), _upload(t2, in2_stage))
+    out_stage, platform = _prepare_output(
+        t1, () -> compiled(_upload(t1, in1_stage), _upload(t2, in2_stage))
+    )
+    return ReactantBinary(compiled, in1_stage, in2_stage, t1, out_stage, platform)
 end
 
 # For g = gradlogp, this JVP is the HVP. The callable and its captures are constant.
diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl
index 36b59fd..b387286 100644
--- a/src/ParallelMCMC.jl
+++ b/src/ParallelMCMC.jl
@@ -40,6 +40,31 @@ ParallelMCMC.needs_host_staging(::ROCArray) = true
 """
 needs_host_staging(::AbstractArray) = false
 
+"""
+    _host_staging_buffer(template::AbstractArray, ::Type{T}, dims::Dims) -> Array{T}
+
+Host buffer for staging transfers to and from arrays like `template`. Device
+extensions can return pinned memory. Defaults to a plain `Array`.
+"""
+function _host_staging_buffer(::AbstractArray, ::Type{T}, dims::Dims) where {T}
+    return Array{T}(undef, dims)
+end
+
+"""
+    _device_array_from_pointer(template::AbstractArray, ::Type{T}, ptr::Ptr{Cvoid}, dims::Dims, platform::AbstractString)
+
+Wrap device memory owned by another runtime as an array like `template`, or
+return `nothing` if that array type cannot address it. `platform` is the owner's
+XLA platform name (`"cuda"`, `"rocm"`, `"cpu"`). The result aliases `ptr` and
+does not own it; copy out of it while the owner still holds the memory.
+Defaults to `nothing`.
+"""
+function _device_array_from_pointer(
+    ::AbstractArray, ::Type, ::Ptr{Cvoid}, ::Dims, ::AbstractString
+)
+    return nothing
+end
+
 #= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks
 and `interface.jl`'s gradient hooks report it. =#
 const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt."
diff --git a/test/test-CUDA-Extension.jl b/test/test-CUDA-Extension.jl
index 4a58896..8a270d5 100644
--- a/test/test-CUDA-Extension.jl
+++ b/test/test-CUDA-Extension.jl
@@ -85,6 +85,23 @@ end
     end
 end
 
+@testset "Staging hooks default to plain host arrays" begin
+    buf = ParallelMCMC._host_staging_buffer(zeros(3), Float32, (2, 2))
+    @test buf isa Matrix{Float32}
+    @test size(buf) == (2, 2)
+
+    @test ParallelMCMC._device_array_from_pointer(
+        zeros(3), Float64, C_NULL, (3,), "cpu"
+    ) === nothing
+
+    x = StagedArray(zeros(Float32, 6))
+    buf2 = ParallelMCMC._host_staging_buffer(x, Float64, (3, 2))
+    @test buf2 isa Matrix{Float64}
+    @test size(buf2) == (3, 2)
+    @test ParallelMCMC._device_array_from_pointer(x, Float32, C_NULL, (6,), "cpu") ===
+        nothing
+end
+
 @testset "CUDAExt loads with CUDA" begin
     cuda_loadable = try
         using CUDA
@@ -102,3 +119,42 @@ end
         end
     end
 end
+
+@testset "CUDAExt device buffer hooks" begin
+    cuda_functional = try
+        using CUDA
+        CUDA.functional() && (CUDA.CuArray([1.0f0]); true)
+    catch
+        false
+    end
+    if !cuda_functional
+        @info "CUDAExt device buffer hooks: CUDA not functional, skipping"
+    else
+        template = CUDA.zeros(Float32, 3)
+
+        host_buf = ParallelMCMC._host_staging_buffer(template, Float32, (2, 3))
+        @test host_buf isa Array{Float32}
+        @test size(host_buf) == (2, 3)
+        @test CUDA.is_pinned(pointer(host_buf))
+
+        # `CuPtr` does not convert to `Ptr`; go through `UInt` like the hook does.
+        src = CUDA.CuArray(Float32[1, 2, 3, 4])
+        raw_ptr = Ptr{Cvoid}(UInt(pointer(src)))
+        dev = ParallelMCMC._device_array_from_pointer(
+            template, Float32, raw_ptr, (4,), "cuda"
+        )
+        @test dev isa CUDA.CuArray{Float32}
+        @test Array(dev) == Array(src)
+
+        # The view aliases the source.
+        dev .= 0.0f0
+        @test all(iszero, Array(src))
+
+        @test ParallelMCMC._device_array_from_pointer(
+            template, Float32, raw_ptr, (4,), "rocm"
+        ) === nothing
+        @test ParallelMCMC._device_array_from_pointer(
+            template, Float32, raw_ptr, (4,), "cpu"
+        ) === nothing
+    end
+end
diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl
index c71bdcc..953874e 100644
--- a/test/test-Reactant-HVP.jl
+++ b/test/test-Reactant-HVP.jl
@@ -14,6 +14,7 @@ hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x)
 logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2)
 gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1)
 logp_r32(x) = -0.25f0 * sum(abs2, x)^2
+logp_batch_r32(X) = vec(-0.25f0 .* sum(abs2, X; dims=1) .^ 2)
 
 logp_gauss(x) = -0.5 * sum(abs2, x)
 gradlogp_gauss(x) = -x
@@ -27,6 +28,28 @@ const CT_R = FlexiChains.FlexiChain{Symbol}
 using Reactant: Reactant
 using Enzyme: Enzyme
 
+# Device-array stand-in for the host-staging path without CUDA. Like
+# `StagedArray` in test-CUDA-Extension.jl, but it also needs a broadcast style
+# because `_prepare_model` fills the batched template with `X_template .= x_template`.
+struct ReactantStagedArray{T,N} <: AbstractArray{T,N}
+    data::Array{T,N}
+end
+Base.size(a::ReactantStagedArray) = size(a.data)
+Base.getindex(a::ReactantStagedArray, i::Int...) = a.data[i...]
+function Base.setindex!(::ReactantStagedArray, v, i::Int...)
+    return error("scalar indexing is unsupported")
+end
+function Base.similar(a::ReactantStagedArray, ::Type{T}, dims::Dims) where {T}
+    return ReactantStagedArray(similar(a.data, T, dims))
+end
+Base.copyto!(a::ReactantStagedArray, src::AbstractArray) = (copyto!(a.data, src); a)
+function Base.BroadcastStyle(::Type{<:ReactantStagedArray})
+    return Broadcast.ArrayStyle{ReactantStagedArray}()
+end
+Base.copyto!(a::ReactantStagedArray, bc::Broadcast.Broadcasted) = (copyto!(a.data, bc); a)
+
+ParallelMCMC.needs_host_staging(::ReactantStagedArray) = true
+
 @testset "Reactant HVP" begin
     @testset "extension is loaded" begin
         @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing
@@ -133,6 +156,18 @@ using Enzyme: Enzyme
             end
         end
 
+        # DEER passes column views to the vector callables.
+        @testset "SubArray inputs" begin
+            model = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoReactant())
+            m_p = ParallelMCMC._prepare_model(model, x, 8, nothing)
+            S = randn(rng, D_R, 2)
+            xv = view(S, :, 1)
+            vv = view(S, :, 2)
+            @test m_p.grad_logdensity(xv) isa Vector{Float64}
+            @test m_p.grad_logdensity(xv) ≈ gradlogp_r(S[:, 1])
+            @test m_p.hvp(xv, vv) ≈ hvp_r(S[:, 1], S[:, 2])
+        end
+
         @testset "Float32 on CPU" begin
             x32 = Float32.(x)
             v32 = Float32.(v)
@@ -198,6 +233,66 @@ using Enzyme: Enzyme
         end
     end
 
+    # `ReactantStagedArray` has no `_device_array_from_pointer` method, so results
+    # come back through the host stage, as they would for a CuArray on the CPU client.
+    @testset "staged template" begin
+        rng = MersenneTwister(74)
+        x0 = ReactantStagedArray(randn(rng, Float32, D_R))
+        model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant())
+        m_p = ParallelMCMC._prepare_model(model, x0, 8, nothing)
+
+        x1 = ReactantStagedArray(randn(rng, Float32, D_R))
+        v1 = ReactantStagedArray(randn(rng, Float32, D_R))
+        x2 = ReactantStagedArray(randn(rng, Float32, D_R))
+        v2 = ReactantStagedArray(randn(rng, Float32, D_R))
+
+        g1 = m_p.grad_logdensity(x1)
+        Hv1 = m_p.hvp(x1, v1)
+        @test g1 isa ReactantStagedArray{Float32,1}
+        @test eltype(g1) === Float32
+        @test g1.data ≈ gradlogp_r(x1.data)
+        @test Hv1 isa ReactantStagedArray{Float32,1}
+        @test Hv1.data ≈ hvp_r(x1.data, v1.data)
+
+        g2 = m_p.grad_logdensity(x2)
+        Hv2 = m_p.hvp(x2, v2)
+        @test g2.data ≈ gradlogp_r(x2.data)
+        @test Hv2.data ≈ hvp_r(x2.data, v2.data)
+
+        # Results must not alias the reused stage.
+        @test g1 !== g2
+        @test g1.data !== g2.data
+        @test Hv1 !== Hv2
+        @test Hv1.data !== Hv2.data
+
+        @testset "batched slots" begin
+            T = 8
+            model_b = DensityModel(
+                logp_r32,
+                AutoReactant(),
+                D_R;
+                logdensity_batch=logp_batch_r32,
+                grad_logdensity_batch=AutoReactant(),
+                hvp=AutoReactant(),
+                hvp_batch=AutoReactant(),
+            )
+            m_pb = ParallelMCMC._prepare_model(model_b, x0, T, nothing)
+
+            X1 = ReactantStagedArray(randn(rng, Float32, D_R, T))
+            V1 = ReactantStagedArray(randn(rng, Float32, D_R, T))
+
+            Gb = m_pb.grad_logdensity_batch(X1)
+            @test Gb isa ReactantStagedArray{Float32,2}
+            @test eltype(Gb) === Float32
+            @test Gb.data ≈ gradlogp_batch_r(X1.data)
+
+            Hvb = m_pb.hvp_batch(X1, V1)
+            Hv_cols = reduce(hcat, [hvp_r(X1.data[:, t], V1.data[:, t]) for t in 1:T])
+            @test Hvb isa ReactantStagedArray{Float32,2}
+            @test Hvb.data ≈ Hv_cols
+        end
+    end
+
     reactant_gpu_ok = try
         using CUDA: CUDA
         CUDA.functional() && (CUDA.CuArray([1.0f0]); true)
@@ -226,6 +321,22 @@ using Enzyme: Enzyme
             Hv = m_p.hvp(x_d, v_d)
             @test Hv isa CUDA.CuArray
             @test Array(Hv) ≈ hvp_r(x_h, v_h)
+
+            # Repeated calls must not alias the reused stages.
+            x_h2 = randn(rng, Float32, D_R)
+            v_h2 = randn(rng, Float32, D_R)
+            x_d2 = CUDA.CuArray(x_h2)
+            v_d2 = CUDA.CuArray(v_h2)
+
+            g2 = m_p.grad_logdensity(x_d2)
+            Hv2 = m_p.hvp(x_d2, v_d2)
+            @test g2 isa CUDA.CuArray
+            @test Array(g2) ≈ gradlogp_r(x_h2)
+            @test Hv2 isa CUDA.CuArray
+            @test Array(Hv2) ≈ hvp_r(x_h2, v_h2)
+
+            @test g !== g2
+            @test Hv !== Hv2
         end
     end
 end

From 16ed9f2a91f4fe6227f7f85ebe403f291326728f Mon Sep 17 00:00:00 2001
From: Ryan Senne <50930199+rsenne@users.noreply.github.com>
Date: Mon, 21 Sep 2026 08:51:51 -0400
Subject: [PATCH 2/4] Address review of the Reactant staging buffers

Copy device results through a copy-and-synchronize hook instead of handing back
a non-owning view. The view was read by an async copyto! on CUDA.jl's stream, so
GC.@preserve held the XLA buffer only until the copy was enqueued, not until it
ran. The extension now owns the copy and synchronizes before returning.

Drop the output staging buffer and the extra compiled call that sized it. The
platform is known at compile time from the default client, and the fallback is a
plain host download. Collapse the two arity-specific callables into one
ReactantCall, and drop a try/catch that could not trigger.

Share one StagedArray test double between the CUDA and Reactant test files.
---
 CHANGELOG.md                |   2 +-
 ext/CUDAExt.jl              |  17 +++---
 ext/ReactantExt.jl          | 103 +++++++++---------------------------
 src/ParallelMCMC.jl         |  18 +++----
 test/staged_array.jl        |  29 ++++++++++
 test/test-CUDA-Extension.jl |  48 +++++------------
 test/test-Reactant-HVP.jl   |  48 +++++------------
 7 files changed, 99 insertions(+), 166 deletions(-)
 create mode 100644 test/staged_array.jl

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63d609b..3fab311 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -48,7 +48,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Batched derivative slots require `logdensity_batch`. An `hvp_batch` without
   a batched gradient now raises an error.
 - `AutoReactant` reuses pinned host buffers for device-array inputs and copies
-  results on the device when the XLA client is on the GPU (#69).
+  results on the device when the XLA client is on the GPU (#73).
 
 ### Removed
 
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
index d5ead9c..7f89fa4 100644
--- a/ext/CUDAExt.jl
+++ b/ext/CUDAExt.jl
@@ -10,14 +10,17 @@ function ParallelMCMC._host_staging_buffer(::CuArray, ::Type{T}, dims::Dims) whe
     return CUDA.pin(Array{T}(undef, dims))
 end
 
-# Non-owning view of an XLA buffer. XLA and CUDA.jl share the device's primary
-# context, so the pointer is usable as a `CuPtr`. Callers copy out of the view
-# while XLA still holds the buffer.
-function ParallelMCMC._device_array_from_pointer(
-    ::CuArray, ::Type{T}, ptr::Ptr{Cvoid}, dims::Dims, platform::AbstractString
+# XLA and CUDA.jl share the device's primary context, so the pointer is usable
+# as a `CuPtr`. The copy runs on CUDA.jl's stream, which XLA does not track, so
+# synchronize before returning: the caller may release the source after this.
+function ParallelMCMC._copy_from_device_pointer!(
+    dest::CuArray{T}, ptr::Ptr{Cvoid}, platform::AbstractString
 ) where {T}
-    platform == "cuda" || return nothing
-    return unsafe_wrap(CuArray, reinterpret(CuPtr{T}, UInt(ptr)), dims; own=false)
+    platform == "cuda" || return false
+    src = unsafe_wrap(CuArray, reinterpret(CuPtr{T}, UInt(ptr)), size(dest); own=false)
+    copyto!(dest, src)
+    CUDA.synchronize()
+    return true
 end
 
 end
diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl
index 5922de2..12ff716 100644
--- a/ext/ReactantExt.jl
+++ b/ext/ReactantExt.jl
@@ -55,21 +55,6 @@ function _in_stage(template::AbstractArray)
     return ParallelMCMC._host_staging_buffer(template, eltype(template), size(template))
 end
 
-# The output eltype can differ from the template's, so run the compiled
-# function once on the template to size the stage. The template is x0, a valid
-# point. The platform is fixed per compiled function.
-function _prepare_output(template::AbstractArray, out_thunk)
-    ParallelMCMC.needs_host_staging(template) || return nothing, ""
-    out = out_thunk()
-    platform = try
-        string(Reactant.XLA.platform_name(Reactant.XLA.client(out)))
-    catch
-        _reactant_client_platform()
-    end
-    out_stage = ParallelMCMC._host_staging_buffer(template, eltype(out), size(out))
-    return out_stage, platform
-end
-
 # Keep promotions performed by the compiled function. Callers hold on to
 # gradients, so the result is a fresh array every call.
 function _promote_like(template::AbstractArray, out_h::Array)
@@ -79,88 +64,48 @@ function _promote_like(template::AbstractArray, out_h::Array)
     return res
 end
 
-# Pointer access needs an unsharded PJRT buffer that is not already on the
-# host. `nothing` means fall back to host staging.
-function _device_view(template::AbstractArray, out, platform::AbstractString)
+# Raw pointer to an XLA buffer, or `nothing` if it cannot be taken: pointer
+# access needs an unsharded PJRT buffer that is not already on the host.
+function _device_pointer(out)
     out isa Reactant.ConcretePJRTArray || return nothing
     Reactant.Sharding.is_sharded(out.sharding) && return nothing
+    wait(out)
     buf = Reactant.get_buffer(out)
     Reactant.XLA.buffer_on_cpu(buf) && return nothing
-    ptr = Reactant.XLA.unsafe_buffer_pointer(buf)
-    return ParallelMCMC._device_array_from_pointer(
-        template, eltype(out), ptr, size(out), platform
-    )
+    return Reactant.XLA.unsafe_buffer_pointer(buf)
 end
 
-# No stage: `Array(out)` already allocates a fresh host array.
-function _download(template::AbstractArray, out, ::Nothing, ::AbstractString)
-    return _promote_like(template, Array(out))
-end
-
-function _download(
-    template::AbstractArray, out, out_stage::AbstractArray, platform::AbstractString
-)
-    wait(out)
-    view = _device_view(template, out, platform)
-    if view !== nothing
+function _download(template::AbstractArray, out, platform::AbstractString)
+    ptr = ParallelMCMC.needs_host_staging(template) ? _device_pointer(out) : nothing
+    if ptr !== nothing
         res = similar(template, eltype(out), size(out))
-        GC.@preserve out copyto!(res, view)
-        return res
+        GC.@preserve out begin
+            ParallelMCMC._copy_from_device_pointer!(res, ptr, platform) && return res
+        end
     end
-    # CPU client, IFRT, sharded, or a device family without pointer wrapping.
-    copyto!(out_stage, out)
-    res = similar(template, eltype(out), size(out))
-    copyto!(res, out_stage)
-    return res
+    return _promote_like(template, Array(out))
 end
 
 # One of these per compiled function, owning its staging buffers. `template` is
-# the first argument; HVP outputs follow `x`, not `v`.
-struct ReactantUnary{F,S1,T,OS}
+# the first argument; HVP outputs follow `x`, not `v`. `stages` has one entry
+# per argument, `nothing` or a host buffer. `platform` is fixed at compile time
+# from the default XLA client.
+struct ReactantCall{F,S,T}
     compiled::F
-    in1_stage::S1
+    stages::S
     template::T
-    out_stage::OS
     platform::String
 end
 
-function (c::ReactantUnary)(x::AbstractArray)
-    xr = _upload(x, c.in1_stage)
-    return _download(c.template, c.compiled(xr), c.out_stage, c.platform)
-end
-
-struct ReactantBinary{F,S1,S2,T,OS}
-    compiled::F
-    in1_stage::S1
-    in2_stage::S2
-    template::T
-    out_stage::OS
-    platform::String
+function (c::ReactantCall)(args::AbstractArray...)
+    return _download(c.template, c.compiled(map(_upload, args, c.stages)...), c.platform)
 end
 
-function (c::ReactantBinary)(x::AbstractArray, v::AbstractArray)
-    xr = _upload(x, c.in1_stage)
-    vr = _upload(v, c.in2_stage)
-    return _download(c.template, c.compiled(xr, vr), c.out_stage, c.platform)
-end
-
-function _compiled(core, t1::AbstractArray)
-    _warn_reactant_host_roundtrip(t1)
-    in1_stage = _in_stage(t1)
-    compiled = @compile core(_upload(t1, in1_stage))
-    out_stage, platform = _prepare_output(t1, () -> compiled(_upload(t1, in1_stage)))
-    return ReactantUnary(compiled, in1_stage, t1, out_stage, platform)
-end
-
-function _compiled(core, t1::AbstractArray, t2::AbstractArray)
-    _warn_reactant_host_roundtrip(t1)
-    in1_stage = _in_stage(t1)
-    in2_stage = _in_stage(t2)
-    compiled = @compile core(_upload(t1, in1_stage), _upload(t2, in2_stage))
-    out_stage, platform = _prepare_output(
-        t1, () -> compiled(_upload(t1, in1_stage), _upload(t2, in2_stage))
-    )
-    return ReactantBinary(compiled, in1_stage, in2_stage, t1, out_stage, platform)
+function _compiled(core, templates::AbstractArray...)
+    _warn_reactant_host_roundtrip(templates[1])
+    stages = map(_in_stage, templates)
+    compiled = @compile core(map(_upload, templates, stages)...)
+    return ReactantCall(compiled, stages, templates[1], _reactant_client_platform())
 end
 
 # For g = gradlogp, this JVP is the HVP. The callable and its captures are constant.
diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl
index b387286..a2f279d 100644
--- a/src/ParallelMCMC.jl
+++ b/src/ParallelMCMC.jl
@@ -51,19 +51,15 @@ function _host_staging_buffer(::AbstractArray, ::Type{T}, dims::Dims) where {T}
 end
 
 """
-    _device_array_from_pointer(template::AbstractArray, ::Type{T}, ptr::Ptr{Cvoid}, dims::Dims, platform::AbstractString)
+    _copy_from_device_pointer!(dest::AbstractArray, ptr::Ptr{Cvoid}, platform::AbstractString) -> Bool
 
-Wrap device memory owned by another runtime as an array like `template`, or
-return `nothing` if that array type cannot address it. `platform` is the owner's
-XLA platform name (`"cuda"`, `"rocm"`, `"cpu"`). The result aliases `ptr` and
-does not own it; copy out of it while the owner still holds the memory.
-Defaults to `nothing`.
+Copy device memory owned by another runtime into `dest`, returning `true` on
+success and `false` if `dest`'s array type cannot address `ptr`. `platform` is
+the owner's XLA platform name ("cuda", "rocm", "cpu"). The copy is complete
+when this returns, so the caller may release the source afterwards. Defaults
+to `false`.
 """
-function _device_array_from_pointer(
-    ::AbstractArray, ::Type, ::Ptr{Cvoid}, ::Dims, ::AbstractString
-)
-    return nothing
-end
+_copy_from_device_pointer!(::AbstractArray, ::Ptr{Cvoid}, ::AbstractString) = false
 
 #= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks
 and `interface.jl`'s gradient hooks report it. =#
diff --git a/test/staged_array.jl b/test/staged_array.jl
new file mode 100644
index 0000000..4fb252b
--- /dev/null
+++ b/test/staged_array.jl
@@ -0,0 +1,29 @@
+#= Device-array stand-in for the host-staging paths, so they can be tested
+without a GPU. Scalar `setindex!` errors, as on `CuArray`. Shared by
+test-CUDA-Extension.jl and test-Reactant-HVP.jl; the name does not match
+`test-*.jl`, so runtests.jl does not include it as a testset. =#
+
+struct StagedArray{T,N} <: AbstractArray{T,N}
+    data::Array{T,N}
+end
+
+Base.size(a::StagedArray) = size(a.data)
+Base.getindex(a::StagedArray, i::Int...) = a.data[i...]
+Base.setindex!(::StagedArray, v, i::Int...) = error("scalar indexing is unsupported")
+function Base.similar(a::StagedArray, ::Type{T}, dims::Dims) where {T}
+    return StagedArray(similar(a.data, T, dims))
+end
+Base.copyto!(a::StagedArray, src::AbstractArray) = (copyto!(a.data, src); a)
+
+# `_prepare_model` fills the batched template with `X_template .= x_template`.
+# The zero-dim method disambiguates against Base's, so scalar fills also work.
+Base.BroadcastStyle(::Type{<:StagedArray}) = Broadcast.ArrayStyle{StagedArray}()
+Base.copyto!(a::StagedArray, bc::Broadcast.Broadcasted) = (copyto!(a.data, bc); a)
+function Base.copyto!(
+    a::StagedArray, bc::Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{0}}
+)
+    copyto!(a.data, bc)
+    return a
+end
+
+ParallelMCMC.needs_host_staging(::StagedArray) = true
diff --git a/test/test-CUDA-Extension.jl b/test/test-CUDA-Extension.jl
index 8a270d5..5888db6 100644
--- a/test/test-CUDA-Extension.jl
+++ b/test/test-CUDA-Extension.jl
@@ -14,19 +14,7 @@ end
     @test ws.zhost === nothing
 end
 
-# Device-array stand-in: scalar `setindex!` errors, as on `CuArray`.
-struct StagedArray{T,N} <: AbstractArray{T,N}
-    data::Array{T,N}
-end
-Base.size(a::StagedArray) = size(a.data)
-Base.getindex(a::StagedArray, i::Int...) = a.data[i...]
-Base.setindex!(::StagedArray, v, i::Int...) = error("scalar indexing is unsupported")
-function Base.similar(a::StagedArray, ::Type{T}, dims::Dims) where {T}
-    return StagedArray(similar(a.data, T, dims))
-end
-Base.copyto!(a::StagedArray, src::AbstractArray) = (copyto!(a.data, src); a)
-
-ParallelMCMC.needs_host_staging(::StagedArray) = true
+isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array.jl"))
 
 @testset "Opting a device array type in" begin
     x = StagedArray(zeros(Float32, 6))
@@ -90,16 +78,13 @@ end
     @test buf isa Matrix{Float32}
     @test size(buf) == (2, 2)
 
-    @test ParallelMCMC._device_array_from_pointer(
-        zeros(3), Float64, C_NULL, (3,), "cpu"
-    ) === nothing
+    @test ParallelMCMC._copy_from_device_pointer!(zeros(3), C_NULL, "cpu") == false
 
     x = StagedArray(zeros(Float32, 6))
     buf2 = ParallelMCMC._host_staging_buffer(x, Float64, (3, 2))
     @test buf2 isa Matrix{Float64}
     @test size(buf2) == (3, 2)
-    @test ParallelMCMC._device_array_from_pointer(x, Float32, C_NULL, (6,), "cpu") ===
-        nothing
+    @test ParallelMCMC._copy_from_device_pointer!(x, C_NULL, "cpu") == false
 end
 
 @testset "CUDAExt loads with CUDA" begin
@@ -140,21 +125,16 @@ end
         # `CuPtr` does not convert to `Ptr`; go through `UInt` like the hook does.
         src = CUDA.CuArray(Float32[1, 2, 3, 4])
         raw_ptr = Ptr{Cvoid}(UInt(pointer(src)))
-        dev = ParallelMCMC._device_array_from_pointer(
-            template, Float32, raw_ptr, (4,), "cuda"
-        )
-        @test dev isa CUDA.CuArray{Float32}
-        @test Array(dev) == Array(src)
-
-        # The view aliases the source.
-        dev .= 0.0f0
-        @test all(iszero, Array(src))
-
-        @test ParallelMCMC._device_array_from_pointer(
-            template, Float32, raw_ptr, (4,), "rocm"
-        ) === nothing
-        @test ParallelMCMC._device_array_from_pointer(
-            template, Float32, raw_ptr, (4,), "cpu"
-        ) === nothing
+        dest = CUDA.CuArray{Float32}(undef, 4)
+        @test ParallelMCMC._copy_from_device_pointer!(dest, raw_ptr, "cuda") == true
+        @test Array(dest) == Array(src)
+
+        # The copy landed: dest does not alias src.
+        dest .= 0.0f0
+        @test Array(src) == Float32[1, 2, 3, 4]
+
+        dest2 = CUDA.CuArray{Float32}(undef, 4)
+        @test ParallelMCMC._copy_from_device_pointer!(dest2, raw_ptr, "rocm") == false
+        @test ParallelMCMC._copy_from_device_pointer!(dest2, raw_ptr, "cpu") == false
     end
 end
diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl
index 953874e..5bc9977 100644
--- a/test/test-Reactant-HVP.jl
+++ b/test/test-Reactant-HVP.jl
@@ -28,27 +28,7 @@ const CT_R = FlexiChains.FlexiChain{Symbol}
 using Reactant: Reactant
 using Enzyme: Enzyme
 
-# Device-array stand-in for the host-staging path without CUDA. Like
-# `StagedArray` in test-CUDA-Extension.jl, but it also needs a broadcast style
-# because `_prepare_model` fills the batched template with `X_template .= x_template`.
-struct ReactantStagedArray{T,N} <: AbstractArray{T,N}
-    data::Array{T,N}
-end
-Base.size(a::ReactantStagedArray) = size(a.data)
-Base.getindex(a::ReactantStagedArray, i::Int...) = a.data[i...]
-function Base.setindex!(::ReactantStagedArray, v, i::Int...)
-    return error("scalar indexing is unsupported")
-end
-function Base.similar(a::ReactantStagedArray, ::Type{T}, dims::Dims) where {T}
-    return ReactantStagedArray(similar(a.data, T, dims))
-end
-Base.copyto!(a::ReactantStagedArray, src::AbstractArray) = (copyto!(a.data, src); a)
-function Base.BroadcastStyle(::Type{<:ReactantStagedArray})
-    return Broadcast.ArrayStyle{ReactantStagedArray}()
-end
-Base.copyto!(a::ReactantStagedArray, bc::Broadcast.Broadcasted) = (copyto!(a.data, bc); a)
-
-ParallelMCMC.needs_host_staging(::ReactantStagedArray) = true
+isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array.jl"))
 
 @testset "Reactant HVP" begin
     @testset "extension is loaded" begin
@@ -233,25 +213,25 @@ ParallelMCMC.needs_host_staging(::ReactantStagedArray) = true
         end
     end
 
-    # `ReactantStagedArray` has no `_device_array_from_pointer` method, so results
-    # come back through the host stage, as they would for a CuArray on the CPU client.
+    # `StagedArray` has no `_copy_from_device_pointer!` method, so results
+    # come back through a plain host download, as they would for a CuArray on the CPU client.
     @testset "staged template" begin
         rng = MersenneTwister(74)
-        x0 = ReactantStagedArray(randn(rng, Float32, D_R))
+        x0 = StagedArray(randn(rng, Float32, D_R))
         model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant())
         m_p = ParallelMCMC._prepare_model(model, x0, 8, nothing)
 
-        x1 = ReactantStagedArray(randn(rng, Float32, D_R))
-        v1 = ReactantStagedArray(randn(rng, Float32, D_R))
-        x2 = ReactantStagedArray(randn(rng, Float32, D_R))
-        v2 = ReactantStagedArray(randn(rng, Float32, D_R))
+        x1 = StagedArray(randn(rng, Float32, D_R))
+        v1 = StagedArray(randn(rng, Float32, D_R))
+        x2 = StagedArray(randn(rng, Float32, D_R))
+        v2 = StagedArray(randn(rng, Float32, D_R))
 
         g1 = m_p.grad_logdensity(x1)
         Hv1 = m_p.hvp(x1, v1)
-        @test g1 isa ReactantStagedArray{Float32,1}
+        @test g1 isa StagedArray{Float32,1}
         @test eltype(g1) === Float32
         @test g1.data ≈ gradlogp_r(x1.data)
-        @test Hv1 isa ReactantStagedArray{Float32,1}
+        @test Hv1 isa StagedArray{Float32,1}
         @test Hv1.data ≈ hvp_r(x1.data, v1.data)
 
         g2 = m_p.grad_logdensity(x2)
@@ -278,17 +258,17 @@ ParallelMCMC.needs_host_staging(::ReactantStagedArray) = true
             )
             m_pb = ParallelMCMC._prepare_model(model_b, x0, T, nothing)
 
-            X1 = ReactantStagedArray(randn(rng, Float32, D_R, T))
-            V1 = ReactantStagedArray(randn(rng, Float32, D_R, T))
+            X1 = StagedArray(randn(rng, Float32, D_R, T))
+            V1 = StagedArray(randn(rng, Float32, D_R, T))
 
             Gb = m_pb.grad_logdensity_batch(X1)
-            @test Gb isa ReactantStagedArray{Float32,2}
+            @test Gb isa StagedArray{Float32,2}
             @test eltype(Gb) === Float32
             @test Gb.data ≈ gradlogp_batch_r(X1.data)
 
             Hvb = m_pb.hvp_batch(X1, V1)
             Hv_cols = reduce(hcat, [hvp_r(X1.data[:, t], V1.data[:, t]) for t in 1:T])
-            @test Hvb isa ReactantStagedArray{Float32,2}
+            @test Hvb isa StagedArray{Float32,2}
             @test Hvb.data ≈ Hv_cols
         end
     end

From 1d4870f2f59afcdfdbcc2b351467b5b878a75fb4 Mon Sep 17 00:00:00 2001
From: Ryan Senne <50930199+rsenne@users.noreply.github.com>
Date: Thu, 24 Sep 2026 11:44:22 -0400
Subject: [PATCH 3/4] add upload

---
 .github/workflows/TestReactant.yml |  1 +
 CHANGELOG.md                       |  4 ++--
 ext/CUDAExt.jl                     | 14 +++++++-----
 ext/ReactantExt.jl                 | 36 +++++++++++++++++-------------
 src/ParallelMCMC.jl                | 13 +++++------
 test/staged_array.jl               | 12 +++++-----
 test/test-CUDA-Extension.jl        |  6 +++--
 test/test-Reactant-HVP.jl          | 14 ++++++++++--
 8 files changed, 61 insertions(+), 39 deletions(-)

diff --git a/.github/workflows/TestReactant.yml b/.github/workflows/TestReactant.yml
index b1e1595..241a6ec 100644
--- a/.github/workflows/TestReactant.yml
+++ b/.github/workflows/TestReactant.yml
@@ -8,6 +8,7 @@ on:
       - "src/**"
       - "ext/ReactantExt.jl"
       - "test/test-Reactant-*.jl"
+      - "test/staged_array.jl"
       - "Project.toml"
       - "test/Project.toml"
       - ".github/workflows/TestReactant.yml"
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3fab311..fc28fcc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -47,8 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   backend in `grad_logdensity` to derive one.
 - Batched derivative slots require `logdensity_batch`. An `hvp_batch` without
   a batched gradient now raises an error.
-- `AutoReactant` reuses pinned host buffers for device-array inputs and copies
-  results on the device when the XLA client is on the GPU (#73).
+- `AutoReactant` stages device-array inputs through reused pinned host buffers,
+  and copies results device-to-device when the XLA client is on the GPU (#73).
 
 ### Removed
 
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
index 7f89fa4..8138bb6 100644
--- a/ext/CUDAExt.jl
+++ b/ext/CUDAExt.jl
@@ -5,14 +5,18 @@ using CUDA: CUDA, CuArray, CuPtr
 
 ParallelMCMC.needs_host_staging(::CuArray) = true
 
-# Pinned host memory: the driver can DMA directly instead of staging the copy.
+#= Pinned, so the device-to-host copy into this buffer can DMA directly.
+`CUDA.pin` returns `nothing` for an already-registered address, so don't
+forward its result. =#
 function ParallelMCMC._host_staging_buffer(::CuArray, ::Type{T}, dims::Dims) where {T}
-    return CUDA.pin(Array{T}(undef, dims))
+    buf = Array{T}(undef, dims)
+    CUDA.pin(buf)
+    return buf
 end
 
-# XLA and CUDA.jl share the device's primary context, so the pointer is usable
-# as a `CuPtr`. The copy runs on CUDA.jl's stream, which XLA does not track, so
-# synchronize before returning: the caller may release the source after this.
+#= XLA and CUDA.jl share the device's primary context, so the pointer is usable
+as a `CuPtr`. The copy runs on CUDA.jl's stream, which XLA does not track;
+synchronize before returning, since the caller may release the source. =#
 function ParallelMCMC._copy_from_device_pointer!(
     dest::CuArray{T}, ptr::Ptr{Cvoid}, platform::AbstractString
 ) where {T}
diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl
index 12ff716..d1c7a42 100644
--- a/ext/ReactantExt.jl
+++ b/ext/ReactantExt.jl
@@ -33,18 +33,25 @@ function _warn_reactant_host_roundtrip(x::AbstractArray)
 end
 
 #=
-Inputs go through `to_rarray` every call. Reactant's `copyto!` into an existing
-ConcreteRArray uploads to a new buffer and then runs a compiled device copy on
-top, so keeping XLA input buffers around would cost more, not less. The host
-staging array is safe to reuse immediately: `ArrayFromHostBuffer` reads it
-inside a `GC.@preserve` that ends when the call returns.
-
-`to_rarray` wants an `Array`; views get collected first, as in
-`DEER._materialize_ad_array`.
+Every call uploads through `to_rarray`. Caching the XLA buffer instead would be
+slower: Reactant's `copyto!` into an existing ConcreteRArray uploads to a new
+buffer and then runs a compiled device copy on top of it.
+
+XLA takes the host stage with `kImmutableOnlyDuringCall` semantics, so the stage
+is free to overwrite once `to_rarray` returns. `to_rarray` wants an `Array`, so
+views get collected first, as in `DEER._materialize_ad_array`.
 =#
 _upload(x::Array, ::Nothing) = Reactant.to_rarray(x)
 _upload(x::AbstractArray, ::Nothing) = Reactant.to_rarray(Array(x))
+#= XLA only ever sees the stage, so a short `x` would leave its tail holding the
+previous call's values and the compiled function would run on them. Uploading
+`x` directly would have failed XLA's shape check instead. =#
 function _upload(x::AbstractArray, stage::AbstractArray)
+    size(x) == size(stage) || throw(
+        DimensionMismatch(
+            "AutoReactant compiled for input size $(size(stage)), got $(size(x))"
+        ),
+    )
     copyto!(stage, x)
     return Reactant.to_rarray(stage)
 end
@@ -55,8 +62,8 @@ function _in_stage(template::AbstractArray)
     return ParallelMCMC._host_staging_buffer(template, eltype(template), size(template))
 end
 
-# Keep promotions performed by the compiled function. Callers hold on to
-# gradients, so the result is a fresh array every call.
+# Keep any promotion the compiled function performed. Callers hold on to
+# gradients, so every call returns a fresh array.
 function _promote_like(template::AbstractArray, out_h::Array)
     template isa Array && eltype(out_h) === eltype(template) && return out_h
     res = similar(template, eltype(out_h), size(out_h))
@@ -64,7 +71,7 @@ function _promote_like(template::AbstractArray, out_h::Array)
     return res
 end
 
-# Raw pointer to an XLA buffer, or `nothing` if it cannot be taken: pointer
+# Device pointer to an XLA buffer, or `nothing` if it cannot be taken: pointer
 # access needs an unsharded PJRT buffer that is not already on the host.
 function _device_pointer(out)
     out isa Reactant.ConcretePJRTArray || return nothing
@@ -86,10 +93,9 @@ function _download(template::AbstractArray, out, platform::AbstractString)
     return _promote_like(template, Array(out))
 end
 
-# One of these per compiled function, owning its staging buffers. `template` is
-# the first argument; HVP outputs follow `x`, not `v`. `stages` has one entry
-# per argument, `nothing` or a host buffer. `platform` is fixed at compile time
-# from the default XLA client.
+# One per compiled function, owning a staging buffer per argument. `template`
+# is the first argument, whose array type the result is rebuilt as.
+# `platform` is read once, at compile time, from the default XLA client.
 struct ReactantCall{F,S,T}
     compiled::F
     stages::S
diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl
index a2f279d..69a5873 100644
--- a/src/ParallelMCMC.jl
+++ b/src/ParallelMCMC.jl
@@ -43,8 +43,8 @@ needs_host_staging(::AbstractArray) = false
 """
     _host_staging_buffer(template::AbstractArray, ::Type{T}, dims::Dims) -> Array{T}
 
-Host buffer for staging transfers to and from arrays like `template`. Device
-extensions can return pinned memory. Defaults to a plain `Array`.
+Host buffer for staging copies between the host and arrays like `template`.
+Device extensions can return pinned memory. Defaults to a plain `Array`.
 """
 function _host_staging_buffer(::AbstractArray, ::Type{T}, dims::Dims) where {T}
     return Array{T}(undef, dims)
@@ -53,11 +53,10 @@ end
 """
     _copy_from_device_pointer!(dest::AbstractArray, ptr::Ptr{Cvoid}, platform::AbstractString) -> Bool
 
-Copy device memory owned by another runtime into `dest`, returning `true` on
-success and `false` if `dest`'s array type cannot address `ptr`. `platform` is
-the owner's XLA platform name ("cuda", "rocm", "cpu"). The copy is complete
-when this returns, so the caller may release the source afterwards. Defaults
-to `false`.
+Copy device memory owned by another runtime into `dest`. `platform` is the
+owner's XLA platform name ("cuda", "rocm", "cpu"). Returns `false` if `dest`'s
+array type cannot address `ptr`, which is the default; on `true` the copy has
+finished, so the caller may release the source.
 """
 _copy_from_device_pointer!(::AbstractArray, ::Ptr{Cvoid}, ::AbstractString) = false
 
diff --git a/test/staged_array.jl b/test/staged_array.jl
index 4fb252b..27e8398 100644
--- a/test/staged_array.jl
+++ b/test/staged_array.jl
@@ -1,7 +1,7 @@
-#= Device-array stand-in for the host-staging paths, so they can be tested
-without a GPU. Scalar `setindex!` errors, as on `CuArray`. Shared by
-test-CUDA-Extension.jl and test-Reactant-HVP.jl; the name does not match
-`test-*.jl`, so runtests.jl does not include it as a testset. =#
+#= Device-array stand-in, so the host-staging paths can be exercised without a
+GPU: scalar `setindex!` errors, as on `CuArray`. Included by
+test-CUDA-Extension.jl and test-Reactant-HVP.jl; the name is off the
+`test-*.jl` pattern so runtests.jl does not pick it up as a testset. =#
 
 struct StagedArray{T,N} <: AbstractArray{T,N}
     data::Array{T,N}
@@ -15,8 +15,8 @@ function Base.similar(a::StagedArray, ::Type{T}, dims::Dims) where {T}
 end
 Base.copyto!(a::StagedArray, src::AbstractArray) = (copyto!(a.data, src); a)
 
-# `_prepare_model` fills the batched template with `X_template .= x_template`.
-# The zero-dim method disambiguates against Base's, so scalar fills also work.
+# `_prepare_model` builds the batched template with `X_template .= x_template`.
+# The zero-dim method only resolves an ambiguity with Base's.
 Base.BroadcastStyle(::Type{<:StagedArray}) = Broadcast.ArrayStyle{StagedArray}()
 Base.copyto!(a::StagedArray, bc::Broadcast.Broadcasted) = (copyto!(a.data, bc); a)
 function Base.copyto!(
diff --git a/test/test-CUDA-Extension.jl b/test/test-CUDA-Extension.jl
index 5888db6..e2dee34 100644
--- a/test/test-CUDA-Extension.jl
+++ b/test/test-CUDA-Extension.jl
@@ -112,8 +112,10 @@ end
     catch
         false
     end
-    if !cuda_functional
-        @info "CUDAExt device buffer hooks: CUDA not functional, skipping"
+    # The extension check is what "CUDAExt loads with CUDA" above asserts; repeat
+    # it as a guard so a load failure reports once instead of failing every hook.
+    if !cuda_functional || Base.get_extension(ParallelMCMC, :CUDAExt) === nothing
+        @info "CUDAExt device buffer hooks: CUDAExt unavailable, skipping"
     else
         template = CUDA.zeros(Float32, 3)
 
diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl
index 5bc9977..90a9921 100644
--- a/test/test-Reactant-HVP.jl
+++ b/test/test-Reactant-HVP.jl
@@ -213,8 +213,8 @@ isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array
         end
     end
 
-    # `StagedArray` has no `_copy_from_device_pointer!` method, so results
-    # come back through a plain host download, as they would for a CuArray on the CPU client.
+    # `StagedArray` has no `_copy_from_device_pointer!` method, so results come
+    # back through a plain host download, as a `CuArray` would on the CPU client.
     @testset "staged template" begin
         rng = MersenneTwister(74)
         x0 = StagedArray(randn(rng, Float32, D_R))
@@ -245,6 +245,16 @@ isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array
         @test Hv1 !== Hv2
         @test Hv1.data !== Hv2.data
 
+        # Everything above would also pass if no stage had been allocated.
+        @test only(m_p.grad_logdensity.stages) isa Vector{Float32}
+        @test size(only(m_p.grad_logdensity.stages)) == (D_R,)
+        @test all(s -> s isa Vector{Float32}, m_p.hvp.stages)
+
+        # XLA only ever sees the stage, so a short input has to be rejected here.
+        @test_throws DimensionMismatch m_p.grad_logdensity(
+            StagedArray(randn(rng, Float32, D_R - 1))
+        )
+
         @testset "batched slots" begin
             T = 8
             model_b = DensityModel(

From be9d3dd0b8a742e0876d050cf2844e37409105e1 Mon Sep 17 00:00:00 2001
From: Ryan Senne <50930199+rsenne@users.noreply.github.com>
Date: Sat, 26 Sep 2026 20:10:38 -0400
Subject: [PATCH 4/4] fix memory leak

---
 ext/ReactantExt.jl        | 21 ++++++++++++++++++++-
 test/test-Reactant-HVP.jl | 12 ++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl
index d1c7a42..f7dfcf1 100644
--- a/ext/ReactantExt.jl
+++ b/ext/ReactantExt.jl
@@ -103,8 +103,27 @@ struct ReactantCall{F,S,T}
     platform::String
 end
 
+#=
+XLA buffers are freed by GC finalizers. Reusing stages creates very little host
+garbage, so GC may not run before device memory fills up. Finalize buffers as
+soon as they are no longer needed; PJRT waits for any pending work before freeing
+them. Finalization runs only once, even if a compiled function returns an input
+wrapped in a new array. Donated inputs no longer own their buffers, so skip them.
+=#
+function _free!(x)
+    x isa Reactant.ConcretePJRTArray || return nothing
+    (x.donated || Reactant.Sharding.is_sharded(x.sharding)) && return nothing
+    finalize(Reactant.get_buffer(x))
+    return nothing
+end
+
 function (c::ReactantCall)(args::AbstractArray...)
-    return _download(c.template, c.compiled(map(_upload, args, c.stages)...), c.platform)
+    inputs = map(_upload, args, c.stages)
+    out = c.compiled(inputs...)
+    res = _download(c.template, out, c.platform)
+    foreach(_free!, inputs)
+    _free!(out)
+    return res
 end
 
 function _compiled(core, templates::AbstractArray...)
diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl
index 90a9921..4d112f8 100644
--- a/test/test-Reactant-HVP.jl
+++ b/test/test-Reactant-HVP.jl
@@ -283,6 +283,18 @@ isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array
         end
     end
 
+    @testset "eager buffer free survives an aliased result" begin
+        ext = Base.get_extension(ParallelMCMC, :ReactantExt)
+        rng = MersenneTwister(75)
+        for template in (randn(rng, Float32, D_R), StagedArray(randn(rng, Float32, D_R)))
+            call = ext._compiled(identity, template)
+            xs = [typeof(template)(randn(rng, Float32, D_R)) for _ in 1:50]
+            ys = map(call, xs)
+            GC.gc(true)
+            @test all(((x, y),) -> collect(y) == collect(x), zip(xs, ys))
+        end
+    end
+
     reactant_gpu_ok = try
         using CUDA: CUDA
         CUDA.functional() && (CUDA.CuArray([1.0f0]); true)