Merge pull request #431 from jpsamaroo/jps/daggerform

Add DaggerForm

Author: ChrisRackauckas

Project.toml

@@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -50,10 +51,11 @@ Unitful = "1.1"
 julia = "1.2"
 
 [extras]
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["OrdinaryDiffEq", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]
+test = ["Dagger", "OrdinaryDiffEq", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]

src/ModelingToolkit.jl

@@ -23,6 +23,8 @@ import LightGraphs: SimpleDiGraph, add_edge!
 
 import TreeViews
 
+using Requires
+
 """
 $(TYPEDEF)
 
@@ -148,4 +150,10 @@ export build_function
 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

@@ -8,6 +8,7 @@ abstract type ParallelForm end
 struct SerialForm <: ParallelForm end
 struct MultithreadedForm <: ParallelForm end
 struct DistributedForm <: ParallelForm end
+struct DaggerForm <: ParallelForm end
 
 """
 `build_function`
@@ -196,6 +197,10 @@ function _build_function(target::JuliaTarget, rhss, args...;
 		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 parallel isa DaggerForm
+		computevars = [Variable(gensym(:MTComputeVar))() for i in axes(rhss,1)]
+        reducevar = Variable(gensym(:MTReduceVar))()
+        _rhss = [getindex(reducevar,i) for i in axes(rhss,1)]
 	elseif rhss isa SparseMatrixCSC
 		_rhss = rhss.nzval
 	else
@@ -290,6 +295,23 @@ function _build_function(target::JuliaTarget, rhss, args...;
 			$(resunpack_exprs...)
 			$(ip_let_expr.args[2])
 		end
+    elseif parallel isa DaggerForm
+        @assert HAS_DAGGER[] "Dagger.jl is not loaded; please do `using Dagger`"
+        dagwrap(x) = x
+        dagwrap(ex::Expr) = dagwrap(ex, Val(ex.head))
+        dagwrap(ex::Expr, ::Val) = ex
+        dagwrap(ex::Expr, ::Val{:call}) = :(Dagger.delayed($(ex.args[1]))($(dagwrap.(ex.args[2:end])...)))
+        new_rhss = dagwrap.(conv.(rhss))
+        delayed_exprs = build_expr(:block, [:($(Symbol(computevars[i])) = Dagger.delayed(identity)($(new_rhss[i]))) for i in axes(computevars,1)])
+        # TODO: treereduce?
+        reduce_expr = quote
+            $(Symbol(reducevar)) = collect(Dagger.delayed(vcat)($(computevars...)))
+        end
+        ip_let_expr.args[2] = quote
+            $delayed_exprs
+            $reduce_expr
+            $(ip_let_expr.args[2])
+        end
 	end
 
     tuple_sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])

src/dagger.jl

@@ -0,0 +1,3 @@
+using .Dagger
+
+HAS_DAGGER[] = true

test/bigsystem.jl

@@ -54,10 +54,14 @@ 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 = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
 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)
@@ -66,15 +70,18 @@ DA  = zeros(N,N)
 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)
 =#