Skip to content
This repository was archived by the owner on Aug 22, 2025. It is now read-only.

Commit dc0900c

Browse files
committed
Better inference
Before: ```julia 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 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) ```
1 parent f10cda4 commit dc0900c

File tree

1 file changed

+7
-3
lines changed

1 file changed

+7
-3
lines changed

src/differentiation/compute_jacobian_ad.jl

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,9 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing;
5252
end
5353

5454
generate_chunked_partials(x,colorvec,N::Integer) = generate_chunked_partials(x,colorvec,Val(N))
55-
function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize
55+
function generate_chunked_partials(x,colorvec,cs::Val{chunksize}) where chunksize
5656
maxcolor = maximum(colorvec)
57-
num_of_chunks = Int(ceil(maxcolor / chunksize))
57+
num_of_chunks = cld(maxcolor, chunksize)
5858
padding_size = (chunksize - (maxcolor % chunksize)) % chunksize
5959

6060
# partials = colorvec .== (1:maxcolor)'
@@ -72,13 +72,17 @@ function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize
7272
for i in 1:num_of_chunks
7373
tmp = Vector{NTuple{chunksize,eltype(x)}}(undef, size(partials,1))
7474
for j in 1:size(partials,1)
75-
tmp[j] = Tuple(@view partials[j,(i-1)*chunksize+1:i*chunksize])
75+
tmp[j] = partials_view_tup(partials, j, i, cs)
7676
end
7777
chunked_partials[i] = tmp
7878
end
7979
chunked_partials
8080
end
8181

82+
@generated function partials_view_tup(partials, j, i, ::Val{chunksize}) where chunksize
83+
:(Base.@ntuple $chunksize k->partials[j,(i-1)*$chunksize+k])
84+
end
85+
8286
function forwarddiff_color_jacobian(f::F,
8387
x::AbstractArray{<:Number};
8488
colorvec = 1:length(x),

0 commit comments

Comments
 (0)