Merge pull request #280 from SciML/threads

Add and test automatic multithreading

Author: ChrisRackauckas

Project.toml

@@ -13,6 +13,7 @@ Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

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=false)
     _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,57 @@
+using ModelingToolkit, LinearAlgebra, SparseArrays
+
+# 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 _DD = 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 = _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
+end
+
+f(du,u,nothing,0.0)
+
+multithreadedf = eval(ModelingToolkit.build_function(du,u,multithread=true)[2])
+_du = rand(N,N,3)
+_u = rand(N,N,3)
+multithreadedf(_du,_u)
+
+jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
+multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,multithread=true)[2])
+
+#_jac = similar(jac,Float64)
+#multithreadedjac(_jac,_u)

test/nonlinearsystem.jl

@@ -52,6 +52,7 @@ jac_func = generate_jacobian(ns)
 f = @eval eval(nlsys_func)
 
 # Intermediate calculations
+a = y - x
 # Define a nonlinear system
 eqs = [0 ~ σ*a,
        0 ~ x*(ρ-z)-y,

test/runtests.jl

@@ -1,18 +1,19 @@
-using ModelingToolkit, Test
+using SafeTestsets, Test
 
-@testset "Parsing Test" begin include("variable_parsing.jl") end
-@testset "Differentiation Test" begin include("derivatives.jl") end
-@testset "Simplify Test" begin include("simplify.jl") end
-@testset "Operation Overloads Test" begin include("operation_overloads.jl") end
-@testset "Direct Usage Test" begin include("direct.jl") end
-@testset "ODESystem Test" begin include("odesystem.jl") end
-@testset "Mass Matrix Test" begin include("mass_matrix.jl") end
-@testset "SDESystem Test" begin include("sdesystem.jl") end
-@testset "NonlinearSystem Test" begin include("nonlinearsystem.jl") end
-@testset "OptimizationSystem Test" begin include("optimizationsystem.jl") end
-@testset "Build Targets Test" begin include("build_targets.jl") end
-@testset "Domain Test" begin include("domains.jl") end
-@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
+@safetestset "Parsing Test" begin include("variable_parsing.jl") end
+@safetestset "Differentiation Test" begin include("derivatives.jl") end
+@safetestset "Simplify Test" begin include("simplify.jl") end
+@safetestset "Operation Overloads Test" begin include("operation_overloads.jl") end
+@safetestset "Direct Usage Test" begin include("direct.jl") end
+@safetestset "ODESystem Test" begin include("odesystem.jl") end
+@safetestset "Mass Matrix Test" begin include("mass_matrix.jl") end
+@safetestset "SDESystem Test" begin include("sdesystem.jl") end
+@safetestset "NonlinearSystem Test" begin include("nonlinearsystem.jl") end
+@safetestset "OptimizationSystem Test" begin include("optimizationsystem.jl") end
+@safetestset "Build Targets Test" begin include("build_targets.jl") end
+@safetestset "Domain Test" begin include("domains.jl") end
+@safetestset "Constraints Test" begin include("constraints.jl") end
+@safetestset "PDE Construction Test" begin include("pde.jl") end
+@safetestset "Test Big System Usage" begin include("bigsystem.jl") end
 #@testset "Latexify recipes Test" begin include("latexify.jl") end
+@testset "Distributed Test" begin include("distributed.jl") end