Don't preallocate GradientConfig in ForwardDiff backend by default

ext/LogDensityProblemsADForwardDiffExt.jl

@@ -18,24 +18,50 @@ end
 # Load DiffResults helpers
 include("DiffResults_helpers.jl")
 
-struct ForwardDiffLogDensity{L, C} <: ADGradientWrapper
+struct ForwardDiffLogDensity{L, C <: ForwardDiff.Chunk, T <: Union{Nothing,ForwardDiff.Tag},
+                             G <: Union{Nothing,ForwardDiff.GradientConfig}} <: ADGradientWrapper
+    "supports zero-order evaluation `logdensity(ℓ, x)`"
     ℓ::L
-    gradientconfig::C
+    "chunk size for ForwardDiff"
+    chunk::C
+    "tag, or `nothing` for the default"
+    tag::T
+    "gradient config, or `nothing` if created for each evaluation"
+    gradient_config::G
 end
 
 function Base.show(io::IO, ℓ::ForwardDiffLogDensity)
     print(io, "ForwardDiff AD wrapper for ", ℓ.ℓ,
-          ", w/ chunk size ", length(ℓ.gradientconfig.seeds))
+          ", w/ chunk size ", ForwardDiff.chunksize(ℓ.chunk))
 end
 
 _chunk(chunk::ForwardDiff.Chunk) = chunk
 _chunk(chunk::Integer) = ForwardDiff.Chunk(chunk)
 
 _default_chunk(ℓ) = _chunk(dimension(ℓ))
 
-_default_gradientconfig(ℓ, chunk, ::Nothing) = _default_gradientconfig(ℓ, chunk, zeros(dimension(ℓ)))
-function _default_gradientconfig(ℓ, chunk, x::AbstractVector)
-    return ForwardDiff.GradientConfig(Base.Fix1(logdensity, ℓ), x, _chunk(chunk))
+function Base.copy(fℓ::ForwardDiffLogDensity{L,C,T,<:ForwardDiff.GradientConfig}) where {L,C,T}
+    @unpack ℓ, chunk, tag, gradient_config = fℓ
+    ForwardDiffLogDensity(ℓ, chunk, tag, copy(gradient_config))
+end
+
+"""
+$(SIGNATURES)
+
+Make a `ForwardDiff.GradientConfig` for type `T`. `tag = nothing` generates the default tag.
+
+Return the function for evaluating log density (with a vector argument) as the second value.
+"""
+function _make_gradient_config(::Type{T}, ℓ, chunk, tag) where T
+    ℓ′ = Base.Fix1(logdensity, ℓ)
+    x = zeros(T, dimension(ℓ))
+    c = _chunk(chunk)
+    gradient_config = if tag ≡ nothing
+        ForwardDiff.GradientConfig(ℓ′, x, c)
+    else
+        ForwardDiff.GradientConfig(ℓ′, x, c, tag)
+    end
+    gradient_config, ℓ′
 end
 
 """
@@ -45,22 +71,44 @@ end
 Wrap a log density that supports evaluation of `Value` to handle `ValueGradient`, using
 `ForwardDiff`.
 
-Keywords are passed on to `ForwardDiff.GradientConfig` to customize the setup. In
-particular, chunk size can be set with a `chunk` keyword argument (accepting an integer or a
-`ForwardDiff.Chunk`), and the underlying vector used by `ForwardDiff` can be set with the
-`x` keyword argument (accepting an `AbstractVector`).
+Keyword arguments:
+
+- `chunk` can be used to set the chunk size, an integer or a `ForwardDiff.Chunk`
+
+- `tag` (default: `nothing`) can be used to set a tag for `ForwardDiff`
+
+- `gradient_config_type` (default: `Union{}`) can be used to preallocate a
+  `ForwardDiff.GradientConfig` for the given type. With the default, one is created for each
+  evaluation.
+
+   Note **pre-allocating a `ForwardDiff.GradienConfig` is not thread-safe**. You can
+   [`copy`](@ref) the results for concurrent evaluation:
+   ```julia
+   ∇ℓ1 = ADgradient(:ForwardDiff, ℓ; gradient_config = my_gradient_config)
+   ∇ℓ2 = copy(∇ℓ1) # you can now use both, in different threads
+   ```
 """
 function ADgradient(::Val{:ForwardDiff}, ℓ;
-                    x::Union{Nothing,AbstractVector} = nothing,
                     chunk::Union{Integer,ForwardDiff.Chunk} = _default_chunk(ℓ),
-                    gradientconfig::ForwardDiff.GradientConfig = _default_gradientconfig(ℓ, chunk, x))
-    ForwardDiffLogDensity(ℓ, gradientconfig)
+                    tag::Union{Nothing,ForwardDiff.Tag} = nothing,
+                    gradient_config_type::Type{T} = Union{}) where {T}
+    gradient_config = if gradient_config_type ≡ Union{}
+        nothing
+    else
+        first(_make_gradient_config(T, ℓ, chunk, tag))
+    end
+    ForwardDiffLogDensity(ℓ, chunk, tag, gradient_config)
 end
 
 function logdensity_and_gradient(fℓ::ForwardDiffLogDensity, x::AbstractVector)
-    @unpack ℓ, gradientconfig = fℓ
+    @unpack ℓ, chunk, tag, gradient_config = fℓ
     buffer = _diffresults_buffer(x)
-    result = ForwardDiff.gradient!(buffer, Base.Fix1(logdensity, ℓ), x, gradientconfig)
+    if gradient_config ≡ nothing
+        gradient_config, ℓ′ = _make_gradient_config(eltype(buffer), ℓ, chunk, tag)
+    else
+        ℓ′ = Base.Fix1(logdensity, ℓ)
+    end
+    result = ForwardDiff.gradient!(buffer, ℓ′, x, gradient_config)
     _diffresults_extract(result)
 end
 

src/LogDensityProblemsAD.jl

@@ -11,7 +11,6 @@ using LogDensityProblems: LogDensityOrder
 
 import UnPack
 
-
 #####
 ##### AD wrappers --- interface and generic code
 #####
@@ -34,6 +33,8 @@ dimension(ℓ::ADGradientWrapper) = dimension(ℓ.ℓ)
 
 Base.parent(ℓ::ADGradientWrapper) = ℓ.ℓ
 
+Base.copy(x::ADGradientWrapper) = x # no-op, except for ForwardDiff
+
 """
 $(SIGNATURES)
 
@@ -57,6 +58,10 @@ ADgradient(:ForwardDiff, P)
 and should mostly be equivalent if the compiler manages to fold the constant.
 
 The function `parent` can be used to retrieve the original argument.
+
+!!! note
+    With the default options, automatic differentiation preserves thread-safety. See
+    exceptions and workarounds in the docstring for each backend.
 """
 ADgradient(kind::Symbol, P; kwargs...) = ADgradient(Val{kind}(), P; kwargs...)
 

test/runtests.jl

@@ -115,13 +115,14 @@ end
             (test_logdensity(x), test_gradient(x))
     end
 
-    # Make sure that other types are supported.
+    # preallocated gradient config
     x = randexp(Float32, 3)
-    ∇ℓ = ADgradient(:ForwardDiff, ℓ; x=x)
+    ∇ℓ = ADgradient(:ForwardDiff, ℓ; gradient_config_type = Float32)
     @test eltype(first(logdensity_and_gradient(∇ℓ, x))) === Float32
     @test @inferred(logdensity(∇ℓ, x)) ≅ test_logdensity(x)
     @test @inferred(logdensity_and_gradient(∇ℓ, x)) ≅
         (test_logdensity(x), test_gradient(x))
+    @test @inferred(copy(∇ℓ)).gradient_config ≢ ∇ℓ.gradient_config
 end
 
 @testset "chunk heuristics for ForwardDiff" begin