Add scheduling, machine info, caching strategies and devices (#9)

Reviewed-on: Rubydragon/MetagraphOptimization.jl#9
Co-authored-by: Anton Reinhard <anton.reinhard@proton.me>
Co-committed-by: Anton Reinhard <anton.reinhard@proton.me>

src/task/compute.jl

@@ -0,0 +1,89 @@
+
+"""
+    compute(t::FusedComputeTask, data)
+
+Compute a [`FusedComputeTask`](@ref). This simply asserts false and should not be called. Fused Compute Tasks generate their expressions directly through the other tasks instead.
+"""
+function compute(t::FusedComputeTask, data)
+    @assert false "This is not implemented and should never be called"
+end
+
+"""
+    get_expression(t::FusedComputeTask, device::AbstractDevice, inExprs::Vector{String}, outExpr::String)
+
+Generate code evaluating a [`FusedComputeTask`](@ref) on `inExprs`, providing the output on `outExpr`.
+`inExprs` should be of the correct types and may be heterogeneous. `outExpr` will be of the type of the output of `T2` of t.
+"""
+function get_expression(t::FusedComputeTask, device::AbstractDevice, inExprs::Vector, outExpr)
+    inExprs1 = Vector()
+    for sym in t.t1_inputs
+        push!(inExprs1, gen_access_expr(device, sym))
+    end
+
+    outExpr1 = gen_access_expr(device, t.t1_output)
+
+    inExprs2 = Vector()
+    for sym in t.t2_inputs
+        push!(inExprs2, gen_access_expr(device, sym))
+    end
+
+    expr1 = get_expression(t.first_task, device, inExprs1, outExpr1)
+    expr2 = get_expression(t.second_task, device, [inExprs2..., outExpr1], outExpr)
+
+    full_expr = Expr(:block, expr1, expr2)
+
+    return full_expr
+end
+
+"""
+    get_expression(node::ComputeTaskNode)
+
+Generate and return code for a given [`ComputeTaskNode`](@ref).
+"""
+function get_expression(node::ComputeTaskNode)
+    @assert length(node.children) <= children(node.task) "Node $(node) has too many children for its task: node has $(length(node.children)) versus task has $(children(node.task))\nNode's children: $(getfield.(node.children, :children))"
+    @assert !ismissing(node.device) "Trying to get expression for an unscheduled ComputeTaskNode\nNode: $(node)"
+
+    inExprs = Vector()
+    for id in getfield.(node.children, :id)
+        push!(inExprs, gen_access_expr(node.device, Symbol(to_var_name(id))))
+    end
+    outExpr = gen_access_expr(node.device, Symbol(to_var_name(node.id)))
+
+    return get_expression(node.task, node.device, inExprs, outExpr)
+end
+
+"""
+    get_expression(node::DataTaskNode)
+
+Generate and return code for a given [`DataTaskNode`](@ref).
+"""
+function get_expression(node::DataTaskNode)
+    @assert length(node.children) == 1 "Trying to call get_expression on a data task node that has $(length(node.children)) children instead of 1"
+
+    # TODO: dispatch to device implementations generating the copy commands
+
+    child = node.children[1]
+    inExpr = eval(gen_access_expr(child.device, Symbol(to_var_name(child.id))))
+    outExpr = eval(gen_access_expr(child.device, Symbol(to_var_name(node.id))))
+    dataTransportExp = Meta.parse("$outExpr = $inExpr")
+
+    return dataTransportExp
+end
+
+"""
+    get_init_expression(node::DataTaskNode, device::AbstractDevice)
+
+Generate and return code for the initial input reading expression for [`DataTaskNode`](@ref)s with 0 children, i.e., entry nodes.
+
+See also: [`get_entry_nodes`](@ref)
+"""
+function get_init_expression(node::DataTaskNode, device::AbstractDevice)
+    @assert isempty(node.children) "Trying to call get_init_expression on a data task node that is not an entry node."
+
+    inExpr = eval(gen_access_expr(device, Symbol("$(to_var_name(node.id))_in")))
+    outExpr = eval(gen_access_expr(device, Symbol(to_var_name(node.id))))
+    dataTransportExp = Meta.parse("$outExpr = $inExpr")
+
+    return dataTransportExp
+end

src/task/create.jl

@@ -3,8 +3,7 @@
 
 Fallback implementation of the copy of an abstract data task, throwing an error.
 """
-copy(t::AbstractDataTask) =
-    error("Need to implement copying for your data tasks!")
+copy(t::AbstractDataTask) = error("Need to implement copying for your data tasks!")
 
 """
     copy(t::AbstractComputeTask)
@@ -12,3 +11,21 @@ copy(t::AbstractDataTask) =
 Return a copy of the given compute task.
 """
 copy(t::AbstractComputeTask) = typeof(t)()
+
+"""
+    copy(t::FusedComputeTask)
+
+Return a copy of th egiven [`FusedComputeTask`](@ref).
+"""
+function copy(t::FusedComputeTask{T1, T2}) where {T1, T2}
+    return FusedComputeTask{T1, T2}(
+        copy(t.first_task),
+        copy(t.second_task),
+        copy(t.t1_inputs),
+        t.t1_output,
+        copy(t.t2_inputs),
+    )
+end
+
+FusedComputeTask{T1, T2}(t1_inputs::Vector{String}, t1_output::String, t2_inputs::Vector{String}) where {T1, T2} =
+    FusedComputeTask{T1, T2}(T1(), T2(), t1_inputs, t1_output, t2_inputs)

src/task/print.jl

@@ -4,6 +4,5 @@
 Print a string representation of the fused compute task to io.
 """
 function show(io::IO, t::FusedComputeTask)
-    (T1, T2) = get_types(t)
-    return print(io, "ComputeFuse(", T1(), ", ", T2(), ")")
+    return print(io, "ComputeFuse($(t.first_task), $(t.second_task))")
 end

src/task/properties.jl

@@ -71,8 +71,7 @@ data(t::AbstractComputeTask) = 0
 Return the compute effort of a fused compute task. 
 """
 function compute_effort(t::FusedComputeTask)
-    (T1, T2) = collect(typeof(t).parameters)
-    return compute_effort(T1()) + compute_effort(T2())
+    return compute_effort(t.first_task) + compute_effort(t.second_task)
 end
 
 """
@@ -81,30 +80,3 @@ end
 Return a tuple of a the fused compute task's components' types.
 """
 get_types(::FusedComputeTask{T1, T2}) where {T1, T2} = (T1, T2)
-
-"""
-    get_expression(t::AbstractTask)
-
-Return an expression evaluating the given task on the :dataIn symbol
-"""
-function get_expression(t::AbstractTask)
-    return quote
-        dataOut = compute($t, dataIn)
-    end
-end
-
-"""
-    get_expression()
-"""
-function get_expression(
-    t::FusedComputeTask,
-    inSymbol::Symbol,
-    outSymbol::Symbol,
-)
-    #TODO
-    computeExp = quote
-        $outSymbol = compute($t, $inSymbol)
-    end
-
-    return computeExp
-end

src/task/type.jl

@@ -26,5 +26,13 @@ A fused compute task made up of the computation of first `T1` and then `T2`.
 
 Also see: [`get_types`](@ref).
 """
-struct FusedComputeTask{T1 <: AbstractComputeTask, T2 <: AbstractComputeTask} <:
-       AbstractComputeTask end
+struct FusedComputeTask{T1 <: AbstractComputeTask, T2 <: AbstractComputeTask} <: AbstractComputeTask
+    first_task::T1
+    second_task::T2
+    # the names of the inputs for T1
+    t1_inputs::Vector{Symbol}
+    # output name of T1
+    t1_output::Symbol
+    # t2_inputs doesn't include the output of t1, that's implicit
+    t2_inputs::Vector{Symbol}
+end