From 0d6f365a7059e64bff9dbde21d858d40ead38d89 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 11 Aug 2021 13:23:02 -0400 Subject: [PATCH 1/5] Make inference's job much easier by avoiding `map` --- 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 0cb521fc..df0474f2 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -14,10 +14,10 @@ void_setindex!(args...) = (setindex!(args...); return) const default_chunk_size = ForwardDiff.pickchunksize -function ForwardColorJacCache(f,x,_chunksize = nothing; +function ForwardColorJacCache(f::F,x,_chunksize = nothing; dx = nothing, colorvec=1:length(x), - sparsity::Union{AbstractArray,Nothing}=nothing) + sparsity::Union{AbstractArray,Nothing}=nothing) where {F} if _chunksize isa Nothing chunksize = ForwardDiff.pickchunksize(maximum(colorvec)) @@ -25,7 +25,11 @@ function ForwardColorJacCache(f,x,_chunksize = nothing; chunksize = _chunksize end - p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) + if x isa Array + p = generate_chunked_partials(x,colorvec,chunksize) + else + p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) + end _t = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}.(vec(x),first(p)) t = ArrayInterface.restructure(x,_t) if dx isa Nothing @@ -50,7 +54,10 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize padding_matrix = BitMatrix(undef, length(x), padding_size) partials = hcat(partials, padding_matrix) - chunked_partials = map(i -> Tuple.(eachrow(partials[:,(i-1)*chunksize+1:i*chunksize])),1:num_of_chunks) + chunked_partials = Vector{Vector{NTuple{chunksize,eltype(x)}}}(undef, num_of_chunks) + for i in 1:num_of_chunks + chunked_partials[i] = Tuple.(eachrow(@view(partials[:,(i-1)*chunksize+1:i*chunksize]))) + end chunked_partials end From e4cb32229bea6488efb461cbe75c6eb6c3b17661 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 11 Aug 2021 14:47:07 -0400 Subject: [PATCH 2/5] A few more array specializations --- src/differentiation/compute_jacobian_ad.jl | 52 ++++++++++++++++++---- 1 file changed, 44 insertions(+), 8 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index df0474f2..0a554212 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -27,11 +27,17 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; if x isa Array p = generate_chunked_partials(x,colorvec,chunksize) + t = similar(x,Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}) + for i in eachindex(t) + t[i] = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}(x[i],first(p)[1]) + end else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) + _t = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}.(vec(x),first(p)) + t = ArrayInterface.restructure(x,_t) end - _t = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}.(vec(x),first(p)) - t = ArrayInterface.restructure(x,_t) + + if dx isa Nothing fx = similar(t) _dx = similar(x) @@ -50,16 +56,25 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize maxcolor = maximum(colorvec) num_of_chunks = Int(ceil(maxcolor / chunksize)) padding_size = (chunksize - (maxcolor % chunksize)) % chunksize - partials = colorvec .== (1:maxcolor)' + + # partials = colorvec .== (1:maxcolor)' + partials = BitMatrix(undef, length(colorvec), maxcolor) + for j in 1:length(colorvec), i in 1:maxcolor + partials[i,j] = colorvec[j] == i + end + padding_matrix = BitMatrix(undef, length(x), padding_size) partials = hcat(partials, padding_matrix) chunked_partials = Vector{Vector{NTuple{chunksize,eltype(x)}}}(undef, num_of_chunks) for i in 1:num_of_chunks - chunked_partials[i] = Tuple.(eachrow(@view(partials[:,(i-1)*chunksize+1:i*chunksize]))) + tmp = Vector{NTuple{chunksize,eltype(x)}}(undef, chunksize) + for j in 1:size(partials,1) + tmp[j] = Tuple(@view partials[j,(i-1)*chunksize+1:i*chunksize]) + end + chunked_partials[i] = tmp end chunked_partials - end @inline function forwarddiff_color_jacobian(f, @@ -287,11 +302,26 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, for i in eachindex(p) partial_i = p[i] - vect .= Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i) + + if vect isa Array + @inbounds @simd ivdep for j in eachindex(vect) + vect = Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}(vecx[j], partial_i[j]) + end + else + vect .= Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i) + end + f(fx,t) if !(sparsity isa Nothing) for j in 1:chunksize - dx .= partials.(fx, j) + + if dx isa Array + @inbounds @simd ivdep for k in eachindex(dx) + dx[k] = partials(fx[k], j) + end + else + dx .= partials.(fx, j) + end if ArrayInterface.fast_scalar_indexing(dx) #dx is implicitly used in vecdx @@ -320,7 +350,13 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, for j in 1:chunksize col_index = (i-1)*chunksize + j (col_index > ncols) && return J - J[:, col_index] .= partials.(vecfx, j) + if J isa Array + @inbounds @simd for k in 1:size(J,1) + J[k, col_index] = partials(vecfx[k], j) + end + else + J[:, col_index] .= partials.(vecfx, j) + end end end end From 9f6d7f0e20e2c0b9432e39f69aa163ee5a8cf656 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 11 Aug 2021 15:33:39 -0400 Subject: [PATCH 3/5] fix loop ordering --- src/differentiation/compute_jacobian_ad.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 0a554212..51c5b75b 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -59,8 +59,8 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize # partials = colorvec .== (1:maxcolor)' partials = BitMatrix(undef, length(colorvec), maxcolor) - for j in 1:length(colorvec), i in 1:maxcolor - partials[i,j] = colorvec[j] == i + for i in 1:maxcolor, j in 1:length(colorvec) + partials[j,i] = colorvec[j] == i end padding_matrix = BitMatrix(undef, length(x), padding_size) From ee7268c0f26f38589503930947da331a8d6363af Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 11 Aug 2021 15:54:26 -0400 Subject: [PATCH 4/5] fix tmp size --- src/differentiation/compute_jacobian_ad.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 51c5b75b..3085dfb9 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -66,9 +66,11 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize padding_matrix = BitMatrix(undef, length(x), padding_size) partials = hcat(partials, padding_matrix) + + #chunked_partials = map(i -> Tuple.(eachrow(partials[:,(i-1)*chunksize+1:i*chunksize])),1:num_of_chunks) chunked_partials = Vector{Vector{NTuple{chunksize,eltype(x)}}}(undef, num_of_chunks) for i in 1:num_of_chunks - tmp = Vector{NTuple{chunksize,eltype(x)}}(undef, chunksize) + tmp = Vector{NTuple{chunksize,eltype(x)}}(undef, size(partials,1)) for j in 1:size(partials,1) tmp[j] = Tuple(@view partials[j,(i-1)*chunksize+1:i*chunksize]) end From 73f201814f6ce08c5a53177ded3251d72a277a1a Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 11 Aug 2021 16:23:06 -0400 Subject: [PATCH 5/5] fix vect indexing --- src/differentiation/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 3085dfb9..59a0467b 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -307,7 +307,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, if vect isa Array @inbounds @simd ivdep for j in eachindex(vect) - vect = Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}(vecx[j], partial_i[j]) + vect[j] = Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}(vecx[j], partial_i[j]) end else vect .= Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i)