diff --git a/docs/src/tutorials/auto_parallel.md b/docs/src/tutorials/auto_parallel.md index 7af8c840cd..780375369a 100644 --- a/docs/src/tutorials/auto_parallel.md +++ b/docs/src/tutorials/auto_parallel.md @@ -26,7 +26,7 @@ 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 @@ -34,19 +34,17 @@ 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 ``` @@ -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! @@ -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! diff --git a/src/direct.jl b/src/direct.jl index 0cd43350fd..7d5a31dcc0 100644 --- a/src/direct.jl +++ b/src/direct.jl @@ -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 diff --git a/test/bigsystem.jl b/test/bigsystem.jl index c94f805c9a..09610a5c18 100644 --- a/test/bigsystem.jl +++ b/test/bigsystem.jl @@ -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])