From 9c75bb620c5b046abb65c7eabd57d07d79ddc688 Mon Sep 17 00:00:00 2001 From: Fredrik Bagge Carlson Date: Tue, 23 Aug 2022 12:03:21 +0200 Subject: [PATCH] function barrier reduces impact of type instability Looping over `t` caused lots of allocations. This PR introduces a function barrier that lessens the impact of this to almost negligible --- src/differentiation/compute_jacobian_ad.jl | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 6a89b24b..818d36a8 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -36,10 +36,8 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; if x isa Array p = generate_chunked_partials(x,colorvec,chunksize) - t = Array{Dual{T,eltype(x),length(first(first(p)))}}(undef,size(x)) - for i in eachindex(t) - t[i] = Dual{T,eltype(x),length(first(first(p)))}(x[i],ForwardDiff.Partials(first(p)[i])) - end + DT = Dual{T,eltype(x),length(first(first(p)))} + t = _get_t(DT, x, p) else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) _t = Dual{T,eltype(x),getsize(chunksize)}.(vec(x),ForwardDiff.Partials.(first(p))) @@ -60,6 +58,15 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; ForwardColorJacCache(t,fx,_dx,p,colorvec,sparsity,getsize(chunksize)) end +# Function barrier for unknown constructor type +function _get_t(::Type{DT}, x, p) where DT + t = similar(x, DT) + for i in eachindex(t) + t[i] = DT(x[i],ForwardDiff.Partials(first(p)[i])) + end + t +end + generate_chunked_partials(x,colorvec,N::Integer) = generate_chunked_partials(x,colorvec,Val(N)) function generate_chunked_partials(x,colorvec,cs::Val{chunksize}) where chunksize maxcolor = maximum(colorvec)