Reuse Reactant staging buffers and copy results on the device

.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"

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` 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
 

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()`.

ext/CUDAExt.jl

@@ -1,8 +1,30 @@
 module CUDAExt
 
 using ParallelMCMC: ParallelMCMC
-using CUDA: CuArray
+using CUDA: CUDA, CuArray, CuPtr
 
 ParallelMCMC.needs_host_staging(::CuArray) = true
 
+#= 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}
+    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;
+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}
+    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

ext/ReactantExt.jl

@@ -25,44 +25,93 @@ 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)
+#=
+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
+
+# 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
 
-# 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 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))
     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))
+# 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
+    Reactant.Sharding.is_sharded(out.sharding) && return nothing
+    wait(out)
+    buf = Reactant.get_buffer(out)
+    Reactant.XLA.buffer_on_cpu(buf) && return nothing
+    return Reactant.XLA.unsafe_buffer_pointer(buf)
+end
+
+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 begin
+            ParallelMCMC._copy_from_device_pointer!(res, ptr, platform) && return res
+        end
+    end
+    return _promote_like(template, Array(out))
+end
+
+# 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
+    template::T
+    platform::String
+end
 
-function _compiled(core, t1::AbstractArray)
-    _warn_reactant_host_roundtrip(t1)
-    compiled = @compile core(_upload(t1))
-    return x -> _from_host(x, compiled(_upload(x)))
+function (c::ReactantCall)(args::AbstractArray...)
+    return _download(c.template, c.compiled(map(_upload, args, c.stages)...), c.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)))
+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.

src/ParallelMCMC.jl

@@ -40,6 +40,26 @@ 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 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)
+end
+
+"""
+    _copy_from_device_pointer!(dest::AbstractArray, ptr::Ptr{Cvoid}, platform::AbstractString) -> Bool
+
+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
+
 #= 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."

test/staged_array.jl

@@ -0,0 +1,29 @@
+#= 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}
+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` 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!(
+    a::StagedArray, bc::Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{0}}
+)
+    copyto!(a.data, bc)
+    return a
+end
+
+ParallelMCMC.needs_host_staging(::StagedArray) = true

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))
@@ -85,6 +73,20 @@ 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._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._copy_from_device_pointer!(x, C_NULL, "cpu") == false
+end
+
 @testset "CUDAExt loads with CUDA" begin
     cuda_loadable = try
         using CUDA
@@ -102,3 +104,39 @@ 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
+    # 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)
+
+        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)))
+        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