Skip to content

Specialize default chunk size on static arrays #143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 6, 2021
Merged

Conversation

ChrisRackauckas
Copy link
Member

Since the size is static and known to be sufficiently small, we might as well directly specialize on it. This does so, and it reduces the allocations and noticeably improves the static array ODE solver speed.

using OrdinaryDiffEq, BenchmarkTools
using StaticArrays, LinearAlgebra

# rober
function rober(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
ff = ODEFunction(rober,tgrad = (u,p,t)->SA[0.0,0.0,0.0])
prob = ODEProblem{false}(rober,SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4])

# Before
@btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 122.500 μs (581 allocations: 53.39 KiB)
# After
@btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 101.600 μs (32 allocations: 10.50 KiB)

using ModelingToolkit
prob2 = ODEProblem{false}(modelingtoolkitize(prob),SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4],jac=true)
@btime solve(prob2,Rodas5(), reltol=1.0e-8, abstol=1.0e-8, saveat = 1000) # 75.000 μs (401 allocations: 51.38 KiB)

It's still better to modelingtoolkitize, but this is still a nice victory.

Since the size is static and known to be sufficiently small, we might as well directly specialize on it. This does so, and it reduces the allocations and noticeably improves the static array ODE solver speed.

```julia
using OrdinaryDiffEq, BenchmarkTools
using StaticArrays, LinearAlgebra

# rober
function rober(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
ff = ODEFunction(rober,tgrad = (u,p,t)->SA[0.0,0.0,0.0])
prob = ODEProblem{false}(rober,SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4])

# Before
@Btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 122.500 μs (581 allocations: 53.39 KiB)
# After
@Btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 101.600 μs (32 allocations: 10.50 KiB)

using ModelingToolkit
prob2 = ODEProblem{false}(modelingtoolkitize(prob),SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4],jac=true)
@Btime solve(prob2,Rodas5(), reltol=1.0e-8, abstol=1.0e-8, saveat = 1000) # 75.000 μs (401 allocations: 51.38 KiB)
```

It's still better to modelingtoolkitize, but this is still a nice victory.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant