JuliaDiff / SparseDiffTools.jl

Fast jacobian computation through sparsity exploitation and matrix coloring
MIT License
237 stars 41 forks source link

Better inference #169

Closed YingboMa closed 2 years ago

YingboMa commented 2 years ago

Before:

julia> itrigs = inference_triggers(tinf)
6-element Vector{InferenceTrigger}:
 Inference triggered to call CommonSolve.solve(::ODEProblem{Vector{Fl, ::Rodas5{3, true, Defa) from eval (./boot.jl:373) inlined into REPL.eval_user_input(::Any, ::REPL.REPLBackend) (/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.7/REPL/src/REPL.jl:150)
 Inference triggered to call OrdinaryDiffEq.build_J_W(::Rodas5{3, true, Defa, ::Vector{Float64}, ::Vector{Float64}, ::SciMLBase.NullParame, ::Float64, ::Float64, ::ODEFunction{true, ty, ::Type, ::Val{true}) from alg_cache (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/caches/rosenbrock_caches.jl:581) with specialization OrdinaryDiffEq.alg_cache(::Rodas5{3, true, Defa, ::Vector{Float64}, ::Vector{Float64}, ::Type{Float64}, ::Type{Float64}, ::Type{Float64}, ::Vector{Float64}, ::Vector{Float64}, ::ODEFunction{true, ty, ::Float64, ::Float64, ::Float64, ::SciMLBase.NullParame, ::Bool, ::Val{true})
 Inference triggered to call setindex!(::Vector{Tuple{Float64, ::Tuple{Bool, Bool, Bo, ::Int64) from generate_chunked_partials (/Users/scheme/src/julia/SparseDiffTools/src/differentiation/compute_jacobian_ad.jl:75) with specialization SparseDiffTools.generate_chunked_partials(::Vector{Float64}, ::UnitRange{Int64}, ::Val{3})
 Inference triggered to call unsafe_copyto!(::Matrix{Float64}, ::Int64, ::Matrix{Float64}, ::Int64, ::Int64) from _copyto_impl! (./array.jl:331) inlined into OrdinaryDiffEq.calc_W!(::Matrix{Float64}, ::OrdinaryDiffEq.ODEIn, ::Nothing, ::OrdinaryDiffEq.Rosen, ::Float64, ::Bool, ::Bool) (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/derivative_utils.jl:503)
 Inference triggered to call DiffEqBase.var"#_#31"(::Nothing, ::Base.Pairs{Symbol, U, ::DefaultLinSolve, ::Vector{Float64}, ::Matrix{Float64}, ::Vector{Float64}, ::Bool) from DefaultLinSolve (/Users/scheme/.julia/packages/DiffEqBase/b1nST/src/linear_nonlinear.jl:94) inlined into OrdinaryDiffEq.perform_step!(::OrdinaryDiffEq.ODEIn, ::OrdinaryDiffEq.Rosen, ::Bool) (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/perform_step/rosenbrock_perform_step.jl:1577)
 Inference triggered to call LinearAlgebra.var"#generic_lufact!#152"(::Bool, ::typeof(LinearAlgebra, ::Matrix{Float64}, ::LinearAlgebra.RowMax) from generic_lufact!##kw (/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.7/LinearAlgebra/src/lu.jl:139) inlined into DiffEqBase.var"#_#31"(::Nothing, ::Base.Pairs{Symbol, U, ::DefaultLinSolve, ::Vector{Float64}, ::Matrix{Float64}, ::Vector{Float64}, ::Bool) (/Users/scheme/.julia/packages/DiffEqBase/b1nST/src/linear_nonlinear.jl:112)

After:

julia> itrigs = inference_triggers(tinf)
5-element Vector{InferenceTrigger}:
 Inference triggered to call CommonSolve.solve(::ODEProblem{Vector{Fl, ::Rodas5{3, true, Defa) from eval (./boot.jl:373) inlined into REPL.eval_user_input(::Any, ::REPL.REPLBackend) (/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.7/REPL/src/REPL.jl:150)
 Inference triggered to call OrdinaryDiffEq.build_J_W(::Rodas5{3, true, Defa, ::Vector{Float64}, ::Vector{Float64}, ::SciMLBase.NullParame, ::Float64, ::Float64, ::ODEFunction{true, ty, ::Type, ::Val{true}) from alg_cache (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/caches/rosenbrock_caches.jl:581) with specialization OrdinaryDiffEq.alg_cache(::Rodas5{3, true, Defa, ::Vector{Float64}, ::Vector{Float64}, ::Type{Float64}, ::Type{Float64}, ::Type{Float64}, ::Vector{Float64}, ::Vector{Float64}, ::ODEFunction{true, ty, ::Float64, ::Float64, ::Float64, ::SciMLBase.NullParame, ::Bool, ::Val{true})
 Inference triggered to call unsafe_copyto!(::Matrix{Float64}, ::Int64, ::Matrix{Float64}, ::Int64, ::Int64) from _copyto_impl! (./array.jl:331) inlined into OrdinaryDiffEq.calc_W!(::Matrix{Float64}, ::OrdinaryDiffEq.ODEIn, ::Nothing, ::OrdinaryDiffEq.Rosen, ::Float64, ::Bool, ::Bool) (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/derivative_utils.jl:503)
 Inference triggered to call DiffEqBase.var"#_#31"(::Nothing, ::Base.Pairs{Symbol, U, ::DefaultLinSolve, ::Vector{Float64}, ::Matrix{Float64}, ::Vector{Float64}, ::Bool) from DefaultLinSolve (/Users/scheme/.julia/packages/DiffEqBase/b1nST/src/linear_nonlinear.jl:94) inlined into OrdinaryDiffEq.perform_step!(::OrdinaryDiffEq.ODEIn, ::OrdinaryDiffEq.Rosen, ::Bool) (/Users/scheme/.julia/packages/OrdinaryDiffEq/6gTvi/src/perform_step/rosenbrock_perform_step.jl:1577)
 Inference triggered to call LinearAlgebra.var"#generic_lufact!#152"(::Bool, ::typeof(LinearAlgebra, ::Matrix{Float64}, ::LinearAlgebra.RowMax) from generic_lufact!##kw (/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.7/LinearAlgebra/src/lu.jl:139) inlined into DiffEqBase.var"#_#31"(::Nothing, ::Base.Pairs{Symbol, U, ::DefaultLinSolve, ::Vector{Float64}, ::Matrix{Float64}, ::Vector{Float64}, ::Bool) (/Users/scheme/.julia/packages/DiffEqBase/b1nST/src/linear_nonlinear.jl:112)
ChrisRackauckas commented 2 years ago

And with all of the latest updates in SciMLBase and OrdinaryDiffEq, this is down to:

julia> itrigs = inference_triggers(tinf)
1-element Vector{InferenceTrigger}:
 Inference triggered to call OrdinaryDiffEq.jacobian2W!(::Matrix{Float64}, ::LinearAlgebra.UniformScaling{Bool}, ::Float64, ::Matrix{Float64}, ::Bool) called from toplevel

Literally one inference trigger, and it's a snoopcompile bug. Awesome.