Merge pull request #798 from SciML/s/threading

Bring back multithreading

Author: shashi

Project.toml

@@ -67,7 +67,6 @@ Unitful = "1.1"
 julia = "1.2"
 
 [extras]
-Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
@@ -79,4 +78,4 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Dagger", "ForwardDiff", "GalacticOptim", "NonlinearSolve", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]
+test = ["ForwardDiff", "GalacticOptim", "NonlinearSolve", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]

docs/src/comparison.md

@@ -24,9 +24,7 @@ excels in many areas due to purposeful design decisions:
 - Parallelism: ModelingToolkit.jl has pervasive parallelism. The
   symbolic simplification via [SymbolicUtils.jl](https://github.com/JuliaSymbolics/SymbolicUtils.jl)
   has built-in parallelism, ModelingToolkit.jl builds functions that
-  parallelizes across threads and multiprocesses across clusters,
-  and it has dynamic scheduling through tools like [Dagger.jl](https://github.com/JuliaParallel/Dagger.jl).
-  ModelingToolkit.jl is compatible with GPU libraries like CUDA.jl.
+  parallelizes across threads. ModelingToolkit.jl is compatible with GPU libraries like CUDA.jl.
 - Scientific Machine Learning (SciML): ModelingToolkit.jl is made to synergize
   with the high performance Julia SciML ecosystem in many ways. At a
   base level, all expressions and built functions are compatible with

src/ModelingToolkit.jl

@@ -296,9 +296,4 @@ export @register
 export modelingtoolkitize
 export @variables, @parameters
 
-const HAS_DAGGER = Ref{Bool}(false)
-function __init__()
-    @require Dagger="d58978e5-989f-55fb-8d15-ea34adc7bf54" include("dagger.jl")
-end
-
 end # module

src/build_function.jl

@@ -1,4 +1,5 @@
 using SymbolicUtils.Code
+using Base.Threads
 
 abstract type BuildTargets end
 struct JuliaTarget <: BuildTargets end
@@ -8,9 +9,10 @@ struct MATLABTarget <: BuildTargets end
 
 abstract type ParallelForm end
 struct SerialForm <: ParallelForm end
-struct MultithreadedForm <: ParallelForm end
-struct DistributedForm <: ParallelForm end
-struct DaggerForm <: ParallelForm end
+struct MultithreadedForm <: ParallelForm
+    ntasks::Int
+end
+MultithreadedForm() = MultithreadedForm(2*nthreads())
 
 """
 `build_function`
@@ -60,18 +62,6 @@ function unflatten_long_ops(op, N=4)
     Rewriters.Fixpoint(Rewriters.Postwalk(Rewriters.Chain([rule1, rule2])))(op)
 end
 
-function observed_let(eqs)
-    process -> ex -> begin
-        isempty(eqs) && return ex
-
-        assignments = map(eq -> :($(process(eq.lhs)) = $(process(eq.rhs))), eqs)
-        letexpr = :(let $(assignments...)
-                    end)
-        # avoid a superfluous `begin ... end` block
-        letexpr.args[2] = ex
-        return letexpr
-    end
-end
 
 # Scalar output
 
@@ -140,7 +130,7 @@ Build function target: JuliaTarget
 function _build_function(target::JuliaTarget, rhss, args...;
                          conv = toexpr, expression = Val{true},
                          checkbounds = false,
-                         linenumbers = false, multithread=nothing,
+                         linenumbers = false,
                          headerfun = addheader, outputidxs=nothing,
                          convert_oop = true, force_SA = false,
                          skipzeros = outputidxs===nothing,
@@ -168,10 +158,6 @@ Special Keyword Argumnets:
   - `SerialForm()`: Serial execution.
   - `MultithreadedForm()`: Multithreaded execution with a static split, evenly
     splitting the number of expressions per thread.
-  - `DistributedForm()`: Multiprocessing using Julia's Distributed with a static
-    schedule, evenly splitting the number of expressions per process.
-  - `DaggerForm()`: Multithreading and multiprocessing using Julia's Dagger.jl
-    for dynamic scheduling and load balancing.
 - `conv`: The conversion function of the Operation to Expr. By default this uses
   the `toexpr` function.
 - `checkbounds`: For whether to enable bounds checking inside of the generated
@@ -194,7 +180,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
                        expression = Val{true},
                        expression_module = @__MODULE__(),
                        checkbounds = false,
-                       linenumbers = false, multithread=nothing,
+                       linenumbers = false,
                        outputidxs=nothing,
                        skipzeros = false,
                        wrap_code = (nothing, nothing),
@@ -204,13 +190,13 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     dargs = map(destructure_arg, [args...])
     i = findfirst(x->x isa DestructuredArgs, dargs)
     similarto = i === nothing ? Array : dargs[i].name
-    oop_expr = Func(dargs, [], _make_array(rhss, similarto))
+    oop_expr = Func(dargs, [], make_array(parallel, rhss, similarto))
     if !isnothing(wrap_code[1])
         oop_expr = wrap_code[1](oop_expr)
     end
 
     out = Sym{Any}(gensym("out"))
-    ip_expr = Func([out, dargs...], [], _set_array(out, outputidxs, rhss, checkbounds, skipzeros))
+    ip_expr = Func([out, dargs...], [], set_array(parallel, out, outputidxs, rhss, checkbounds, skipzeros))
 
     if !isnothing(wrap_code[2])
         ip_expr = wrap_code[2](ip_expr)
@@ -224,6 +210,24 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     end
 end
 
+function make_array(s, arr, similarto)
+    Base.@warn("Parallel form of $(typeof(s)) not implemented")
+    _make_array(arr, similarto)
+end
+
+function make_array(s::SerialForm, arr, similarto)
+    _make_array(arr, similarto)
+end
+
+function make_array(s::MultithreadedForm, arr, similarto)
+    per_task = ceil(Int, length(arr) / s.ntasks)
+    slices = collect(Iterators.partition(arr, per_task))
+    arrays = map(slices) do slice
+        _make_array(slice, similarto)
+    end
+    SpawnFetch{Multithreaded}(arrays, vcat)
+end
+
 function _make_array(rhss::AbstractSparseArray, similarto)
     arr = map(x->_make_array(x, similarto), rhss)
     if !(arr isa AbstractSparseArray)
@@ -233,19 +237,59 @@ function _make_array(rhss::AbstractSparseArray, similarto)
     end
 end
 
+function _make_array(rhss::AbstractArray, similarto)
+    arr = map(x->_make_array(x, similarto), rhss)
+    # Ugh reshaped array of a sparse array when mapped gives a sparse array
+    if arr isa AbstractSparseArray
+        _make_array(arr, similarto)
+    else
+        MakeArray(arr, similarto)
+    end
+end
+
+_make_array(x, similarto) = x
+
 ## In-place version
-function _set_array(out, outputidxs, rhss::AbstractArray, checkbounds, skipzeros)
-    if rhss isa Union{SparseVector, SparseMatrixCSC}
-        return SetArray(checkbounds, LiteralExpr(:($out.nzval)), rhss.nzval)
-    elseif outputidxs === nothing
+
+function set_array(p, args...)
+    Base.@warn("Parallel form of $(typeof(p)) not implemented")
+    _set_array(args...)
+end
+
+function set_array(s::SerialForm, args...)
+    _set_array(args...)
+end
+
+function set_array(s::MultithreadedForm, out, outputidxs, rhss, checkbounds, skipzeros)
+    if rhss isa AbstractSparseArray
+        return set_array(LiteralExpr(:($out.nzval)),
+                         nothing,
+                         rhss.nzval,
+                         checkbounds,
+                         skipzeros)
+    end
+    if outputidxs === nothing
         outputidxs = collect(eachindex(rhss))
     end
+    per_task = ceil(Int, length(rhss) / s.ntasks)
+    # TODO: do better partitioning when skipzeros is present
+    slices = collect(Iterators.partition(zip(outputidxs, rhss), per_task))
+    arrays = map(slices) do slice
+        idxs, vals = first.(slice), last.(slice)
+        _set_array(out, idxs, vals, checkbounds, skipzeros)
+    end
+    SpawnFetch{Multithreaded}(arrays, @inline noop(args...) = nothing)
+end
 
+function _set_array(out, outputidxs, rhss::AbstractArray, checkbounds, skipzeros)
+    if outputidxs === nothing
+        outputidxs = collect(eachindex(rhss))
+    end
     # sometimes outputidxs is a Tuple
     ii = findall(i->!(rhss[i] isa AbstractArray) && !(skipzeros && _iszero(rhss[i])), eachindex(outputidxs))
     jj = findall(i->rhss[i] isa AbstractArray, eachindex(outputidxs))
     exprs = []
-    push!(exprs, SetArray(checkbounds, out, AtIndex.(vec(collect(outputidxs[ii])), vec(rhss[ii]))))
+    push!(exprs, SetArray(!checkbounds, out, AtIndex.(vec(collect(outputidxs[ii])), vec(rhss[ii]))))
     for j in jj
         push!(exprs, _set_array(LiteralExpr(:($out[$j])), nothing, rhss[j], checkbounds, skipzeros))
     end
@@ -257,18 +301,6 @@ end
 _set_array(out, outputidxs, rhs, checkbounds, skipzeros) = rhs
 
 
-function _make_array(rhss::AbstractArray, similarto)
-    arr = map(x->_make_array(x, similarto), rhss)
-    # Ugh reshaped array of a sparse array when mapped gives a sparse array
-    if arr isa AbstractSparseArray
-        _make_array(arr, similarto)
-    else
-        MakeArray(arr, similarto)
-    end
-end
-
-_make_array(x, similarto) = x
-
 function vars_to_pairs(name,vs::Union{Tuple, AbstractArray}, symsdict=Dict())
     vs_names = tosymbol.(vs)
     for (v,k) in zip(vs_names, vs)

src/dagger.jl

@@ -80,18 +80,9 @@ FiniteDiff.finite_difference_jacobian!(J2,(du,u)->f!(du,u,nothing,nothing),u)
 maximum(J2 .- Array(J)) < 1e-5
 =#
 
-using Distributed
-addprocs(4)
-distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DistributedForm())[2])
-
-using Dagger
-daggerf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DaggerForm())[2])
-
 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])
-daggerjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DaggerForm())[2])
 
 MyA = zeros(N,N)
 AMx = zeros(N,N)
@@ -101,19 +92,13 @@ _u = rand(N,N,3)
 
 f(_du,_u,nothing,0.0)
 multithreadedf(_du,_u)
-#distributedf(_du,_u)
-#daggerf(_du,_u)
 
 #=
 using BenchmarkTools
 @btime f(_du,_u,nothing,0.0)
 @btime multithreadedf(_du,_u)
-@btime distributedf(_du,_u)
-@btime daggerf(_du,_u)
 
 _jac = similar(jac,Float64)
 @btime serialjac(_jac,_u)
 @btime multithreadedjac(_jac,_u)
-@btime distributedjac(_jac,_u)
-@btime daggerjac(_jac,_u)
 =#

test/bigsystem.jl

@@ -14,11 +14,15 @@ end
 
 h_str = ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g])
 h_oop = eval(h_str[1])
+h_str_par = ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g], parallel=ModelingToolkit.MultithreadedForm())
+h_oop_par = eval(h_str_par[1])
 h_ip! = eval(h_str[2])
 h_ip_skip! = eval(ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g], skipzeros=true, fillzeros=false)[2])
+h_ip_skip_par! = eval(ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g], skipzeros=true, parallel=ModelingToolkit.MultithreadedForm(), fillzeros=false)[2])
 inputs = ([1], [2], [3, 4, 5], [6], [7], [8])
 
 @test h_oop(inputs...) == h_julia(inputs...)
+@test h_oop_par(inputs...) == h_julia(inputs...)
 out_1 = similar(h, Int)
 out_2 = similar(out_1)
 h_ip!(out_1, inputs...)
@@ -30,6 +34,12 @@ h_ip_skip!(out_1, inputs...)
 out_1[3] = 0
 @test out_1 == out_2
 
+fill!(out_1, 10)
+h_ip_skip_par!(out_1, inputs...)
+@test out_1[3] == 10
+out_1[3] = 0
+@test out_1 == out_2
+
 # Multiple input matrix, some unused arguments
 h_skip = [a + b + c1; c2 + c3 + g] # skip d, e
 h_julia_skip(a, b, c, d, e, g) = [a[1] + b[1] + c[1]; c[2] + c[3] + g[1]]