From e104cfad94534441ba9c2778e8a486a3b46c1d2d Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Wed, 8 Dec 2021 12:28:36 -0500 Subject: [PATCH 1/2] Remove inline and specialize the function evaluations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I believe the reason for the inlining was because the functions weren't specializing, and so this will fix that. Improves this ODE example a ton: ```julia function rober2(u,p,t) y₁,y₂,y₃ = u k₁,k₂,k₃ = p du1 = -k₁*y₁+k₃*y₂*y₃ du2 = k₁*y₁-k₂*y₂^2-k₃*y₂*y₃ du3 = k₂*y₂^2 SA[du1,du2,du3] end prob2 = ODEProblem{false}(rober2,SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4]) @btime sol = solve(prob2,Rosenbrock23(),save_everystep=false) ``` --- src/differentiation/compute_jacobian_ad.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 9d9b0d99..a2c7d08d 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -79,13 +79,13 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize chunked_partials end -@inline function forwarddiff_color_jacobian(f, +function forwarddiff_color_jacobian(f::F, x::AbstractArray{<:Number}; colorvec = 1:length(x), sparsity = nothing, jac_prototype = nothing, chunksize = nothing, - dx = sparsity === nothing && jac_prototype === nothing ? nothing : copy(x)) #if dx is nothing, we will estimate dx at the cost of a function call + dx = sparsity === nothing && jac_prototype === nothing ? nothing : copy(x)) where {F} #if dx is nothing, we will estimate dx at the cost of a function call if sparsity === nothing && jac_prototype === nothing cfg = if chunksize === nothing @@ -95,7 +95,7 @@ end ForwardDiff.JacobianConfig(f, x) end else - ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize))) + ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk{getsize(chunksize)}()) end return ForwardDiff.jacobian(f, x, cfg) end @@ -105,13 +105,13 @@ end return forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype) end -@inline function forwarddiff_color_jacobian(J::AbstractArray{<:Number}, f, +function forwarddiff_color_jacobian(J::AbstractArray{<:Number}, f::F, x::AbstractArray{<:Number}; colorvec = 1:length(x), sparsity = nothing, jac_prototype = nothing, chunksize = nothing, - dx = similar(x, size(J, 1))) #dx kwarg can be used to avoid re-allocating dx every time + dx = similar(x, size(J, 1))) where {F} #dx kwarg can be used to avoid re-allocating dx every time if sparsity === nothing && jac_prototype === nothing cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize))) return ForwardDiff.jacobian(f, x, cfg) @@ -119,7 +119,7 @@ end return forwarddiff_color_jacobian(J,f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity)) end -function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing) +function forwarddiff_color_jacobian(f::F,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing) where F if jac_prototype isa Nothing ? ArrayInterface.ismutable(x) : ArrayInterface.ismutable(jac_prototype) # Whenever J is mutable, we mutate it to avoid allocations @@ -136,7 +136,7 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw end # When J is mutable, this version of forwarddiff_color_jacobian will mutate J to avoid allocations -function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache) +function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f::F,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache) where F t = jac_cache.t dx = jac_cache.dx p = jac_cache.p From 95a104f25b1263f6ed64904950c84100fa5704a7 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 8 Dec 2021 12:54:39 -0500 Subject: [PATCH 2/2] Update test_acyclic.jl --- test/test_acyclic.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_acyclic.jl b/test/test_acyclic.jl index 7ae98622..1bb0fc60 100644 --- a/test/test_acyclic.jl +++ b/test/test_acyclic.jl @@ -4,7 +4,7 @@ using Graphs: SimpleGraph using Test using Random -Random.seed!(100) +Random.seed!(90) # println("Starting acyclic coloring test...") #= Test data =#