From e0be439fb28c294c92b7a2769e95ef99077f81fb Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 9 Nov 2020 11:00:16 -0800
Subject: [PATCH] Make seeds and extract_jacobian gpu-friendly

Use broadcast/macro consistently

Fix jac

Add AllocationsTest.jl
---
 Project.toml            |  2 +-
 src/apiutils.jl         | 22 ++++++++--------------
 src/jacobian.jl         | 19 ++++++++++---------
 test/AllocationsTest.jl | 33 +++++++++++++++++++++++++++++++++
 test/runtests.jl        |  6 ++++++
 5 files changed, 58 insertions(+), 24 deletions(-)
 create mode 100644 test/AllocationsTest.jl

diff --git a/Project.toml b/Project.toml
index 0d27764b..9c380865 100644
--- a/Project.toml
+++ b/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"
diff --git a/src/apiutils.jl b/src/apiutils.jl
index ae5d78ad..654c64a3 100644
--- a/src/apiutils.jl
+++ b/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
diff --git a/src/jacobian.jl b/src/jacobian.jl
index 39332b42..c24f294e 100644
--- a/src/jacobian.jl
+++ b/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
 
diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl
new file mode 100644
index 00000000..a16a0103
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index e19c9527..0b9b1d8b 100644
--- a/test/runtests.jl
+++ b/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