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 931068b..fc28fcc 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` 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/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..8138bb6 100644 --- a/ext/CUDAExt.jl +++ b/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 diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 32a0858..f7dfcf1 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -25,44 +25,112 @@ 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. -# 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) +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 + +# 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 +# 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 + #= -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. +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. =# -_upload(x::AbstractArray) = Reactant.to_rarray(_host(x)) +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 _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...) + 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, 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. diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 36b59fd..69a5873 100644 --- a/src/ParallelMCMC.jl +++ b/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." diff --git a/test/staged_array.jl b/test/staged_array.jl new file mode 100644 index 0000000..27e8398 --- /dev/null +++ b/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 diff --git a/test/test-CUDA-Extension.jl b/test/test-CUDA-Extension.jl index 4a58896..e2dee34 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)) @@ -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 diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index c71bdcc..4d112f8 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,8 @@ const CT_R = FlexiChains.FlexiChain{Symbol} using Reactant: Reactant using Enzyme: Enzyme +isdefined(@__MODULE__, :StagedArray) || include(joinpath(@__DIR__, "staged_array.jl")) + @testset "Reactant HVP" begin @testset "extension is loaded" begin @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing @@ -133,6 +136,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 +213,88 @@ using Enzyme: Enzyme end end + # `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)) + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x0, 8, nothing) + + 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 StagedArray{Float32,1} + @test eltype(g1) === Float32 + @test g1.data ≈ gradlogp_r(x1.data) + @test Hv1 isa StagedArray{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 + + # 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( + 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 = 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 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 StagedArray{Float32,2} + @test Hvb.data ≈ Hv_cols + 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) @@ -226,6 +323,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