Skip to content

Commit

Permalink
make the chunksize determination static when static
Browse files Browse the repository at this point in the history
  • Loading branch information
ChrisRackauckas committed Dec 11, 2021
1 parent 4077245 commit 0fc3154
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion src/alg_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -171,10 +171,22 @@ function DiffEqBase.prepare_alg(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorit
if typeof(alg) <: OrdinaryDiffEqImplicitExtrapolationAlgorithm
return alg # remake fails, should get fixed
else
remake(alg,chunk_size=Val{ForwardDiff.pickchunksize(x)}())
L = ArrayInterface.known_length(typeof(u0))
if L === nothing # dynamic sized
cs = ForwardDiff.pickchunksize(x)
remake(alg,chunk_size=cs)
else # statically sized
cs = pick_static_chunksize(Val{L}())
remake(alg,chunk_size=cs)
end
end
end

@generated function pick_static_chunksize(::Val{chunksize}) where chunksize
x = ForwardDiff.pickchunksize(chunksize)
:(Val{$x}())
end

function DiffEqBase.prepare_alg(alg::CompositeAlgorithm,u0,p,prob)
algs = map(alg -> DiffEqBase.prepare_alg(alg, u0, p, prob), alg.algs)
CompositeAlgorithm(algs, alg.choice_function)
Expand Down

0 comments on commit 0fc3154

Please sign in to comment.