Make seed! and extract_jacobian! gpu-friendly

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "0.10.13"
+version = "0.10.14"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"

src/apiutils.jl

@@ -55,36 +55,30 @@ end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
                seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
-    for i in eachindex(duals)
-        duals[i] = Dual{T,V,N}(x[i], seed)
-    end
+    duals .= Dual{T,V,N}.(x, Ref(seed))
     return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
                seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    for i in 1:N
-        duals[i] = Dual{T,V,N}(x[i], seeds[i])
-    end
+    dual_inds = 1:N
+    duals[dual_inds] .= Dual{T,V,N}.(view(x,dual_inds), seeds)
     return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
     offset = index - 1
-    for i in 1:N
-        j = i + offset
-        duals[j] = Dual{T,V,N}(x[j], seed)
-    end
+    dual_inds = (1:N) .+ offset
+    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), Ref(seed))
     return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
     offset = index - 1
-    for i in 1:chunksize
-        j = i + offset
-        duals[j] = Dual{T,V,N}(x[j], seeds[i])
-    end
+    seed_inds = 1:chunksize
+    dual_inds = seed_inds .+ offset
+    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), getindex.(Ref(seeds), seed_inds))
     return duals
 end

src/jacobian.jl

@@ -111,9 +111,10 @@ end
 
 function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T}
     out_reshaped = reshape(result, length(ydual), n)
-    for col in 1:size(out_reshaped, 2), row in 1:size(out_reshaped, 1)
-        out_reshaped[row, col] = partials(T, ydual[row], col)
-    end
+    ydual_reshaped = vec(ydual)
+    # Use closure to avoid GPU broadcasting with Type
+    partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
+    out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n))
     return result
 end
 
@@ -123,13 +124,13 @@ function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::Abstract
 end
 
 function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T}
+    ydual_reshaped = vec(ydual)
     offset = index - 1
-    for i in 1:chunksize
-        col = i + offset
-        for row in eachindex(ydual)
-            result[row, col] = partials(T, ydual[row], i)
-        end
-    end
+    irange = 1:chunksize
+    col = irange .+ offset
+    # Use closure to avoid GPU broadcasting with Type
+    partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
+    result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange))
     return result
 end
 

test/AllocationsTest.jl

@@ -0,0 +1,33 @@
+module AllocationsTest
+
+using ForwardDiff
+
+include(joinpath(dirname(@__FILE__), "utils.jl"))
+
+@testset "Test seed! allocations" begin
+    x = rand(1000)
+    cfg = ForwardDiff.GradientConfig(nothing, x)
+    duals = cfg.duals
+    seeds = cfg.seeds
+    seed = cfg.seeds[1]
+
+    alloc = @allocated ForwardDiff.seed!(duals, x, seeds)
+    alloc = @allocated ForwardDiff.seed!(duals, x, seeds)
+    @test alloc == 0
+
+    alloc = @allocated ForwardDiff.seed!(duals, x, seed)
+    alloc = @allocated ForwardDiff.seed!(duals, x, seed)
+    @test alloc == 0
+
+    index = 1
+    alloc = @allocated ForwardDiff.seed!(duals, x, index, seeds)
+    alloc = @allocated ForwardDiff.seed!(duals, x, index, seeds)
+    @test alloc == 0
+
+    index = 1
+    alloc = @allocated ForwardDiff.seed!(duals, x, index, seed)
+    alloc = @allocated ForwardDiff.seed!(duals, x, index, seed)
+    @test alloc == 0
+end
+
+end

test/runtests.jl

@@ -31,3 +31,9 @@ println("done (took $t seconds).")
 println("Testing miscellaneous functionality...")
 t = @elapsed include("MiscTest.jl")
 println("done (took $t seconds).")
+
+if VERSION >= v"1.5-"
+    println("Testing allocations...")
+    t = @elapsed include("AllocationsTest.jl")
+    println("done (took $t seconds).")
+end