Skip to content

fix automated parallelism with sparsejacobian #516

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.

Sign up for GitHub

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 2 commits into from
Jul 13, 2020
Merged

fix automated parallelism with sparsejacobian #516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c00be38
fix automated parallelism with sparsejacobian
ChrisRackauckas Jul 13, 2020
6db17ce
test sparsejacobian
ChrisRackauckas Jul 13, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 32 additions & 24 deletions docs/src/tutorials/auto_parallel.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,27 +26,25 @@ const X = reshape([i for i in 1:N for j in 1:N],N,N)
const Y = reshape([j for i in 1:N for j in 1:N],N,N)
const α₁ = 1.0.*(X.>=4*N/5)

const Mx = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
const Mx = Array(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]))
const My = copy(Mx)
Mx[2,1] = 2.0
Mx[end-1,end] = 2.0
My[1,2] = 2.0
My[end,end-1] = 2.0

# Define the discretized PDE as an ODE function
function f!(du,u,p,t)
A = @view u[:,:,1]
B = @view u[:,:,2]
C = @view u[:,:,3]
dA = @view du[:,:,1]
dB = @view du[:,:,2]
dC = @view du[:,:,3]
mul!(MyA,My,A)
mul!(AMx,A,Mx)
@. DA = _DD*(MyA + AMx)
@. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
@. dB = α₂ - β₂*B - r₁*A*B + r₂*C
@. dC = α₃ - β₃*C + r₁*A*B - r₂*C
function f(u,p,t)
A = u[:,:,1]
B = u[:,:,2]
C = u[:,:,3]
MyA = My*A
AMx = A*Mx
DA = @. _DD*(MyA + AMx)
dA = @. DA + α₁ - β₁*A - r₁*A*B + r₂*C
dB = @. α₂ - β₂*B - r₁*A*B + r₂*C
dC = @. α₃ - β₃*C + r₁*A*B - r₂*C
cat(dA,dB,dC,dims=3)
end
```

Expand All @@ -55,18 +53,25 @@ model function:

```julia
# Define the initial condition as normal arrays
@variables du[1:N,1:N,1:3] u[1:N,1:N,1:3] MyA[1:N,1:N] AMx[1:N,1:N] DA[1:N,1:N]
f!(du,u,nothing,0.0)
@variables u[1:N,1:N,1:3]
du = simplify.(f(u,nothing,0.0))
```

The output, here the in-place modified `du`, is a symbolic representation of
each output of the function. We can then utilize this in the ModelingToolkit
functionality. For example, let's compute the sparse Jacobian function and
compile a fast multithreaded version:
functionality. For example, let's build a parallel version of `f` first:

```julia
jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u)))
fjac = eval(ModelingToolkit.build_function(jac,u,parallel=ModelingToolkit.MultithreadedForm())[2])
fastf = eval(ModelingToolkit.build_function(du,u,
parallel=ModelingToolkit.MultithreadedForm())[2])
```

Now let's compute the sparse Jacobian function and compile a fast multithreaded version:

```julia
jac = ModelingToolkit.sparsejacobian(vec(du),vec(u))
fjac = eval(ModelingToolkit.build_function(jac,u,
parallel=ModelingToolkit.MultithreadedForm())[2])
```

It takes awhile for this to generate, but the results will be worth it!
Expand All @@ -81,16 +86,19 @@ MyA = zeros(N,N);
AMx = zeros(N,N);
DA = zeros(N,N);
prob = ODEProblem(f!,u0,(0.0,10.0))
prob_jac = ODEProblem(ODEFunction(f!,jac = (du,u,p,t) -> fjac(du,u), jac_prototype = similar(jac,Float64)),u0,(0.0,10.0))
fastprob = ODEProblem(ODEFunction((du,u,p,t)->fastf(du,u),
jac = (du,u,p,t) -> fjac(du,u),
jac_prototype = similar(jac,Float64)),
u0,(0.0,10.0))
```

Let's see the timing difference:

```julia
using BenchmarkTools
@btime solve(prob, TRBDF2(autodiff=false)) # 5.995 s (85786 allocations: 149.26 MiB)
@btime solve(prob_jac, TRBDF2()) # 250.409 ms (7411 allocations: 97.73 MiB)
@btime solve(prob, TRBDF2()) # 33.073 s (895404 allocations: 23.87 GiB)
@btime solve(fastprob, TRBDF2()) # 209.670 ms (8208 allocations: 109.25 MiB)
```

Boom, and automatic 24x acceleration that grows as the size of the problem
Boom, an automatic 157x acceleration that grows as the size of the problem
increases!
2 changes: 2 additions & 0 deletions src/direct.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,3 +208,5 @@ function simplified_expr(eq::Equation)
end

simplified_expr(eq::AbstractArray) = simplified_expr.(eq)
simplified_expr(x::Integer) = x
simplified_expr(x::AbstractFloat) = x
2 changes: 1 addition & 1 deletion test/bigsystem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit
using Dagger
daggerf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DaggerForm())[2])

jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
jac = ModelingToolkit.sparsejacobian(vec(du),vec(u))
serialjac = eval(ModelingToolkit.build_function(vec(jac),u)[2])
multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.MultithreadedForm())[2])
distributedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DistributedForm())[2])
Expand Down