Make seeds and extract_jacobian gpu-friendly

Use broadcast/macro consistently

Fix jac

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"

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), 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