Add DaggerForm

Author: jpsamaroo

Project.toml

@@ -16,6 +16,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"

src/ModelingToolkit.jl

@@ -20,6 +20,8 @@ import SymbolicUtils: to_symbolic, FnType
 
 import TreeViews
 
+using Requires
+
 """
 $(TYPEDEF)
 
@@ -133,4 +135,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`
@@ -165,6 +166,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
@@ -222,6 +227,18 @@ 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`"
+        delayed_exprs = build_expr(:block, [:($(Symbol(computevars[i])) = Dagger.delayed(identity)($(conv(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