add and test automatic multithreading

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -5,7 +5,7 @@ using StaticArrays, LinearAlgebra, SparseArrays
 using Latexify, Unitful
 using MacroTools
 
-using MacroTools
+using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines
 import GeneralizedGenerated
 using DocStringExtensions

src/build_function.jl

@@ -47,7 +47,7 @@ end
 function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
                          conv = simplified_expr, expression = Val{true};
                          checkbounds = false, constructor=nothing,
-                         linenumbers = true)
+                         linenumbers = true, multithread=true)
     _vs = map(x-> x isa Operation ? x.op : x, vs)
     _ps = map(x-> x isa Operation ? x.op : x, ps)
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
@@ -67,6 +67,21 @@ function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
 
     ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))
 
+    if multithread
+        lens = Int(ceil(length(ip_let_expr.args[2].args)/Threads.nthreads()))
+        threaded_exprs = vcat([quote
+           Threads.@spawn begin
+              $(ip_let_expr.args[2].args[((i-1)*lens+1):i*lens]...)
+           end
+        end for i in 1:Threads.nthreads()-1],
+           quote
+              Threads.@spawn begin
+                 $(ip_let_expr.args[2].args[((Threads.nthreads()-1)*lens+1):end]...)
+              end
+           end)
+        ip_let_expr.args[2] =  ModelingToolkit.build_expr(:block, threaded_exprs)
+    end
+
     tuple_sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
 
     if rhss isa Matrix

test/bigsystem.jl

@@ -0,0 +1,50 @@
+using ModelingToolkit, LinearAlgebra
+
+# Define the constants for the PDE
+const α₂ = 1.0
+const α₃ = 1.0
+const β₁ = 1.0
+const β₂ = 1.0
+const β₃ = 1.0
+const r₁ = 1.0
+const r₂ = 1.0
+const D = 100.0
+const γ₁ = 0.1
+const γ₂ = 0.1
+const γ₃ = 0.1
+const N = 8
+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 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 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]
+
+# 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 = D*(MyA + AMx)
+  @. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
+  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
+  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
+end
+
+f(du,u,nothing,0.0)
+
+ModelingToolkit.build_function(du,u,multithread=true);
+jac = ModelingToolkit.jacobian(vec(du),vec(u),simplify=false)
+ModelingToolkit.build_function(vec(jac),u,multithread=true);

test/runtests.jl

@@ -15,4 +15,5 @@ using ModelingToolkit, Test
 @testset "Constraints Test" begin include("constraints.jl") end
 @testset "PDE Construction Test" begin include("pde.jl") end
 @testset "Distributed Test" begin include("distributed.jl") end
+@testset "Test Big System Usage" begin include("bigsystem.jl") end
 #@testset "Latexify recipes Test" begin include("latexify.jl") end