Make seeds and extract_jacobian gpu-friendly

Use broadcast/macro consistently

Fix jac

Add AllocationsTest.jl

Author: charleskawczynski

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "0.10.12"
+version = "0.10.13"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"
@@ -28,7 +28,8 @@ DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Calculus", "DiffTests", "LinearAlgebra", "SparseArrays", "Test", "InteractiveUtils"]
+test = ["Calculus", "DiffTests", "LinearAlgebra", "SparseArrays", "Test", "InteractiveUtils", "BenchmarkTools"]

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,27 @@
+module AllocationsTest
+
+using ForwardDiff
+using BenchmarkTools
+
+include(joinpath(dirname(@__FILE__), "utils.jl"))
+
+@testset "Test seed! allocations" begin
+    x = rand(1000)
+    cfg = ForwardDiff.GradientConfig(nothing, x)
+
+    balloc = @ballocated ForwardDiff.seed!($(cfg.duals), $x, $(cfg.seeds))
+    @test balloc == 0
+
+    balloc = @ballocated ForwardDiff.seed!($(cfg.duals), $x, $(cfg.seeds[1]))
+    @test balloc == 0
+
+    index = 1
+    balloc = @ballocated ForwardDiff.seed!($(cfg.duals), $x, $index, $(cfg.seeds))
+    @test balloc == 0
+
+    index = 1
+    balloc = @ballocated ForwardDiff.seed!($(cfg.duals), $x, $index, $(cfg.seeds[1]))
+    @test balloc == 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