Big slowdown and type instabilities of AD through ODE solver

[hersle]

Example:

using OrdinaryDiffEq, ProfileCanvas, ForwardDiff
function rober!(du, u, p, t)
    y₁, y₂, y₃ = u
    k₁, k₂, k₃ = p
    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
    du[2] = k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃
    du[3] = k₂ * y₂^2
    nothing
end
prob = ODEProblem(rober!, [1.0, 0.0, 0.0], (0.0, 1.0), [0.04, 3e7, 1e4])
function y_final(y_initial)
    newprob = remake(prob, u0 = y_initial)
    sol = solve(newprob, KenCarp4(); adaptive = false, dt = 1e-4, save_everystep = false)
    return sol[end]
end
@profview ForwardDiff.jacobian(y_final, [1.0, 0.0, 0.0])

First run the example with new package versions:

(bug) pkg> add OrdinaryDiffEq@6.101.0 LinearSolve@3.28.0 ForwardDiff ProfileCanvas

(bug) pkg> status
Status `~/Downloads/bug/Project.toml`
⌅ [f6369f11] ForwardDiff v0.10.38
  [7ed4a6bd] LinearSolve v3.28.0
  [1dea7af3] OrdinaryDiffEq v6.101.0
  [efd6af41] ProfileCanvas v0.1.6

Then run the same example with older package versions (at least before #2567?):

(bug) pkg> add OrdinaryDiffEq@6.93.0 LinearSolve@3.17.0 ForwardDiff ProfileCanvas

(bug) pkg> status
Status `~/Downloads/bug/Project.toml`
⌅ [f6369f11] ForwardDiff v0.10.38
⌃ [7ed4a6bd] LinearSolve v3.17.0
⌃ [1dea7af3] OrdinaryDiffEq v6.93.0
  [efd6af41] ProfileCanvas v0.1.6

There are a lot of type instabilities with the latest package versions!

[hersle]

Timing with new versions:

julia> @time ForwardDiff.jacobian(y_final, [1.0, 0.0, 0.0])
  0.738241 seconds (11.90 M allocations: 544.145 MiB, 7.43% gc time)
3×3 Matrix{Float64}:
 0.968768     0.259881     0.259953
 1.33748e-5  -9.96571e-5  -9.96455e-5
 0.0312191    0.740219     0.740146

Timing with old versions:

julia> @time ForwardDiff.jacobian(y_final, [1.0, 0.0, 0.0])
  0.035497 seconds (250.17 k allocations: 25.190 MiB, 31.49% gc time)
3×3 Matrix{Float64}:
 0.968768     0.259881     0.259953
 1.33748e-5  -9.96571e-5  -9.96455e-5
 0.0312191    0.740219     0.740146

With the latest versions, the examples run 20x slower!

[hersle]

I think the example runs fast also for OrdinaryDiffEq@6.101.0 (latest) and LinearSolve@3.19.2 (older version), but seems to slow down with LinearSolve@3.20.0, which introduced SciML/LinearSolve.jl#629.

[jClugstor]

I think I know where the type instabilities are coming from. I'll see if I can fix this up.

[hersle]

Thanks! With LinearSolve@3.30 this example has improved, but it is still 5x slower than with 3.19.2:

julia> @time ForwardDiff.jacobian(y_final, [1.0, 0.0, 0.0])
  0.188933 seconds (4.90 M allocations: 204.536 MiB, 14.45% gc time)

Now the red type-instability "# 6" points to LinearSolveForwardDiffExt.jl:119.

[ChrisRackauckas]

The partials size needs to be inferred. @jClugstor can you enforce that the partials use static arrays? That would solve it. Or pass them around as tuples instead of arrays, and keep the size information as Val type to use to reconstruct

[jClugstor]

I think this should help: SciML/LinearSolve.jl#727

I already use the Dual type to construct the Duals, and that includes the size of the partials so I can just reuse that. Then changing the tuple construction to NTuple with the same size as the Partials helps a lot too since the construction of the tuple was taking a lot of time. It's still quite a bit slower because of the need to create all of the arrays.