Merge pull request #348 from SciML/distributed

Distributed parallelism build targets

Author: ChrisRackauckas

src/build_function.jl

@@ -4,6 +4,11 @@ struct StanTarget <: BuildTargets end
 struct CTarget <: BuildTargets end
 struct MATLABTarget <: BuildTargets end
 
+abstract type ParallelForm end
+struct SerialForm <: ParallelForm end
+struct MultithreadedForm <: ParallelForm end
+struct DistributedForm <: ParallelForm end
+
 """
 `build_function`
 
@@ -57,7 +62,7 @@ function build_function(args...;target = JuliaTarget(),kwargs...)
 end
 
 function addheader(ex, fargs, iip; X=gensym(:MTIIPVar))
-  if iip  
+  if iip
     wrappedex = :(
         ($X,$(fargs.args...)) -> begin
         $ex
@@ -69,28 +74,28 @@ function addheader(ex, fargs, iip; X=gensym(:MTIIPVar))
       ($(fargs.args...),) -> begin
         $ex
       end
-    )  
+    )
   end
   wrappedex
 end
 
 function add_integrator_header(ex, fargs, iip; X=gensym(:MTIIPVar))
   integrator = gensym(:MTKIntegrator)
-  if iip  
+  if iip
     wrappedex = :(
-        $integrator -> begin 
+        $integrator -> begin
         ($X,$(fargs.args...)) = (($integrator).u,($integrator).u,($integrator).p,($integrator).t)
         $ex
         nothing
       end
     )
   else
     wrappedex = :(
-      $integrator -> begin 
+      $integrator -> begin
       ($(fargs.args...),) = (($integrator).u,($integrator).p,($integrator).t)
         $ex
       end
-    )  
+    )
   end
   wrappedex
 end
@@ -114,7 +119,7 @@ function _build_function(target::JuliaTarget, op::Operation, args...;
 
     fargs = Expr(:tuple,argnames...)
     oop_ex = headerfun(bounds_block, fargs, false)
-        
+
     if !linenumbers
         oop_ex = striplines(oop_ex)
     end
@@ -129,8 +134,15 @@ end
 function _build_function(target::JuliaTarget, rhss, args...;
                          conv = simplified_expr, expression = Val{true},
                          checkbounds = false, constructor=nothing,
-                         linenumbers = false, multithread=false, 
-                         headerfun=addheader, outputidxs=nothing)
+                         linenumbers = false, multithread=nothing,
+                         headerfun=addheader, outputidxs=nothing,
+						 parallel=SerialForm())
+
+	if multithread isa Bool
+		@warn("multithraded is deprecated for the parallel argument. See the documentation.")
+		parallel = multithread ? MultithreadedForm() : SerialForm()
+	end
+
     argnames = [gensym(:MTKArg) for i in 1:length(args)]
     arg_pairs = map(vars_to_pairs,zip(argnames,args))
     ls = reduce(vcat,first.(arg_pairs))
@@ -143,23 +155,39 @@ function _build_function(target::JuliaTarget, rhss, args...;
 
     oidx = isnothing(outputidxs) ? (i -> i) : (i -> outputidxs[i])
     X = gensym(:MTIIPVar)
+
+	rhs_length = rhss isa SparseMatrixCSC ? length(rhss.nzval) : length(rhss)
+
+	if parallel isa DistributedForm
+		numworks = Distributed.nworkers()
+		reducevars = [Variable(gensym(:MTReduceVar))() for i in 1:numworks]
+		lens = Int(ceil(rhs_length/numworks))
+		finalsize = rhs_length - (numworks-1)*lens
+		_rhss = vcat(reduce(vcat,[[getindex(reducevars[i],j) for j in 1:lens] for i in 1:numworks-1],init=Expr[]),
+						 [getindex(reducevars[end],j) for j in 1:finalsize])
+	elseif rhss isa SparseMatrixCSC
+		_rhss = rhss.nzval
+	else
+		_rhss = rhss
+	end
+
 	if eltype(eltype(rhss)) <: AbstractArray # Array of arrays of arrays
-		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j][$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)],init=Expr[])) for (i,rhsel) ∈ enumerate(rhss)],init=Expr[])
+		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j][$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)],init=Expr[])) for (i,rhsel) ∈ enumerate(_rhss)],init=Expr[])
 	elseif eltype(eltype(rhss)) <: SparseMatrixCSC # Array of arrays of arrays
-		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j].nzval[$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)])) for (i,rhsel) ∈ enumerate(rhss)])
+		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j].nzval[$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)])) for (i,rhsel) ∈ enumerate(_rhss)])
 	elseif eltype(rhss) <: SparseMatrixCSC # Array of sparse matrices
-		ip_sys_exprs = reduce(vcat,[vec([:($X[$i].nzval[$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(rhss)])
+		ip_sys_exprs = reduce(vcat,[vec([:($X[$i].nzval[$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(_rhss)])
     elseif eltype(rhss) <: AbstractArray # Array of arrays
-		ip_sys_exprs = reduce(vcat,[vec([:($X[$i][$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(rhss)], init = Expr[])
+		ip_sys_exprs = reduce(vcat,[vec([:($X[$i][$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(_rhss)], init = Expr[])
     elseif rhss isa SparseMatrixCSC
-        ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss.nzval)]
+        ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(_rhss)]
     else
-        ip_sys_exprs = [:($X[$(oidx(i))] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
+        ip_sys_exprs = [:($X[$(oidx(i))] = $(conv(rhs))) for (i, rhs) ∈ enumerate(_rhss)]
     end
 
     ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))
 
-    if multithread
+    if parallel isa MultithreadedForm
         lens = Int(ceil(length(ip_let_expr.args[2].args)/Threads.nthreads()))
         threaded_exprs = vcat([quote
            Threads.@spawn begin
@@ -172,7 +200,29 @@ function _build_function(target::JuliaTarget, rhss, args...;
               end
            end)
         ip_let_expr.args[2] =  ModelingToolkit.build_expr(:block, threaded_exprs)
-    end
+    elseif parallel isa DistributedForm
+		numworks = Distributed.nworkers()
+		lens = Int(ceil(length(ip_let_expr.args[2].args)/numworks))
+		spawnvars = [gensym(:MTSpawnVar) for i in 1:numworks]
+		rhss_flat = rhss isa SparseMatrixCSC ? rhss.nzval : rhss
+		spawnvectors = vcat(
+					   [build_expr(:vect, [conv(rhs) for rhs ∈ rhss_flat[((i-1)*lens+1):i*lens]]) for i in 1:numworks-1],
+					   build_expr(:vect, [conv(rhs) for rhs ∈ rhss_flat[((numworks-1)*lens+1):end]]))
+
+        spawn_exprs = [quote
+           $(spawnvars[i]) = ModelingToolkit.Distributed.remotecall($(i+1)) do
+              $(spawnvectors[i])
+           end
+        end for i in 1:numworks]
+        spawn_exprs = ModelingToolkit.build_expr(:block, spawn_exprs)
+		resunpack_exprs = [:($(Symbol(reducevars[iter])) = fetch($(spawnvars[iter]))) for iter in 1:numworks]
+
+		ip_let_expr.args[2] = quote
+			$spawn_exprs
+			$(resunpack_exprs...)
+			$(ip_let_expr.args[2])
+		end
+	end
 
     tuple_sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
 
@@ -214,7 +264,7 @@ function _build_function(target::JuliaTarget, rhss, args...;
 
     oop_ex = headerfun(oop_body_block, fargs, false)
     iip_ex = headerfun(ip_bounds_block, fargs, true; X=X)
-    
+
     if !linenumbers
         oop_ex = striplines(oop_ex)
         iip_ex = striplines(iip_ex)

src/differentials.jl

@@ -36,8 +36,8 @@ $(SIGNATURES)
 
 TODO
 """
-function expand_derivatives(O::Operation)
-    @. O.args = expand_derivatives(O.args)
+function expand_derivatives(O::Operation,simplify=true)
+    @. O.args = expand_derivatives(O.args,simplify)
 
     if isa(O.op, Differential)
         (D, o) = (O.op, O.args[1])
@@ -47,14 +47,16 @@ function expand_derivatives(O::Operation)
         isa(o, Operation)   || return O
         isa(o.op, Variable) && return O
 
-        return sum(1:length(o.args)) do i
-            derivative(o, i) * expand_derivatives(D(o.args[i]))
-        end |> simplify
+        x = sum(1:length(o.args)) do i
+            derivative(o, i) * expand_derivatives(D(o.args[i]),simplify)
+        end
+
+        return simplify ? ModelingToolkit.simplify(x) : x
     end
 
-    return simplify(O)
+    return simplify ? ModelingToolkit.simplify(O) : O
 end
-expand_derivatives(x) = x
+expand_derivatives(x,args...) = x
 
 # Don't specialize on the function here
 """

src/direct.jl

@@ -7,8 +7,7 @@ A helper function for computing the gradient of an expression with respect to
 an array of variable expressions.
 """
 function gradient(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v)(O)) for v in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v)(O),simplify) for v in vars]
 end
 
 """
@@ -20,8 +19,7 @@ A helper function for computing the Jacobian of an array of expressions with res
 an array of variable expressions.
 """
 function jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v)(O)) for O in ops, v in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v)(O),simplify) for O in ops, v in vars]
 end
 
 """
@@ -33,8 +31,7 @@ A helper function for computing the Hessian of an expression with respect to
 an array of variable expressions.
 """
 function hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
-    out = [expand_derivatives(Differential(v2)(Differential(v1)(O))) for v1 in vars, v2 in vars]
-    simplify ? ModelingToolkit.simplify.(out) : out
+    [expand_derivatives(Differential(v2)(Differential(v1)(O)),simplify) for v1 in vars, v2 in vars]
 end
 
 function simplified_expr(O::Operation)

src/systems/diffeqs/abstractodesystem.jl

@@ -197,7 +197,7 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     checkbounds = false, sparse = false,
-                                    linenumbers = true, multithread=false,
+                                    linenumbers = true, parallel=SerialForm(),
                                     kwargs...) where iip
 ```
 
@@ -209,10 +209,10 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     checkbounds = false, sparse = false,
-                                    linenumbers = true, multithread=false,
+                                    linenumbers = true, parallel=SerialForm(),
                                     kwargs...) where iip
     f = ODEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
-                        linenumbers=linenumbers,multithread=multithread,
+                        linenumbers=linenumbers,parallel=parallel,
                         sparse=sparse)
     u0 = varmap_to_vars(u0map,states(sys))
     p = varmap_to_vars(parammap,parameters(sys))

src/systems/diffeqs/sdesystem.jl

@@ -161,7 +161,7 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,p=parammap;
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     checkbounds = false, sparse = false,
-                                    linenumbers = true, multithread=false,
+                                    linenumbers = true, parallel=SerialForm(),
                                     kwargs...)
 ```
 
@@ -172,11 +172,11 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,parammap=DiffEqBa
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     checkbounds = false, sparse = false,
-                                    linenumbers = true, multithread=false,
+                                    linenumbers = true, parallel=SerialForm(),
                                     kwargs...) where iip
 
     f = SDEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
-                        linenumbers=linenumbers,multithread=multithread,
+                        linenumbers=linenumbers,parallel=parallel,
                         sparse=sparse)
     u0 = varmap_to_vars(u0map,states(sys))
     p = varmap_to_vars(parammap,parameters(sys))
@@ -185,4 +185,4 @@ end
 
 function DiffEqBase.SDEProblem(sys::SDESystem, args...; kwargs...)
     SDEProblem{true}(sys, args...; kwargs...)
-end
+end

src/systems/nonlinear/nonlinearsystem.jl

@@ -71,7 +71,7 @@ function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
                                           parammap=DiffEqBase.NullParameters();
                                           jac = false, sparse=false,
                                           checkbounds = false,
-                                          linenumbers = true, multithread=false,
+                                          linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
 ```
 
@@ -82,13 +82,13 @@ function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
                                           parammap=DiffEqBase.NullParameters();
                                           jac = false, sparse=false,
                                           checkbounds = false,
-                                          linenumbers = true, multithread=false,
+                                          linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
 
     f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
-                              multithread=multithread,sparse=sparse,expression=Val{false})
+                              parallel=parallel,sparse=sparse,expression=Val{false})
     u0 = varmap_to_vars(u0map,dvs)
     p = varmap_to_vars(parammap,ps)
     NonlinearProblem(f,u0,tspan,p;kwargs...)

src/systems/optimization/optimizationsystem.jl

@@ -80,7 +80,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
                                           u0=nothing, lb=nothing, ub=nothing,
                                           hes = false, sparse = false,
                                           checkbounds = false,
-                                          linenumbers = true, multithread=false,
+                                          linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
 ```
 
@@ -92,13 +92,13 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
                                           u0=nothing, lb=nothing, ub=nothing,
                                           hes = false, sparse = false,
                                           checkbounds = false,
-                                          linenumbers = true, multithread=false,
+                                          linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
 
     f = generate_function(sys,checkbounds=checkbounds,linenumbers=linenumbers,
-                              multithread=multithread,expression=Val{false})
+                              parallel=parallel,expression=Val{false})
     u0 = varmap_to_vars(u0,dvs)
     p = varmap_to_vars(parammap,ps)
     lb = varmap_to_vars(lb,dvs)

test/bigsystem.jl

@@ -45,13 +45,36 @@ end
 
 f(du,u,nothing,0.0)
 
-multithreadedf = eval(ModelingToolkit.build_function(du,u,multithread=true)[2])
+multithreadedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.MultithreadedForm())[2])
 _du = rand(N,N,3)
 _u = rand(N,N,3)
 multithreadedf(_du,_u)
 
+using Distributed
+addprocs(4)
+distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DistributedForm())[2])
+
 jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
-multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,multithread=true)[2])
+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])
+
+MyA = zeros(N,N)
+AMx = zeros(N,N)
+DA  = zeros(N,N)
+
+f(_du,_u,nothing,0.0)
+multithreadedf(_du,_u)
+distributedf(_du,_u)
+
+#=
+using BenchmarkTools
+@btime f(_du,_u,nothing,0.0)
+@btime multithreadedf(_du,_u)
+@btime distributedf(_du,_u)
 
-#_jac = similar(jac,Float64)
-#multithreadedjac(_jac,_u)
+_jac = similar(jac,Float64)
+@btime serialjac(_jac,_u)
+@btime multithreadedjac(_jac,_u)
+@btime distributedjac(_jac,_u)
+=#