Cost Estimation interface (#14)

See issue #13

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

src/MetagraphOptimization.jl

@@ -5,6 +5,7 @@ A module containing tools to work on DAGs.
 """
 module MetagraphOptimization
 
+# graph types
 export DAG
 export Node
 export Edge
@@ -18,6 +19,7 @@ export FusedComputeTask
 export PossibleOperations
 export GraphProperties
 
+# graph functions
 export make_node
 export make_edge
 export insert_node
@@ -27,10 +29,13 @@ export is_exit_node
 export parents
 export children
 export compute
+export data
+export compute_effort
 export get_properties
 export get_exit_node
 export is_valid, is_scheduled
 
+# graph operation related
 export Operation
 export AppliedOperation
 export NodeFusion
@@ -42,6 +47,10 @@ export can_pop
 export reset_graph!
 export get_operations
 
+# ABC model
+export ParticleValue
+export ParticleA, ParticleB, ParticleC
+export ABCProcessDescription, ABCProcessInput, ABCModel
 export ComputeTaskP
 export ComputeTaskS1
 export ComputeTaskS2
@@ -49,14 +58,17 @@ export ComputeTaskV
 export ComputeTaskU
 export ComputeTaskSum
 
+# code generation related
 export execute
 export parse_dag, parse_process
 export gen_process_input
 export get_compute_function
-export ParticleValue
-export ParticleA, ParticleB, ParticleC
-export ABCProcessDescription, ABCProcessInput, ABCModel
 
+# estimator
+export cost_type, graph_cost, operation_effect
+export GlobalMetricEstimator, CDCost
+
+# machine info
 export Machine
 export get_machine_info
 
@@ -121,6 +133,9 @@ include("task/compute.jl")
 include("task/print.jl")
 include("task/properties.jl")
 
+include("estimator/interface.jl")
+include("estimator/global_metric.jl")
+
 include("models/interface.jl")
 include("models/print.jl")
 

src/estimator/global_metric.jl

@@ -0,0 +1,65 @@
+
+"""
+   CDCost
+
+Representation of a [`DAG`](@ref)'s cost as estimated by the [`GlobalMetricEstimator`](@ref).
+
+# Fields:
+`.data`: The total data transfer.\\
+`.computeEffort`: The total compute effort.\\
+`.computeIntensity`: The compute intensity, will always equal `.computeEffort / .data`.
+
+
+!!! note
+    Note that the `computeIntensity` doesn't necessarily make sense in the context of only operation costs. 
+    For example, for node fusions this will always be 0, since the computeEffort is zero. 
+    It will still work as intended when adding/subtracting to/from a `graph_cost` estimate.
+"""
+const CDCost = NamedTuple{(:data, :computeEffort, :computeIntensity), Tuple{Float64, Float64, Float64}}
+
+function +(cost1::CDCost, cost2::CDCost)::CDCost
+    d = cost1.data + cost2.data
+    ce = computeEffort = cost1.computeEffort + cost2.computeEffort
+    return (data = d, computeEffort = ce, computeIntensity = ce / d)::CDCost
+end
+
+function -(cost1::CDCost, cost2::CDCost)::CDCost
+    d = cost1.data - cost2.data
+    ce = computeEffort = cost1.computeEffort - cost2.computeEffort
+    return (data = d, computeEffort = ce, computeIntensity = ce / d)::CDCost
+end
+
+struct GlobalMetricEstimator <: AbstractEstimator end
+
+function cost_type(estimator::GlobalMetricEstimator)
+    return CDCost
+end
+
+function graph_cost(estimator::GlobalMetricEstimator, graph::DAG)
+    properties = get_properties(graph)
+    return (
+        data = properties.data,
+        computeEffort = properties.computeEffort,
+        computeIntensity = properties.computeIntensity,
+    )::CDCost
+end
+
+function operation_effect(estimator::GlobalMetricEstimator, graph::DAG, operation::NodeFusion)
+    return (data = -data(operation.input[2].task), computeEffort = 0.0, computeIntensity = 0.0)::CDCost
+end
+
+function operation_effect(estimator::GlobalMetricEstimator, graph::DAG, operation::NodeReduction)
+    s = length(operation.input) - 1
+    return (
+        data = s * -data(operation.input[1].task),
+        computeEffort = s * -compute_effort(operation.input[1].task),
+        computeIntensity = typeof(operation.input) <: DataTaskNode ? 0.0 : Inf,
+    )::CDCost
+end
+
+function operation_effect(estimator::GlobalMetricEstimator, graph::DAG, operation::NodeSplit)
+    s = length(operation.input.parents) - 1
+    d = s * data(operation.input.task)
+    ce = s * compute_effort(operation.input.task)
+    return (data = d, computeEffort = ce, computeIntensity = ce / d)::CDCost
+end

src/estimator/interface.jl

@@ -0,0 +1,44 @@
+
+"""
+    AbstractEstimator
+
+Abstract base type for an estimator. An estimator estimates the cost of a graph or the difference an operation applied to a graph will make to its cost.
+
+Interface functions are
+- [`graph_cost`](@ref)
+- [`operation_effect`](@ref)
+"""
+abstract type AbstractEstimator end
+
+"""
+    cost_type(estimator::AbstractEstimator)
+
+Interface function returning a specific estimator's cost type, i.e., the type returned by its implementation of [`graph_cost`](@ref) and [`operation_effect`](@ref).
+"""
+function cost_type end
+
+"""
+    graph_cost(estimator::AbstractEstimator, graph::DAG)
+
+Get the total estimated cost of the graph. The cost's data type can be chosen by the implementation, but should have usable comparison operators (<, <=, >, >=, ==) and basic math operators (+, -, *, /).
+"""
+function graph_cost end
+
+"""
+    operation_effect(estimator::AbstractEstimator, graph::DAG, operation::Operation)
+
+Get the estimated effect on the cost of the graph, such that `graph_cost(estimator, graph) + operation_effect(estimator, graph, operation) ~= graph_cost(estimator, graph_with_operation_applied)`. There is no hard requirement for this, but the better the estimate, the better an optimization algorithm will be.
+
+!!! note
+    There is a default implementation of this function, applying the operation, calling [`graph_cost`](@ref), then popping the operation again.
+    
+    It can be much faster to overload this function for a specific estimator and directly compute the effects from the operation if possible.
+"""
+function operation_effect(estimator::AbstractEstimator, graph::DAG, operation::Operation)
+    # This is currently not stably working, see issue #16
+    cost = graph_cost(estimator, graph)
+    push_operation!(graph, operation)
+    cost_after = graph_cost(estimator, graph)
+    pop_operation!(graph)
+    return cost_after - cost
+end

src/models/abc/compute.jl

@@ -55,7 +55,8 @@ function compute(::ComputeTaskS2, data1::ParticleValue{P}, data2::ParticleValue{
     @assert isapprox(data1.p.momentum.py, -data2.p.momentum.py, rtol = 0.001, atol = sqrt(eps())) "py: $(data1.p.momentum.py) vs. $(data2.p.momentum.py)"
     @assert isapprox(data1.p.momentum.pz, -data2.p.momentum.pz, rtol = 0.001, atol = sqrt(eps())) "pz: $(data1.p.momentum.pz) vs. $(data2.p.momentum.pz)"
     =#
-    return data1.v * inner_edge(data1.p) * data2.v
+    inner = inner_edge(data1.p)
+    return data1.v * inner * data2.v
 end
 
 """

src/models/abc/particle.jl

@@ -1,5 +1,7 @@
 using QEDbase
 
+import QEDbase.mass
+
 """
     ABCModel <: AbstractPhysicsModel
 
@@ -87,9 +89,9 @@ For 2 given (non-equal) particle types, return the third of ABC.
 """
 function interaction_result(t1::Type{T1}, t2::Type{T2}) where {T1 <: ABCParticle, T2 <: ABCParticle}
     @assert t1 != t2
-    if t1 != Type{ParticleA} && t2 != Type{ParticleA}
+    if t1 != ParticleA && t2 != ParticleA
         return ParticleA
-    elseif t1 != Type{ParticleB} && t2 != Type{ParticleB}
+    elseif t1 != ParticleB && t2 != ParticleB
         return ParticleB
     else
         return ParticleC
@@ -161,7 +163,6 @@ Takes 4 effective FLOP.
 function preserve_momentum(p1::ABCParticle, p2::ABCParticle)
     t3 = interaction_result(typeof(p1), typeof(p2))
     p3 = t3(p1.momentum + p2.momentum)
-
     return p3
 end
 

src/models/abc/properties.jl

@@ -3,35 +3,35 @@
 
 Return the compute effort of an S1 task.
 """
-compute_effort(t::ComputeTaskS1) = 11
+compute_effort(t::ComputeTaskS1) = 11.0
 
 """
     compute_effort(t::ComputeTaskS2)
 
 Return the compute effort of an S2 task.
 """
-compute_effort(t::ComputeTaskS2) = 12
+compute_effort(t::ComputeTaskS2) = 12.0
 
 """
     compute_effort(t::ComputeTaskU)
 
 Return the compute effort of a U task.
 """
-compute_effort(t::ComputeTaskU) = 1
+compute_effort(t::ComputeTaskU) = 1.0
 
 """
     compute_effort(t::ComputeTaskV)
 
 Return the compute effort of a V task.
 """
-compute_effort(t::ComputeTaskV) = 6
+compute_effort(t::ComputeTaskV) = 6.0
 
 """
     compute_effort(t::ComputeTaskP)
 
 Return the compute effort of a P task.
 """
-compute_effort(t::ComputeTaskP) = 0
+compute_effort(t::ComputeTaskP) = 0.0
 
 """
     compute_effort(t::ComputeTaskSum)
@@ -41,7 +41,7 @@ Return the compute effort of a Sum task.
 Note: This is a constant compute effort, even though sum scales with the number of its inputs. Since there is only ever a single sum node in a graph generated from the ABC-Model,
 this doesn't matter.
 """
-compute_effort(t::ComputeTaskSum) = 1
+compute_effort(t::ComputeTaskSum) = 1.0
 
 """
     show(io::IO, t::DataTask)

src/models/abc/types.jl

@@ -4,7 +4,7 @@
 Task representing a specific data transfer in the ABC Model.
 """
 struct DataTask <: AbstractDataTask
-    data::UInt64
+    data::Float64
 end
 
 """

src/task/properties.jl

@@ -49,7 +49,7 @@ end
 
 Return the compute effort of a data task, always zero, regardless of the specific task.
 """
-compute_effort(t::AbstractDataTask) = 0
+compute_effort(t::AbstractDataTask) = 0.0
 
 """
     data(t::AbstractDataTask)
@@ -63,7 +63,7 @@ data(t::AbstractDataTask) = getfield(t, :data)
 
 Return the data of a compute task, always zero, regardless of the specific task.
 """
-data(t::AbstractComputeTask) = 0
+data(t::AbstractComputeTask) = 0.0
 
 """
     compute_effort(t::FusedComputeTask)