From 23a30630cef43b611ceec5b24617765cf429e9b5 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 05:29:33 -0400 Subject: [PATCH 1/6] setup chunking structure --- src/compute_jacobian_ad.jl | 35 +++++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index fbfee2c5..804e0949 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -1,4 +1,4 @@ -using ForwardDiff: Dual, jacobian, partials +using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD struct ForwardColorJacCache{T,T2,T3,T4,T5} t::T @@ -8,10 +8,24 @@ struct ForwardColorJacCache{T,T2,T3,T4,T5} color::T5 end -function ForwardColorJacCache(f,x; +function default_chunk_size(maxcolor) + if maxcolor < DEFAULT_CHUNK_THRESHOLD + Val(maxcolor) + else + Val(DEFAULT_CHUNK_THRESHOLD) + end +end + +function ForwardColorJacCache(f,x,_chunksize = nothing; dx = nothing, color=1:length(x)) + if _chunksize === nothing + chunksize = default_chunk_size(maximum(color)) + else + chunksize = _chunksize + end + t = zeros(Dual{typeof(f), eltype(x), maximum(color)},length(x)) if dx === nothing @@ -22,6 +36,16 @@ function ForwardColorJacCache(f,x; _dx = dx end + p = generate_chunked_partials(x,color,chunksize) + ForwardColorJacCache(t,fx,_dx,p,color) +end + +generate_chunked_partials(x,color,N::Integer) = generate_chunked_partials(x,color,Val(N)) +function generate_chunked_partials(x,color,::Val{N}) where N + + # TODO: should only go up to the chunksize each time, and should + # generate p[i] different parts, each with less than the chunksize + partials_array = Array{Bool}(undef, length(x), maximum(color)) for color_i in 1:maximum(color) for i in 1:length(x) @@ -33,22 +57,20 @@ function ForwardColorJacCache(f,x; end end p = Tuple.(eachrow(partials_array)) - - ForwardColorJacCache(t,fx,_dx,p,color) end function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}; dx = nothing, - color) + color = 1:length(x)) forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color)) end function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, - jac_cache = ForwardColorJacCache(f,x)) + jac_cache::ForwardColorJacCache) t = jac_cache.t fx = jac_cache.fx @@ -56,6 +78,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, p = jac_cache.p color = jac_cache.color + # TODO: Should compute on each p[i] and decompress t .= Dual{typeof(f)}.(x, p) f(fx, t) From 213b1ee594af589467713f7c3cf927647d402d91 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 05:32:25 -0400 Subject: [PATCH 2/6] set the partial length to the chunksize --- src/compute_jacobian_ad.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index 804e0949..af335b53 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -16,6 +16,9 @@ function default_chunk_size(maxcolor) end end +getsize(::Val{N}) where N = N +getsize(N::Integer) = N + function ForwardColorJacCache(f,x,_chunksize = nothing; dx = nothing, color=1:length(x)) @@ -26,13 +29,13 @@ function ForwardColorJacCache(f,x,_chunksize = nothing; chunksize = _chunksize end - t = zeros(Dual{typeof(f), eltype(x), maximum(color)},length(x)) + t = zeros(Dual{typeof(f), eltype(x), getsize(chunksize)},length(x)) if dx === nothing fx = similar(t) _dx = similar(x) else - fx = zeros(Dual{typeof(f), eltype(dx), maximum(color)},length(dx)) + fx = zeros(Dual{typeof(f), eltype(dx), getsize(chunksize)},length(dx)) _dx = dx end From b605c549ad2714a507c690966e92fc1f51360220 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 16:01:03 -0400 Subject: [PATCH 3/6] Update src/compute_jacobian_ad.jl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Mathieu Besançon --- src/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index af335b53..7f113def 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -66,7 +66,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}; dx = nothing, - color = 1:length(x)) + color = eachindex(x)) forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color)) end From 4898f9afe4412e2b00158dae613d7ccdbae8d83f Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 21:57:00 -0400 Subject: [PATCH 4/6] Update src/compute_jacobian_ad.jl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Mathieu Besançon --- src/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index 7f113def..e25cbe05 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -49,7 +49,7 @@ function generate_chunked_partials(x,color,::Val{N}) where N # TODO: should only go up to the chunksize each time, and should # generate p[i] different parts, each with less than the chunksize - partials_array = Array{Bool}(undef, length(x), maximum(color)) + partials_array = BitMatrix{Bool}(undef, length(x), maximum(color)) for color_i in 1:maximum(color) for i in 1:length(x) if color[i]==color_i From 7d2916ab3f4c3c099d82b120ab314cebb610b872 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 21:57:07 -0400 Subject: [PATCH 5/6] Update src/compute_jacobian_ad.jl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Mathieu Besançon --- src/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index e25cbe05..5f1e261e 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -51,7 +51,7 @@ function generate_chunked_partials(x,color,::Val{N}) where N partials_array = BitMatrix{Bool}(undef, length(x), maximum(color)) for color_i in 1:maximum(color) - for i in 1:length(x) + for i in eachindex(x) if color[i]==color_i partials_array[i,color_i] = true else From 14a0f9e921d075a4e425774473a904c1efff00d1 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 19 Jun 2019 22:06:05 -0400 Subject: [PATCH 6/6] Update src/compute_jacobian_ad.jl --- src/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute_jacobian_ad.jl b/src/compute_jacobian_ad.jl index 5f1e261e..fde48049 100644 --- a/src/compute_jacobian_ad.jl +++ b/src/compute_jacobian_ad.jl @@ -49,7 +49,7 @@ function generate_chunked_partials(x,color,::Val{N}) where N # TODO: should only go up to the chunksize each time, and should # generate p[i] different parts, each with less than the chunksize - partials_array = BitMatrix{Bool}(undef, length(x), maximum(color)) + partials_array = BitMatrix(undef, length(x), maximum(color)) for color_i in 1:maximum(color) for i in eachindex(x) if color[i]==color_i