From f1edce258a028db608653d64abca2548a9e154cc Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@wandelbots.com>
Date: Thu, 31 Aug 2023 18:24:48 +0200
Subject: [PATCH 1/8] Start adding code generation

---
 src/MetagraphOptimization.jl |   8 +++
 src/code_gen/main.jl         |  32 +++++++++++
 src/graph/properties.jl      |  15 +++++
 src/models/abc/compute.jl    | 106 +++++++++++++++++++++++++++++++++++
 src/models/abc/particle.jl   |  49 ++++++++++++++++
 src/task/properties.jl       |  51 +++++++++++++----
 6 files changed, 251 insertions(+), 10 deletions(-)
 create mode 100644 src/code_gen/main.jl
 create mode 100644 src/models/abc/compute.jl
 create mode 100644 src/models/abc/particle.jl

diff --git a/src/MetagraphOptimization.jl b/src/MetagraphOptimization.jl
index 1e8c0c9..fadf456 100644
--- a/src/MetagraphOptimization.jl
+++ b/src/MetagraphOptimization.jl
@@ -50,6 +50,10 @@ export ComputeTaskV
 export ComputeTaskU
 export ComputeTaskSum
 
+export gen_code
+export ParticleValue
+export Particle
+
 export ==, in, show, isempty, delete!, length
 
 export bytes_to_human_readable
@@ -104,12 +108,16 @@ include("operation/validate.jl")
 include("properties/create.jl")
 include("properties/utility.jl")
 
+include("code_gen/main.jl")
+
 include("task/create.jl")
 include("task/compare.jl")
 include("task/print.jl")
 include("task/properties.jl")
 
 include("models/abc/types.jl")
+include("models/abc/particle.jl")
+include("models/abc/compute.jl")
 include("models/abc/properties.jl")
 include("models/abc/parse.jl")
 
diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
new file mode 100644
index 0000000..1180bc6
--- /dev/null
+++ b/src/code_gen/main.jl
@@ -0,0 +1,32 @@
+using DataStructures
+
+function gen_code(graph::DAG)
+    code = Vector{Expr}()
+    sizehint!(code, length(graph.nodes))
+
+    nodeQueue = PriorityQueue{Node, Int}()
+    inputSyms = Vector{Symbol}()
+
+    for node in get_entry_nodes(graph)
+        enqueue!(nodeQueue, node => 1)
+        push!(inputSyms, Symbol("data_$(replace(string(node.id), "-"=>"_"))_in"))
+    end
+
+    node = nothing
+    while !isempty(nodeQueue)
+        prio = peek(nodeQueue)[2]
+        node = dequeue!(nodeQueue)
+
+        push!(code, get_expression(node))
+        for parent in node.parents
+            if (!haskey(nodeQueue, parent))
+                enqueue!(nodeQueue, parent => prio + length(parent.children))
+            end
+        end
+    end
+
+    # node is now the last node we looked at -> the output node
+    outSym = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+    
+    return (code = Expr(:block, code...), inputSymbols = inputSyms, outputSymbol = outSym)
+end
diff --git a/src/graph/properties.jl b/src/graph/properties.jl
index 211d8be..2fd89db 100644
--- a/src/graph/properties.jl
+++ b/src/graph/properties.jl
@@ -27,3 +27,18 @@ function get_exit_node(graph::DAG)
     end
     @assert false "The given graph has no exit node! It is either empty or not acyclic!"
 end
+
+"""
+    get_entry_nodes(graph::DAG)
+
+Return a vector of the graph's entry nodes.
+"""
+function get_entry_nodes(graph::DAG)
+    result = Vector{Node}()
+    for node in graph.nodes
+        if (is_entry_node(node))
+            push!(result, node)
+        end
+    end
+    return result
+end
diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
new file mode 100644
index 0000000..f931756
--- /dev/null
+++ b/src/models/abc/compute.jl
@@ -0,0 +1,106 @@
+
+# Compute Particle, nothing to be done (0 FLOP)
+function compute(::ComputeTaskP, data::ParticleValue)
+    return data
+end
+
+# generate code evaluating ComputeTaskP on inSymbol, providing the output on outSymbol
+function get_expression(::ComputeTaskP, inSymbol::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskP(), $inSymbol)")
+end
+
+# Compute outer edge
+function compute(::ComputeTaskU, data::ParticleValue)
+    return ParticleValue(data.p, data.v * outer_edge(data.p))
+end
+
+# generate code evaluating ComputeTaskU on inSymbol, providing the output on outSymbol
+# inSymbol should be of type ParticleValue, outSymbol will be of type ParticleValue
+function get_expression(::ComputeTaskU, inSymbol::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskU(), $inSymbol)")
+end
+
+# compute vertex
+function compute(::ComputeTaskV, data1, data2)
+    # calculate new particle from the two input particles
+    p3 = preserve_momentum(data1.p, data2.p)
+    dataOut = ParticleValue(p3, data1.v * vertex() * data2.v)
+    return dataOut
+end
+
+function get_expression(::ComputeTaskV, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskV(), $inSymbol1, $inSymbol2)")
+end
+
+# compute final inner edge (no output particle)
+function compute(::ComputeTaskS2, data1, data2)
+    return data1.v * inner_edge(data1.p) * data2.v
+end
+
+function get_expression(::ComputeTaskS2, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskS2(), $inSymbol1, $inSymbol2)")
+end
+
+# compute inner edge
+function compute(::ComputeTaskS1, data)
+    return (particle = data.p, v = data.v * inner_edge(data.p))
+end
+
+function get_expression(::ComputeTaskS1, inSymbol::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskS1(), $inSymbol)")
+end
+
+function compute(::ComputeTaskSum, data::Vector{Float64})
+    return sum(data)
+end
+
+function get_expression(::ComputeTaskSum, inSymbols::Vector{Symbol}, outSymbol::Symbol)
+    return quote
+        $outSymbol = compute(ComputeTaskSum(), [$(inSymbols...)])
+    end
+end
+
+function get_expression(node::ComputeTaskNode)
+    t = typeof(node.task)
+    if (t <: ComputeTaskU || t <: ComputeTaskP || t <: ComputeTaskS1) # single input
+        @assert length(node.children) == 1
+        symbolIn = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        return get_expression(t(), symbolIn, symbolOut)
+    elseif (t <: ComputeTaskS2 || t <: ComputeTaskV) # double input
+        @assert length(node.children) == 2
+        symbolIn1 = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        symbolIn2 = Symbol("data_$(replace(string(node.children[2].id), "-"=>"_"))")
+        symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        return get_expression(t(), symbolIn1, symbolIn2, symbolOut)
+    elseif (t <: ComputeTaskSum) # vector input
+        @assert length(node.children) > 0
+        inSymbols = Vector{Symbol}()
+        for child in node.children
+            push!(inSymbols, Symbol("data_$(replace(string(child.id), "-"=>"_"))"))
+        end
+        outSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        return get_expression(t(), inSymbols, outSymbol)
+    elseif (t <: FusedComputeTask)
+        # uuuuuh
+    else
+        error("Unknown compute task")
+    end
+end
+
+function get_expression(node::DataTaskNode)
+    # TODO: do things to transport data from/to gpu, between numa nodes, etc.
+    @assert length(node.children) <= 1
+
+    inSymbol = nothing
+    if (length(node.children) == 1)
+        inSymbol = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+    else
+        inSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))_in")
+    end
+    outSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+
+    dataTransportExp = Meta.parse("$outSymbol = $inSymbol")
+
+    return dataTransportExp
+end
diff --git a/src/models/abc/particle.jl b/src/models/abc/particle.jl
new file mode 100644
index 0000000..d515b06
--- /dev/null
+++ b/src/models/abc/particle.jl
@@ -0,0 +1,49 @@
+
+struct Particle
+    P0::Float64
+    P1::Float64
+    P2::Float64
+    P3::Float64
+
+    m::Float64
+end
+
+struct ParticleValue
+    p::Particle
+    v::Float64
+end
+
+function square(p::Particle)
+    return p.P0 * p.P0 - p.P1 * p.P1 - p.P2 * p.P2 - p.P3 * p.P3
+end
+
+function inner_edge(p::Particle)
+    return 1.0 / (square(p) - p.m * p.m)
+end
+
+function outer_edge(p::Particle)
+    return 1.0
+end
+
+function vertex()
+    i = 1.0
+    lambda = 1.0/137.0
+    return i * lambda
+end
+
+# calculate new particle from two given interacting ones
+function preserve_momentum(p1::Particle, p2::Particle)
+    # TODO: is this correct?
+
+    p3 = Particle(
+        p1.P0 + p2.P0,
+        p1.P1 + p2.P1,
+        p1.P2 + p2.P2,
+        p1.P3 + p2.P3,
+        1.0
+    )
+
+    # m3 = sqrt(-  PC * PC / c^2)
+
+    return p3
+end
diff --git a/src/task/properties.jl b/src/task/properties.jl
index dd120e9..11aa59b 100644
--- a/src/task/properties.jl
+++ b/src/task/properties.jl
@@ -7,6 +7,24 @@ function compute(t::AbstractTask; data...)
     return error("Need to implement compute()")
 end
 
+"""
+    compute(t::FusedComputeTask; data...)
+
+Compute a fused compute task.
+"""
+function compute(t::FusedComputeTask; data...)
+    (T1, T2) = collect(typeof(t).parameters)
+
+    return compute(T2(), compute(T1(), data))
+end
+
+"""
+    compute(t::AbstractDataTask; data...)
+
+The compute function of a data task, always the identity function, regardless of the specific task.
+"""
+compute(t::AbstractDataTask; data...) = data
+
 """
     compute_effort(t::AbstractTask)
 
@@ -33,13 +51,6 @@ Return the compute effort of a data task, always zero, regardless of the specifi
 """
 compute_effort(t::AbstractDataTask) = 0
 
-"""
-    compute(t::AbstractDataTask; data...)
-
-The compute function of a data task, always the identity function, regardless of the specific task.
-"""
-compute(t::AbstractDataTask; data...) = data
-
 """
     data(t::AbstractDataTask)
 
@@ -64,12 +75,32 @@ function compute_effort(t::FusedComputeTask)
     return compute_effort(T1()) + compute_effort(T2())
 end
 
-# actual compute functions for the tasks can stay undefined for now
-# compute(t::ComputeTaskU, data::Any) = mycomputation(data)
-
 """
     get_types(::FusedComputeTask{T1, T2})
 
 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

From 7a1a97dac84aac7db55dcd93df2478897eb64f38 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@wandelbots.com>
Date: Thu, 31 Aug 2023 18:47:05 +0200
Subject: [PATCH 2/8] Add basic execution function

---
 src/MetagraphOptimization.jl |  6 ++---
 src/code_gen/main.jl         | 36 ++++++++++++++++++++++++++---
 src/models/abc/compute.jl    | 45 ++++++++++++++++++++++++++++--------
 src/models/abc/particle.jl   |  4 ++--
 src/task/properties.jl       |  6 ++++-
 5 files changed, 78 insertions(+), 19 deletions(-)

diff --git a/src/MetagraphOptimization.jl b/src/MetagraphOptimization.jl
index fadf456..9c09392 100644
--- a/src/MetagraphOptimization.jl
+++ b/src/MetagraphOptimization.jl
@@ -50,7 +50,7 @@ export ComputeTaskV
 export ComputeTaskU
 export ComputeTaskSum
 
-export gen_code
+export execute
 export ParticleValue
 export Particle
 
@@ -108,8 +108,6 @@ include("operation/validate.jl")
 include("properties/create.jl")
 include("properties/utility.jl")
 
-include("code_gen/main.jl")
-
 include("task/create.jl")
 include("task/compare.jl")
 include("task/print.jl")
@@ -121,4 +119,6 @@ include("models/abc/compute.jl")
 include("models/abc/properties.jl")
 include("models/abc/parse.jl")
 
+include("code_gen/main.jl")
+
 end # module MetagraphOptimization
diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index 1180bc6..919b2a7 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -9,7 +9,10 @@ function gen_code(graph::DAG)
 
     for node in get_entry_nodes(graph)
         enqueue!(nodeQueue, node => 1)
-        push!(inputSyms, Symbol("data_$(replace(string(node.id), "-"=>"_"))_in"))
+        push!(
+            inputSyms,
+            Symbol("data_$(replace(string(node.id), "-"=>"_"))_in"),
+        )
     end
 
     node = nothing
@@ -27,6 +30,33 @@ function gen_code(graph::DAG)
 
     # node is now the last node we looked at -> the output node
     outSym = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
-    
-    return (code = Expr(:block, code...), inputSymbols = inputSyms, outputSymbol = outSym)
+
+    return (
+        code = Expr(:block, code...),
+        inputSymbols = inputSyms,
+        outputSymbol = outSym,
+    )
+end
+
+function execute(graph::DAG, input::Vector{Particle})
+    (code, inputSymbols, outputSymbol) = gen_code(graph)
+
+    @assert length(input) == length(inputSymbols)
+
+    assignInputs = Vector{Expr}()
+    for i in 1:length(input)
+        push!(
+            assignInputs,
+            Meta.parse(
+                "$(inputSymbols[i]) = ParticleValue(Particle($(input[i]).P0, $(input[i]).P1, $(input[i]).P2, $(input[i]).P3, $(input[i]).m), 1.0)",
+            ),
+        )
+    end
+
+    assignInputs = Expr(:block, assignInputs...)
+    eval(assignInputs)
+    eval(code)
+
+    eval(Meta.parse("result = $outputSymbol"))
+    return result
 end
diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
index f931756..d3c7f84 100644
--- a/src/models/abc/compute.jl
+++ b/src/models/abc/compute.jl
@@ -28,8 +28,15 @@ function compute(::ComputeTaskV, data1, data2)
     return dataOut
 end
 
-function get_expression(::ComputeTaskV, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
-    return Meta.parse("$outSymbol = compute(ComputeTaskV(), $inSymbol1, $inSymbol2)")
+function get_expression(
+    ::ComputeTaskV,
+    inSymbol1::Symbol,
+    inSymbol2::Symbol,
+    outSymbol::Symbol,
+)
+    return Meta.parse(
+        "$outSymbol = compute(ComputeTaskV(), $inSymbol1, $inSymbol2)",
+    )
 end
 
 # compute final inner edge (no output particle)
@@ -37,8 +44,15 @@ function compute(::ComputeTaskS2, data1, data2)
     return data1.v * inner_edge(data1.p) * data2.v
 end
 
-function get_expression(::ComputeTaskS2, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
-    return Meta.parse("$outSymbol = compute(ComputeTaskS2(), $inSymbol1, $inSymbol2)")
+function get_expression(
+    ::ComputeTaskS2,
+    inSymbol1::Symbol,
+    inSymbol2::Symbol,
+    outSymbol::Symbol,
+)
+    return Meta.parse(
+        "$outSymbol = compute(ComputeTaskS2(), $inSymbol1, $inSymbol2)",
+    )
 end
 
 # compute inner edge
@@ -54,7 +68,11 @@ function compute(::ComputeTaskSum, data::Vector{Float64})
     return sum(data)
 end
 
-function get_expression(::ComputeTaskSum, inSymbols::Vector{Symbol}, outSymbol::Symbol)
+function get_expression(
+    ::ComputeTaskSum,
+    inSymbols::Vector{Symbol},
+    outSymbol::Symbol,
+)
     return quote
         $outSymbol = compute(ComputeTaskSum(), [$(inSymbols...)])
     end
@@ -64,20 +82,26 @@ function get_expression(node::ComputeTaskNode)
     t = typeof(node.task)
     if (t <: ComputeTaskU || t <: ComputeTaskP || t <: ComputeTaskS1) # single input
         @assert length(node.children) == 1
-        symbolIn = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        symbolIn =
+            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
         symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
         return get_expression(t(), symbolIn, symbolOut)
     elseif (t <: ComputeTaskS2 || t <: ComputeTaskV) # double input
         @assert length(node.children) == 2
-        symbolIn1 = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
-        symbolIn2 = Symbol("data_$(replace(string(node.children[2].id), "-"=>"_"))")
+        symbolIn1 =
+            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        symbolIn2 =
+            Symbol("data_$(replace(string(node.children[2].id), "-"=>"_"))")
         symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
         return get_expression(t(), symbolIn1, symbolIn2, symbolOut)
     elseif (t <: ComputeTaskSum) # vector input
         @assert length(node.children) > 0
         inSymbols = Vector{Symbol}()
         for child in node.children
-            push!(inSymbols, Symbol("data_$(replace(string(child.id), "-"=>"_"))"))
+            push!(
+                inSymbols,
+                Symbol("data_$(replace(string(child.id), "-"=>"_"))"),
+            )
         end
         outSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
         return get_expression(t(), inSymbols, outSymbol)
@@ -94,7 +118,8 @@ function get_expression(node::DataTaskNode)
 
     inSymbol = nothing
     if (length(node.children) == 1)
-        inSymbol = Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        inSymbol =
+            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
     else
         inSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))_in")
     end
diff --git a/src/models/abc/particle.jl b/src/models/abc/particle.jl
index d515b06..e336de6 100644
--- a/src/models/abc/particle.jl
+++ b/src/models/abc/particle.jl
@@ -27,7 +27,7 @@ end
 
 function vertex()
     i = 1.0
-    lambda = 1.0/137.0
+    lambda = 1.0 / 137.0
     return i * lambda
 end
 
@@ -40,7 +40,7 @@ function preserve_momentum(p1::Particle, p2::Particle)
         p1.P1 + p2.P1,
         p1.P2 + p2.P2,
         p1.P3 + p2.P3,
-        1.0
+        1.0,
     )
 
     # m3 = sqrt(-  PC * PC / c^2)
diff --git a/src/task/properties.jl b/src/task/properties.jl
index 11aa59b..4b1d889 100644
--- a/src/task/properties.jl
+++ b/src/task/properties.jl
@@ -96,7 +96,11 @@ end
 """
     get_expression()
 """
-function get_expression(t::FusedComputeTask, inSymbol::Symbol, outSymbol::Symbol)
+function get_expression(
+    t::FusedComputeTask,
+    inSymbol::Symbol,
+    outSymbol::Symbol,
+)
     #TODO
     computeExp = quote
         $outSymbol = compute($t, $inSymbol)

From 0f78053ccfbbcd3b04e764738fb42f928d1fc8ed Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@proton.me>
Date: Tue, 5 Sep 2023 12:14:41 +0200
Subject: [PATCH 3/8] Fix topoligical ordering on the graph

---
 docs/src/lib/internals/models.md | 21 +++++++++++++++++
 src/MetagraphOptimization.jl     |  2 ++
 src/code_gen/main.jl             | 39 +++++++++++++++++++++++++-------
 src/models/abc/compute.jl        | 31 ++++++++++---------------
 src/models/abc/create.jl         | 30 ++++++++++++++++++++++++
 src/node/print.jl                |  9 ++++++++
 6 files changed, 105 insertions(+), 27 deletions(-)
 create mode 100644 src/models/abc/create.jl

diff --git a/docs/src/lib/internals/models.md b/docs/src/lib/internals/models.md
index d28254d..1f7877f 100644
--- a/docs/src/lib/internals/models.md
+++ b/docs/src/lib/internals/models.md
@@ -9,6 +9,13 @@ Pages = ["models/abc/types.jl"]
 Order   = [:type, :constant]
 ```
 
+### Particle
+```@autodocs
+Modules = [MetagraphOptimization]
+Pages = ["models/abc/particle.jl"]
+Order   = [:type, :constant]
+```
+
 ### Parse
 ```@autodocs
 Modules = [MetagraphOptimization]
@@ -23,6 +30,20 @@ Pages = ["models/abc/properties.jl"]
 Order   = [:function]
 ```
 
+### Create
+```@autodocs
+Modules = [MetagraphOptimization]
+Pages = ["models/abc/create.jl]
+Order = [:function]
+```
+
+### Compute
+```@autodocs
+Modules = [MetagraphOptimization]
+Pages = ["models/abc/compute.jl]
+Order = [:function]
+```
+
 ## QED-Model
 
 *To be added*
diff --git a/src/MetagraphOptimization.jl b/src/MetagraphOptimization.jl
index 9c09392..093536a 100644
--- a/src/MetagraphOptimization.jl
+++ b/src/MetagraphOptimization.jl
@@ -51,6 +51,7 @@ export ComputeTaskU
 export ComputeTaskSum
 
 export execute
+export gen_particles
 export ParticleValue
 export Particle
 
@@ -116,6 +117,7 @@ include("task/properties.jl")
 include("models/abc/types.jl")
 include("models/abc/particle.jl")
 include("models/abc/compute.jl")
+include("models/abc/create.jl")
 include("models/abc/properties.jl")
 include("models/abc/parse.jl")
 
diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index 919b2a7..e392510 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -7,29 +7,30 @@ function gen_code(graph::DAG)
     nodeQueue = PriorityQueue{Node, Int}()
     inputSyms = Vector{Symbol}()
 
+    # use a priority equal to the number of unseen children -> 0 are nodes that can be added
     for node in get_entry_nodes(graph)
-        enqueue!(nodeQueue, node => 1)
-        push!(
-            inputSyms,
-            Symbol("data_$(replace(string(node.id), "-"=>"_"))_in"),
-        )
+        enqueue!(nodeQueue, node => 0)
+        push!(inputSyms, Symbol("data_$(to_var_name(node.id))_in"))
     end
 
     node = nothing
     while !isempty(nodeQueue)
-        prio = peek(nodeQueue)[2]
+        @assert peek(nodeQueue)[2] == 0
         node = dequeue!(nodeQueue)
 
         push!(code, get_expression(node))
         for parent in node.parents
+            # reduce the priority of all parents by one
             if (!haskey(nodeQueue, parent))
-                enqueue!(nodeQueue, parent => prio + length(parent.children))
+                enqueue!(nodeQueue, parent => length(parent.children) - 1)
+            else
+                nodeQueue[parent] = nodeQueue[parent] - 1
             end
         end
     end
 
     # node is now the last node we looked at -> the output node
-    outSym = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+    outSym = Symbol("data_$(to_var_name(node.id))")
 
     return (
         code = Expr(:block, code...),
@@ -38,6 +39,28 @@ function gen_code(graph::DAG)
     )
 end
 
+function execute(generated_code, input::Vector{Particle})
+    (code, inputSymbols, outputSymbol) = generated_code
+    @assert length(input) == length(inputSymbols)
+
+    assignInputs = Vector{Expr}()
+    for i in 1:length(input)
+        push!(
+            assignInputs,
+            Meta.parse(
+                "$(inputSymbols[i]) = ParticleValue(Particle($(input[i]).P0, $(input[i]).P1, $(input[i]).P2, $(input[i]).P3, $(input[i]).m), 1.0)",
+            ),
+        )
+    end
+
+    assignInputs = Expr(:block, assignInputs...)
+    eval(assignInputs)
+    eval(code)
+
+    eval(Meta.parse("result = $outputSymbol"))
+    return result
+end
+
 function execute(graph::DAG, input::Vector{Particle})
     (code, inputSymbols, outputSymbol) = gen_code(graph)
 
diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
index d3c7f84..e0143dd 100644
--- a/src/models/abc/compute.jl
+++ b/src/models/abc/compute.jl
@@ -21,7 +21,7 @@ function get_expression(::ComputeTaskU, inSymbol::Symbol, outSymbol::Symbol)
 end
 
 # compute vertex
-function compute(::ComputeTaskV, data1, data2)
+function compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
     # calculate new particle from the two input particles
     p3 = preserve_momentum(data1.p, data2.p)
     dataOut = ParticleValue(p3, data1.v * vertex() * data2.v)
@@ -57,7 +57,7 @@ end
 
 # compute inner edge
 function compute(::ComputeTaskS1, data)
-    return (particle = data.p, v = data.v * inner_edge(data.p))
+    return ParticleValue(data.p, data.v * inner_edge(data.p))
 end
 
 function get_expression(::ComputeTaskS1, inSymbol::Symbol, outSymbol::Symbol)
@@ -82,28 +82,22 @@ function get_expression(node::ComputeTaskNode)
     t = typeof(node.task)
     if (t <: ComputeTaskU || t <: ComputeTaskP || t <: ComputeTaskS1) # single input
         @assert length(node.children) == 1
-        symbolIn =
-            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
-        symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        symbolIn = Symbol("data_$(to_var_name(node.children[1].id))")
+        symbolOut = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), symbolIn, symbolOut)
     elseif (t <: ComputeTaskS2 || t <: ComputeTaskV) # double input
         @assert length(node.children) == 2
-        symbolIn1 =
-            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
-        symbolIn2 =
-            Symbol("data_$(replace(string(node.children[2].id), "-"=>"_"))")
-        symbolOut = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        symbolIn1 = Symbol("data_$(to_var_name(node.children[1].id))")
+        symbolIn2 = Symbol("data_$(to_var_name(node.children[2].id))")
+        symbolOut = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), symbolIn1, symbolIn2, symbolOut)
     elseif (t <: ComputeTaskSum) # vector input
         @assert length(node.children) > 0
         inSymbols = Vector{Symbol}()
         for child in node.children
-            push!(
-                inSymbols,
-                Symbol("data_$(replace(string(child.id), "-"=>"_"))"),
-            )
+            push!(inSymbols, Symbol("data_$(to_var_name(child.id))"))
         end
-        outSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+        outSymbol = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), inSymbols, outSymbol)
     elseif (t <: FusedComputeTask)
         # uuuuuh
@@ -118,12 +112,11 @@ function get_expression(node::DataTaskNode)
 
     inSymbol = nothing
     if (length(node.children) == 1)
-        inSymbol =
-            Symbol("data_$(replace(string(node.children[1].id), "-"=>"_"))")
+        inSymbol = Symbol("data_$(to_var_name(node.children[1].id))")
     else
-        inSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))_in")
+        inSymbol = Symbol("data_$(to_var_name(node.id))_in")
     end
-    outSymbol = Symbol("data_$(replace(string(node.id), "-"=>"_"))")
+    outSymbol = Symbol("data_$(to_var_name(node.id))")
 
     dataTransportExp = Meta.parse("$outSymbol = $inSymbol")
 
diff --git a/src/models/abc/create.jl b/src/models/abc/create.jl
new file mode 100644
index 0000000..58b11bb
--- /dev/null
+++ b/src/models/abc/create.jl
@@ -0,0 +1,30 @@
+
+"""
+    Particle(rng)
+
+Return a randomly generated particle.
+"""
+function Particle(rng)
+    return Particle(
+        rand(rng, Float64),
+        rand(rng, Float64),
+        rand(rng, Float64),
+        rand(rng, Float64),
+        rand(rng, Float64),
+    )
+end
+
+"""
+    gen_particles(n::Int)
+
+Return a Vector of `n` randomly generated [`Particle`](@ref)s.
+"""
+function gen_particles(n::Int)
+    particles = Vector{Particle}()
+    sizehint!(particles, n)
+    rng = MersenneTwister(0)
+    for i in 1:n
+        push!(particles, Particle(rng))
+    end
+    return particles
+end
diff --git a/src/node/print.jl b/src/node/print.jl
index 3a6ee1a..c39c1b5 100644
--- a/src/node/print.jl
+++ b/src/node/print.jl
@@ -15,3 +15,12 @@ Print a short string representation of the edge to io.
 function show(io::IO, e::Edge)
     return print(io, "Edge(", e.edge[1], ", ", e.edge[2], ")")
 end
+
+"""
+    to_var_name(id::UUID)
+
+Return the uuid as a string usable as a variable name in code generation.
+"""
+function to_var_name(id::UUID)
+    return replace(string(id), "-" => "_")
+end

From d1666de43271c7855d792d481930a4eab59b39f6 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@wandelbots.com>
Date: Thu, 7 Sep 2023 15:15:21 +0200
Subject: [PATCH 4/8] Add accurate arithmetic for summation, fix order of input
 particles

---
 Project.toml                    |  1 +
 examples/profiling_utilities.jl | 23 ++++++++++++
 src/code_gen/main.jl            | 38 ++++++++++++++------
 src/graph/mute.jl               |  3 ++
 src/graph/print.jl              |  1 +
 src/models/abc/compute.jl       | 62 ++++++++++++++++++++++++++++-----
 src/models/abc/create.jl        | 31 ++++++++++-------
 src/models/abc/parse.jl         |  7 +++-
 src/models/abc/particle.jl      | 27 ++++++++++----
 src/models/abc/properties.jl    | 14 ++++++++
 src/node/create.jl              |  8 +++--
 src/node/type.jl                |  3 ++
 src/operation/apply.jl          | 23 ++++--------
 src/operation/clean.jl          |  8 +----
 14 files changed, 183 insertions(+), 66 deletions(-)

diff --git a/Project.toml b/Project.toml
index 5e8350a..a6000e3 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,7 @@ authors = ["Anton Reinhard <anton.reinhard@proton.me>"]
 version = "0.1.0"
 
 [deps]
+AccurateArithmetic = "22286c92-06ac-501d-9306-4abd417d9753"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
diff --git a/examples/profiling_utilities.jl b/examples/profiling_utilities.jl
index 6aa7727..eb99b88 100644
--- a/examples/profiling_utilities.jl
+++ b/examples/profiling_utilities.jl
@@ -34,3 +34,26 @@ function test_random_walk(g::DAG, n::Int64)
 
     return reset_graph!(g)
 end
+
+function reduce_all!(g::DAG)
+    reset_graph!(g)
+
+    opt = get_operations(g)
+    while (!isempty(opt.nodeReductions))
+        push_operation!(g, pop!(opt.nodeReductions))
+
+        if (isempty(opt.nodeReductions))
+            opt = get_operations(g)
+        end
+    end
+    return nothing
+end
+
+function reduce_one!(g::DAG)
+    opt = get_operations(g)
+    if !isempty(opt.nodeReductions)
+        push_operation!(g, pop!(opt.nodeReductions))
+    end
+    opt = get_operations(g)
+    return nothing
+end
diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index e392510..f43d185 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -5,12 +5,12 @@ function gen_code(graph::DAG)
     sizehint!(code, length(graph.nodes))
 
     nodeQueue = PriorityQueue{Node, Int}()
-    inputSyms = Vector{Symbol}()
+    inputSyms = Dict{String, Symbol}()
 
     # use a priority equal to the number of unseen children -> 0 are nodes that can be added
     for node in get_entry_nodes(graph)
         enqueue!(nodeQueue, node => 0)
-        push!(inputSyms, Symbol("data_$(to_var_name(node.id))_in"))
+        push!(inputSyms, node.name => Symbol("data_$(to_var_name(node.id))_in"))
     end
 
     node = nothing
@@ -39,16 +39,26 @@ function gen_code(graph::DAG)
     )
 end
 
-function execute(generated_code, input::Vector{Particle})
+function execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = generated_code
     @assert length(input) == length(inputSymbols)
 
     assignInputs = Vector{Expr}()
-    for i in 1:length(input)
+    for (name, symbol) in inputSymbols
+        type = nothing
+        if startswith("A", name)
+            type = A
+        elseif startswith("B", name)
+            type = B
+        else
+            type = C
+        end
+        index = parse(Int, name[2:end])
+
         push!(
             assignInputs,
             Meta.parse(
-                "$(inputSymbols[i]) = ParticleValue(Particle($(input[i]).P0, $(input[i]).P1, $(input[i]).P2, $(input[i]).P3, $(input[i]).m), 1.0)",
+                "$(symbol) = ParticleValue(Particle($(input[type][index]).P0, $(input[type][index]).P1, $(input[type][index]).P2, $(input[type][index]).P3, $(type)), 1.0)",
             ),
         )
     end
@@ -61,17 +71,25 @@ function execute(generated_code, input::Vector{Particle})
     return result
 end
 
-function execute(graph::DAG, input::Vector{Particle})
+function execute(graph::DAG, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = gen_code(graph)
 
-    @assert length(input) == length(inputSymbols)
-
     assignInputs = Vector{Expr}()
-    for i in 1:length(input)
+    for (name, symbol) in inputSymbols
+        type = nothing
+        if startswith(name, "A")
+            type = A
+        elseif startswith(name, "B")
+            type = B
+        else
+            type = C
+        end
+        index = parse(Int, name[2:end])
+
         push!(
             assignInputs,
             Meta.parse(
-                "$(inputSymbols[i]) = ParticleValue(Particle($(input[i]).P0, $(input[i]).P1, $(input[i]).P2, $(input[i]).P3, $(input[i]).m), 1.0)",
+                "$(symbol) = ParticleValue(Particle($(input[type][index]).P0, $(input[type][index]).P1, $(input[type][index]).P2, $(input[type][index]).P3, $(type)), 1.0)",
             ),
         )
     end
diff --git a/src/graph/mute.jl b/src/graph/mute.jl
index e030fbd..77fd83d 100644
--- a/src/graph/mute.jl
+++ b/src/graph/mute.jl
@@ -144,6 +144,8 @@ function remove_edge!(
     # 1: mute
     pre_length1 = length(node1.parents)
     pre_length2 = length(node2.children)
+
+    #TODO: filter is very slow
     filter!(x -> x != node2, node1.parents)
     filter!(x -> x != node1, node2.children)
 
@@ -201,6 +203,7 @@ function invalidate_caches!(graph::DAG, operation::NodeFusion)
     delete!(graph.possibleOperations, operation)
 
     # delete the operation from all caches of nodes involved in the operation
+    # TODO: filter is very slow
     filter!(!=(operation), operation.input[1].nodeFusions)
     filter!(!=(operation), operation.input[3].nodeFusions)
 
diff --git a/src/graph/print.jl b/src/graph/print.jl
index 73df06c..c7e66e4 100644
--- a/src/graph/print.jl
+++ b/src/graph/print.jl
@@ -23,6 +23,7 @@ end
 Print the given graph to io. If there are too many nodes it will print only a summary of them.
 """
 function show(io::IO, graph::DAG)
+    apply_all!(graph)
     println(io, "Graph:")
     print(io, "  Nodes: ")
 
diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
index e0143dd..6e08a1b 100644
--- a/src/models/abc/compute.jl
+++ b/src/models/abc/compute.jl
@@ -1,3 +1,4 @@
+using AccurateArithmetic
 
 # Compute Particle, nothing to be done (0 FLOP)
 function compute(::ComputeTaskP, data::ParticleValue)
@@ -40,7 +41,8 @@ function get_expression(
 end
 
 # compute final inner edge (no output particle)
-function compute(::ComputeTaskS2, data1, data2)
+function compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
+    # data1 and data2 particles should be equal in a calculation with valid inputs, so it doesn't matter which one is used for inner_edge()
     return data1.v * inner_edge(data1.p) * data2.v
 end
 
@@ -56,7 +58,7 @@ function get_expression(
 end
 
 # compute inner edge
-function compute(::ComputeTaskS1, data)
+function compute(::ComputeTaskS1, data::ParticleValue)
     return ParticleValue(data.p, data.v * inner_edge(data.p))
 end
 
@@ -65,7 +67,8 @@ function get_expression(::ComputeTaskS1, inSymbol::Symbol, outSymbol::Symbol)
 end
 
 function compute(::ComputeTaskSum, data::Vector{Float64})
-    return sum(data)
+    # use an error correcting sum since the vectors may get very large
+    return sum_kbn(data)
 end
 
 function get_expression(
@@ -78,29 +81,70 @@ function get_expression(
     end
 end
 
+function compute(t::FusedComputeTask, data)
+    @assert false "This is not implemented and should never be called"
+end
+
+# expects the inSymbols ordered
+function get_expression(
+    t::FusedComputeTask,
+    inSymbols::Vector{Symbol},
+    outSymbol::Symbol,
+)
+    (T1, T2) = get_types(t)
+    c1 = children(T1())
+    c2 = children(T2())
+
+    expr1 = nothing
+    expr2 = nothing
+
+    # TODO need to figure out how to know which inputs belong to which subtask
+    # since we order the vectors with the child nodes we can't just split
+    if (c1 == 1)
+        expr1 = get_expression(T1(), inSymbols[begin], :intermediate)
+    elseif (c1 == 2)
+        expr1 =
+            get_expression(T1(), inSymbols[begin], inSymbols[2], :intermediate)
+    else
+        expr1 = get_expression(T1(), inSymbols[begin:c1], :intermediate)
+    end
+
+    if (c2 == 1)
+        expr2 = get_expression(T2(), :intermediate, outSymbol)
+    elseif c2 == 2
+        expr2 =
+            get_expression(T2(), :intermediate, inSymbols[c1 + 1], outSymbol)
+    else
+        expr2 = get_expression(
+            T2(),
+            :intermediate * inSymbols[(c1 + 1):end],
+            outSymbol,
+        )
+    end
+
+    return Expr(:block, expr1, expr2)
+end
+
 function get_expression(node::ComputeTaskNode)
     t = typeof(node.task)
+    @assert length(node.children) == children(node.task) || t <: ComputeTaskSum
+
     if (t <: ComputeTaskU || t <: ComputeTaskP || t <: ComputeTaskS1) # single input
-        @assert length(node.children) == 1
         symbolIn = Symbol("data_$(to_var_name(node.children[1].id))")
         symbolOut = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), symbolIn, symbolOut)
     elseif (t <: ComputeTaskS2 || t <: ComputeTaskV) # double input
-        @assert length(node.children) == 2
         symbolIn1 = Symbol("data_$(to_var_name(node.children[1].id))")
         symbolIn2 = Symbol("data_$(to_var_name(node.children[2].id))")
         symbolOut = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), symbolIn1, symbolIn2, symbolOut)
-    elseif (t <: ComputeTaskSum) # vector input
-        @assert length(node.children) > 0
+    elseif (t <: ComputeTaskSum || t <: FusedComputeTask) # vector input
         inSymbols = Vector{Symbol}()
         for child in node.children
             push!(inSymbols, Symbol("data_$(to_var_name(child.id))"))
         end
         outSymbol = Symbol("data_$(to_var_name(node.id))")
         return get_expression(t(), inSymbols, outSymbol)
-    elseif (t <: FusedComputeTask)
-        # uuuuuh
     else
         error("Unknown compute task")
     end
diff --git a/src/models/abc/create.jl b/src/models/abc/create.jl
index 58b11bb..d5902ff 100644
--- a/src/models/abc/create.jl
+++ b/src/models/abc/create.jl
@@ -4,14 +4,16 @@
 
 Return a randomly generated particle.
 """
-function Particle(rng)
-    return Particle(
-        rand(rng, Float64),
-        rand(rng, Float64),
-        rand(rng, Float64),
-        rand(rng, Float64),
-        rand(rng, Float64),
-    )
+function Particle(rng, type::ParticleType)
+
+    p1 = rand(rng, Float64)
+    p2 = rand(rng, Float64)
+    p3 = rand(rng, Float64)
+    m = mass(type)
+
+    p4 = sqrt(p1^2 + p2^2 + p3^2 + m^2)
+
+    return Particle(p1, p2, p3, p4, type)
 end
 
 """
@@ -19,12 +21,15 @@ end
 
 Return a Vector of `n` randomly generated [`Particle`](@ref)s.
 """
-function gen_particles(n::Int)
-    particles = Vector{Particle}()
-    sizehint!(particles, n)
+function gen_particles(ns::Dict{ParticleType, Int})
+    particles = Dict{ParticleType, Vector{Particle}}()
+
     rng = MersenneTwister(0)
-    for i in 1:n
-        push!(particles, Particle(rng))
+    for (type, n) in ns
+        particles[type] = Vector{Particle}()
+        for i in 1:n
+            push!(particles[type], Particle(rng, type))
+        end
     end
     return particles
 end
diff --git a/src/models/abc/parse.jl b/src/models/abc/parse.jl
index d78b23b..9afc23b 100644
--- a/src/models/abc/parse.jl
+++ b/src/models/abc/parse.jl
@@ -86,7 +86,12 @@ function parse_abc(filename::String, verbose::Bool = false)
         end
         if occursin(regex_a, node)
             # add nodes and edges for the state reading to u(P(Particle))
-            data_in = insert_node!(graph, make_node(DataTask(4)), false, false) # read particle data node
+            data_in = insert_node!(
+                graph,
+                make_node(DataTask(4), string(node)),
+                false,
+                false,
+            ) # read particle data node
             compute_P =
                 insert_node!(graph, make_node(ComputeTaskP()), false, false) # compute P node
             data_Pu = insert_node!(graph, make_node(DataTask(6)), false, false) # transfer data from P to u
diff --git a/src/models/abc/particle.jl b/src/models/abc/particle.jl
index e336de6..3f64be4 100644
--- a/src/models/abc/particle.jl
+++ b/src/models/abc/particle.jl
@@ -1,3 +1,7 @@
+@enum ParticleType A = 1 B = 2 C = 3 ALL = 6
+
+const PARTICLE_MASSES =
+    Dict{ParticleType, Float64}(A => 1.0, B => 1.0, C => 0.0)
 
 struct Particle
     P0::Float64
@@ -5,7 +9,7 @@ struct Particle
     P2::Float64
     P3::Float64
 
-    m::Float64
+    type::ParticleType
 end
 
 struct ParticleValue
@@ -13,12 +17,25 @@ struct ParticleValue
     v::Float64
 end
 
+mass(t::ParticleType) = PARTICLE_MASSES[t]
+
+function remaining_type(t1::ParticleType, t2::ParticleType)
+    @assert t1 != t2
+    if t1 != A && t2 != A
+        return A
+    elseif t1 != B && t2 != B
+        return B
+    else
+        return C
+    end
+end
+
 function square(p::Particle)
     return p.P0 * p.P0 - p.P1 * p.P1 - p.P2 * p.P2 - p.P3 * p.P3
 end
 
 function inner_edge(p::Particle)
-    return 1.0 / (square(p) - p.m * p.m)
+    return 1.0 / (square(p) - mass(p.type) * mass(p.type))
 end
 
 function outer_edge(p::Particle)
@@ -33,17 +50,13 @@ end
 
 # calculate new particle from two given interacting ones
 function preserve_momentum(p1::Particle, p2::Particle)
-    # TODO: is this correct?
-
     p3 = Particle(
         p1.P0 + p2.P0,
         p1.P1 + p2.P1,
         p1.P2 + p2.P2,
         p1.P3 + p2.P3,
-        1.0,
+        remaining_type(p1.type, p2.type),
     )
 
-    # m3 = sqrt(-  PC * PC / c^2)
-
     return p3
 end
diff --git a/src/models/abc/properties.jl b/src/models/abc/properties.jl
index fcc254b..f780a48 100644
--- a/src/models/abc/properties.jl
+++ b/src/models/abc/properties.jl
@@ -100,3 +100,17 @@ show(io::IO, t::ComputeTaskSum) = print("ComputeSum")
 Copy the data task and return it.
 """
 copy(t::DataTask) = DataTask(t.data)
+
+children(::DataTask) = 1
+children(::ComputeTaskS1) = 1
+children(::ComputeTaskS2) = 2
+children(::ComputeTaskP) = 1
+children(::ComputeTaskU) = 1
+children(::ComputeTaskV) = 2
+
+# TODO: this is kind of bad because it means we can't fuse with a sum task
+children(::ComputeTaskSum) = -1 # wildcard for "n" children
+function children(t::FusedComputeTask)
+    (T1, T2) = get_types(t)
+    return children(T1()) + children(T2()) - 1 # one of the inputs is the output of T1 and thus not a child of the node
+end
diff --git a/src/node/create.jl b/src/node/create.jl
index 127f3c2..1335eaa 100644
--- a/src/node/create.jl
+++ b/src/node/create.jl
@@ -1,5 +1,5 @@
 
-DataTaskNode(t::AbstractDataTask) = DataTaskNode(
+DataTaskNode(t::AbstractDataTask, name = "") = DataTaskNode(
     t,
     Vector{Node}(),
     Vector{Node}(),
@@ -7,6 +7,7 @@ DataTaskNode(t::AbstractDataTask) = DataTaskNode(
     missing,
     missing,
     missing,
+    name,
 )
 ComputeTaskNode(t::AbstractComputeTask) = ComputeTaskNode(
     t,
@@ -36,6 +37,7 @@ copy(n::DataTaskNode) = DataTaskNode(
     copy(n.nodeReduction),
     copy(n.nodeSplit),
     copy(n.nodeFusion),
+    copy(n.name),
 )
 
 """
@@ -52,8 +54,8 @@ end
 
 Construct and return a new [`DataTaskNode`](@ref) with the given task.
 """
-function make_node(t::AbstractDataTask)
-    return DataTaskNode(t)
+function make_node(t::AbstractDataTask, name::String = "")
+    return DataTaskNode(t, name)
 end
 
 """
diff --git a/src/node/type.jl b/src/node/type.jl
index e8a353a..06bf308 100644
--- a/src/node/type.jl
+++ b/src/node/type.jl
@@ -52,6 +52,9 @@ mutable struct DataTaskNode <: Node
 
     # the node fusion involving this node, if it exists
     nodeFusion::Union{Operation, Missing}
+
+    # for input nodes we need a name for the node to distinguish between them
+    name::String
 end
 
 """
diff --git a/src/operation/apply.jl b/src/operation/apply.jl
index b6da322..2974289 100644
--- a/src/operation/apply.jl
+++ b/src/operation/apply.jl
@@ -160,7 +160,6 @@ function node_fusion!(
     # clear snapshot
     get_snapshot_diff(graph)
 
-
     # save children and parents
     n1_children = children(n1)
     n3_parents = parents(n3)
@@ -181,26 +180,18 @@ function node_fusion!(
         ComputeTaskNode(FusedComputeTask{typeof(n1.task), typeof(n3.task)}())
     insert_node!(graph, new_node)
 
-    # use a set for combined children of n1 and n3 to not get duplicates
-    n1and3_children = Set{Node}()
-
-    # remove edges from n1 children to n1
     for child in n1_children
         remove_edge!(graph, child, n1)
-        push!(n1and3_children, child)
-    end
-
-    # remove edges from n3 children to n3
-    for child in n3_children
-        remove_edge!(graph, child, n3)
-        push!(n1and3_children, child)
-    end
-
-    for child in n1and3_children
         insert_edge!(graph, child, new_node)
     end
 
-    # "repoint" parents of n3 from new node
+    for child in n3_children
+        remove_edge!(graph, child, n3)
+        if !(child in n1_children)
+            insert_edge!(graph, child, new_node)
+        end
+    end
+
     for parent in n3_parents
         remove_edge!(graph, n3, parent)
         insert_edge!(graph, new_node, parent)
diff --git a/src/operation/clean.jl b/src/operation/clean.jl
index e42a2a3..70ebbd3 100644
--- a/src/operation/clean.jl
+++ b/src/operation/clean.jl
@@ -71,14 +71,8 @@ function find_reductions!(graph::DAG, node::Node)
     partners_ = partners(node)
     delete!(partners_, node)
     for partner in partners_
-        if partner ∉ graph.nodes
-            error("Partner is not part of the graph")
-        end
-
+        @assert partner in graph.nodes
         if can_reduce(node, partner)
-            if Set(node.children) != Set(partner.children)
-                error("Not equal children")
-            end
             if reductionVector === nothing
                 # only when there's at least one reduction partner, insert the vector
                 reductionVector = Vector{Node}()

From e59d24ebe5262f26784e55190048d55fcea09574 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@wandelbots.com>
Date: Thu, 7 Sep 2023 18:23:36 +0200
Subject: [PATCH 5/8] Add code gen documentation

---
 src/code_gen/main.jl         |  24 +++++++
 src/models/abc/compute.jl    | 136 +++++++++++++++++++++++++++--------
 src/models/abc/create.jl     |   3 +
 src/models/abc/particle.jl   |  58 ++++++++++++++-
 src/models/abc/properties.jl |  53 +++++++++++++-
 5 files changed, 240 insertions(+), 34 deletions(-)

diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index f43d185..e22b590 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -1,5 +1,16 @@
 using DataStructures
 
+"""
+    gen_code(graph::DAG)
+
+Generate the code for a given graph. The return value is a tuple of:
+
+    - `code::Expr`: The julia expression containing the code for the whole graph.
+    - `inputSymbols::Dict{String, Symbol}`: A dictionary of symbols mapping the names of the input nodes of the graph to the symbols their inputs should be provided on.
+    - `outputSymbol::Symbol`: The symbol of the final calculated value
+
+See also: [`execute`](@ref)
+"""
 function gen_code(graph::DAG)
     code = Vector{Expr}()
     sizehint!(code, length(graph.nodes))
@@ -39,6 +50,11 @@ function gen_code(graph::DAG)
     )
 end
 
+"""
+    execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
+
+Execute the given generated_code (as returned by [`gen_code`](@ref)) on the given input particles.
+"""
 function execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = generated_code
     @assert length(input) == length(inputSymbols)
@@ -71,6 +87,14 @@ function execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
     return result
 end
 
+"""
+    execute(graph::DAG, input::Dict{ParticleType, Vector{Particle}})
+
+Execute the given generated_code (as returned by [`gen_code`](@ref)) on the given input particles.
+The input particles should be sorted correctly into the dictionary to their according [`ParticleType`](@ref)s.
+
+See also: [`gen_particles`](@ref)
+"""
 function execute(graph::DAG, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = gen_code(graph)
 
diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
index 6e08a1b..3fd3250 100644
--- a/src/models/abc/compute.jl
+++ b/src/models/abc/compute.jl
@@ -1,34 +1,97 @@
 using AccurateArithmetic
 
-# Compute Particle, nothing to be done (0 FLOP)
+"""
+    compute(::ComputeTaskP, data::ParticleValue)
+
+Return the particle and value as is. 0 FLOP.
+"""
 function compute(::ComputeTaskP, data::ParticleValue)
     return data
 end
 
-# generate code evaluating ComputeTaskP on inSymbol, providing the output on outSymbol
-function get_expression(::ComputeTaskP, inSymbol::Symbol, outSymbol::Symbol)
-    return Meta.parse("$outSymbol = compute(ComputeTaskP(), $inSymbol)")
-end
+"""
+    compute(::ComputeTaskU, data::ParticleValue)
 
-# Compute outer edge
+Compute an outer edge. Return the particle value with the same particle and the value multiplied by an outer_edge factor.
+"""
 function compute(::ComputeTaskU, data::ParticleValue)
     return ParticleValue(data.p, data.v * outer_edge(data.p))
 end
 
-# generate code evaluating ComputeTaskU on inSymbol, providing the output on outSymbol
-# inSymbol should be of type ParticleValue, outSymbol will be of type ParticleValue
-function get_expression(::ComputeTaskU, inSymbol::Symbol, outSymbol::Symbol)
-    return Meta.parse("$outSymbol = compute(ComputeTaskU(), $inSymbol)")
-end
+"""
+    compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
 
-# compute vertex
+Compute a vertex. Preserve momentum and particle types (AB->C etc.) to create resulting particle, multiply values together and times a vertex factor.
+"""
 function compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
-    # calculate new particle from the two input particles
     p3 = preserve_momentum(data1.p, data2.p)
     dataOut = ParticleValue(p3, data1.v * vertex() * data2.v)
     return dataOut
 end
 
+"""
+    compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
+
+Compute a final inner edge (2 input particles, no output particle).
+
+For valid inputs, both input particles should have the same momenta at this point.
+"""
+function compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
+    return data1.v * inner_edge(data1.p) * data2.v
+end
+
+"""
+    compute(::ComputeTaskS1, data::ParticleValue)
+
+Compute inner edge (1 input particle, 1 output particle).
+"""
+function compute(::ComputeTaskS1, data::ParticleValue)
+    return ParticleValue(data.p, data.v * inner_edge(data.p))
+end
+
+"""
+    compute(::ComputeTaskSum, data::Vector{Float64})
+
+Compute a sum over the vector. Use an algorithm that accounts for accumulated errors in long sums with potentially large differences in magnitude of the summands.
+"""
+function compute(::ComputeTaskSum, data::Vector{Float64})
+    return sum_kbn(data)
+end
+
+"""
+    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(::ComputeTaskP, inSymbol::Symbol, outSymbol::Symbol)
+
+Generate and return code evaluating [`ComputeTaskP`](@ref) on `inSymbol`, providing the output on `outSymbol`.
+"""
+function get_expression(::ComputeTaskP, inSymbol::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskP(), $inSymbol)")
+end
+
+"""
+    get_expression(::ComputeTaskU, inSymbol::Symbol, outSymbol::Symbol)
+
+Generate code evaluating [`ComputeTaskU`](@ref) on `inSymbol`, providing the output on `outSymbol`.
+`inSymbol` should be of type [`ParticleValue`](@ref), `outSymbol` will be of type [`ParticleValue`](@ref).
+"""
+function get_expression(::ComputeTaskU, inSymbol::Symbol, outSymbol::Symbol)
+    return Meta.parse("$outSymbol = compute(ComputeTaskU(), $inSymbol)")
+end
+
+"""
+    get_expression(::ComputeTaskV, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
+
+Generate code evaluating [`ComputeTaskV`](@ref) on `inSymbol1` and `inSymbol2`, providing the output on `outSymbol`.
+`inSymbol1` and `inSymbol2` should be of type [`ParticleValue`](@ref), `outSymbol` will be of type [`ParticleValue`](@ref).
+"""
 function get_expression(
     ::ComputeTaskV,
     inSymbol1::Symbol,
@@ -40,12 +103,12 @@ function get_expression(
     )
 end
 
-# compute final inner edge (no output particle)
-function compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
-    # data1 and data2 particles should be equal in a calculation with valid inputs, so it doesn't matter which one is used for inner_edge()
-    return data1.v * inner_edge(data1.p) * data2.v
-end
+"""
+    get_expression(::ComputeTaskS2, inSymbol1::Symbol, inSymbol2::Symbol, outSymbol::Symbol)
 
+Generate code evaluating [`ComputeTaskS2`](@ref) on `inSymbol1` and `inSymbol2`, providing the output on `outSymbol`.
+`inSymbol1` and `inSymbol2` should be of type [`ParticleValue`](@ref), `outSymbol` will be of type `Float64`.
+"""
 function get_expression(
     ::ComputeTaskS2,
     inSymbol1::Symbol,
@@ -57,20 +120,22 @@ function get_expression(
     )
 end
 
-# compute inner edge
-function compute(::ComputeTaskS1, data::ParticleValue)
-    return ParticleValue(data.p, data.v * inner_edge(data.p))
-end
+"""
+    get_expression(::ComputeTaskS1, inSymbol::Symbol, outSymbol::Symbol)
 
+Generate code evaluating [`ComputeTaskS1`](@ref) on `inSymbol`, providing the output on `outSymbol`.
+`inSymbol` should be of type [`ParticleValue`](@ref), `outSymbol` will be of type [`ParticleValue`](@ref).
+"""
 function get_expression(::ComputeTaskS1, inSymbol::Symbol, outSymbol::Symbol)
     return Meta.parse("$outSymbol = compute(ComputeTaskS1(), $inSymbol)")
 end
 
-function compute(::ComputeTaskSum, data::Vector{Float64})
-    # use an error correcting sum since the vectors may get very large
-    return sum_kbn(data)
-end
+"""
+    get_expression(::ComputeTaskSum, inSymbols::Vector{Symbol}, outSymbol::Symbol)
 
+Generate code evaluating [`ComputeTaskSum`](@ref) on `inSymbols`, providing the output on `outSymbol`.
+`inSymbols` should be of type [`Float64`], `outSymbol` will be of type [`Float64`].
+"""
 function get_expression(
     ::ComputeTaskSum,
     inSymbols::Vector{Symbol},
@@ -81,11 +146,12 @@ function get_expression(
     end
 end
 
-function compute(t::FusedComputeTask, data)
-    @assert false "This is not implemented and should never be called"
-end
+"""
+    get_expression(t::FusedComputeTask, inSymbols::Vector{Symbol}, outSymbol::Symbol)
 
-# expects the inSymbols ordered
+Generate code evaluating a [`FusedComputeTask`](@ref) on `inSymbols`, providing the output on `outSymbol`.
+`inSymbols` should be of the correct types and may be heterogeneous. `outSymbol` will be of the type of the output of `T2` of t.
+"""
 function get_expression(
     t::FusedComputeTask,
     inSymbols::Vector{Symbol},
@@ -125,6 +191,11 @@ function get_expression(
     return Expr(:block, expr1, expr2)
 end
 
+"""
+    get_expression(node::ComputeTaskNode)
+
+Generate and return code for a given [`ComputeTaskNode`](@ref).
+"""
 function get_expression(node::ComputeTaskNode)
     t = typeof(node.task)
     @assert length(node.children) == children(node.task) || t <: ComputeTaskSum
@@ -150,6 +221,11 @@ function get_expression(node::ComputeTaskNode)
     end
 end
 
+"""
+    get_expression(node::DataTaskNode)
+
+Generate and return code for a given [`DataTaskNode`](@ref).
+"""
 function get_expression(node::DataTaskNode)
     # TODO: do things to transport data from/to gpu, between numa nodes, etc.
     @assert length(node.children) <= 1
diff --git a/src/models/abc/create.jl b/src/models/abc/create.jl
index d5902ff..cf9e39d 100644
--- a/src/models/abc/create.jl
+++ b/src/models/abc/create.jl
@@ -11,6 +11,7 @@ function Particle(rng, type::ParticleType)
     p3 = rand(rng, Float64)
     m = mass(type)
 
+    # keep the momenta of the particles on-shell
     p4 = sqrt(p1^2 + p2^2 + p3^2 + m^2)
 
     return Particle(p1, p2, p3, p4, type)
@@ -20,6 +21,8 @@ end
     gen_particles(n::Int)
 
 Return a Vector of `n` randomly generated [`Particle`](@ref)s.
+
+Note: This does not take into account the preservation of momenta required for an actual valid process!
 """
 function gen_particles(ns::Dict{ParticleType, Int})
     particles = Dict{ParticleType, Vector{Particle}}()
diff --git a/src/models/abc/particle.jl b/src/models/abc/particle.jl
index 3f64be4..7b125d3 100644
--- a/src/models/abc/particle.jl
+++ b/src/models/abc/particle.jl
@@ -1,8 +1,23 @@
-@enum ParticleType A = 1 B = 2 C = 3 ALL = 6
+"""
+    ParticleType
 
+A Particle Type in the ABC Model as an enum, with types `A`, `B` and `C`.
+"""
+@enum ParticleType A = 1 B = 2 C = 3
+
+"""
+    PARTICLE_MASSES
+
+A constant dictionary containing the masses of the different [`ParticleType`](@ref)s.
+"""
 const PARTICLE_MASSES =
     Dict{ParticleType, Float64}(A => 1.0, B => 1.0, C => 0.0)
 
+"""
+    Particle
+
+A struct describing a particle of the ABC-Model. It has the 4 momentum parts P0...P3 and a [`ParticleType`](@ref).
+"""
 struct Particle
     P0::Float64
     P1::Float64
@@ -12,13 +27,28 @@ struct Particle
     type::ParticleType
 end
 
+"""
+    ParticleValue
+
+A struct describing a particle during a calculation of a Feynman Diagram, together with the value that's being calculated.
+"""
 struct ParticleValue
     p::Particle
     v::Float64
 end
 
+"""
+    mass(t::ParticleType)
+    
+Return the mass (at rest) of the given particle type.
+"""
 mass(t::ParticleType) = PARTICLE_MASSES[t]
 
+"""
+    remaining_type(t1::ParticleType, t2::ParticleType)
+
+For 2 given (non-equal) particle types, return the third of ABC.
+"""
 function remaining_type(t1::ParticleType, t2::ParticleType)
     @assert t1 != t2
     if t1 != A && t2 != A
@@ -30,25 +60,49 @@ function remaining_type(t1::ParticleType, t2::ParticleType)
     end
 end
 
+"""
+    square(p::Particle)
+
+Return the square of the particle's momentum as a `Float` value.
+"""
 function square(p::Particle)
     return p.P0 * p.P0 - p.P1 * p.P1 - p.P2 * p.P2 - p.P3 * p.P3
 end
 
+"""
+    inner_edge(p::Particle)
+
+Return the factor of the inner edge with the given (virtual) particle.
+"""
 function inner_edge(p::Particle)
     return 1.0 / (square(p) - mass(p.type) * mass(p.type))
 end
 
+"""
+    outer_edge(p::Particle)
+
+Return the factor of the outer edge with the given (real) particle.
+"""
 function outer_edge(p::Particle)
     return 1.0
 end
 
+"""
+    vertex()
+
+Return the factor of a vertex.
+"""
 function vertex()
     i = 1.0
     lambda = 1.0 / 137.0
     return i * lambda
 end
 
-# calculate new particle from two given interacting ones
+"""
+    preserve_momentum(p1::Particle, p2::Particle)
+
+Calculate and return a new particle from two given interacting ones at a vertex.
+"""
 function preserve_momentum(p1::Particle, p2::Particle)
     p3 = Particle(
         p1.P0 + p2.P0,
diff --git a/src/models/abc/properties.jl b/src/models/abc/properties.jl
index f780a48..73b6367 100644
--- a/src/models/abc/properties.jl
+++ b/src/models/abc/properties.jl
@@ -101,15 +101,64 @@ Copy the data task and return it.
 """
 copy(t::DataTask) = DataTask(t.data)
 
+"""
+    children(::DataTask)
+
+Return the number of children of a data task (always 1).
+"""
 children(::DataTask) = 1
+
+"""
+    children(::ComputeTaskS1)
+
+Return the number of children of a ComputeTaskS1 (always 1).
+"""
 children(::ComputeTaskS1) = 1
+
+"""
+    children(::ComputeTaskS2)
+
+Return the number of children of a ComputeTaskS2 (always 2).
+"""
 children(::ComputeTaskS2) = 2
+
+"""
+    children(::ComputeTaskP)
+
+Return the number of children of a ComputeTaskP (always 1).
+"""
 children(::ComputeTaskP) = 1
+
+"""
+    children(::ComputeTaskU)
+
+Return the number of children of a ComputeTaskU (always 1).
+"""
 children(::ComputeTaskU) = 1
+
+"""
+    children(::ComputeTaskV)
+
+Return the number of children of a ComputeTaskV (always 2).
+"""
 children(::ComputeTaskV) = 2
 
-# TODO: this is kind of bad because it means we can't fuse with a sum task
-children(::ComputeTaskSum) = -1 # wildcard for "n" children
+
+"""
+    children(::ComputeTaskSum)
+
+Return the number of children of a ComputeTaskSum, since this is variable and the task doesn't know
+how many children it will sum over, return a wildcard -1.
+
+TODO: this is kind of bad because it means we can't fuse with a sum task
+"""
+children(::ComputeTaskSum) = -1
+
+"""
+    children(t::FusedComputeTask)
+
+Return the number of children of a FusedComputeTask. It's the sum of the children of both tasks minus one.
+"""
 function children(t::FusedComputeTask)
     (T1, T2) = get_types(t)
     return children(T1()) + children(T2()) - 1 # one of the inputs is the output of T1 and thus not a child of the node

From 27c4b8ba34eab4b13446cad6628d0ff5e8585036 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@wandelbots.com>
Date: Thu, 7 Sep 2023 18:46:41 +0200
Subject: [PATCH 6/8] Use real ABC-Model compute effort and data transfer
 numbers

---
 src/models/abc/compute.jl    | 14 ++++++++-
 src/models/abc/parse.jl      | 60 ++++++++++++++++++++++++++++--------
 src/models/abc/particle.jl   | 14 +++++++++
 src/models/abc/properties.jl | 10 +++---
 src/node/create.jl           |  2 +-
 test/unit_tests_graph.jl     | 12 ++++----
 test/unit_tests_tasks.jl     | 10 +++---
 7 files changed, 92 insertions(+), 30 deletions(-)

diff --git a/src/models/abc/compute.jl b/src/models/abc/compute.jl
index 3fd3250..03fc82a 100644
--- a/src/models/abc/compute.jl
+++ b/src/models/abc/compute.jl
@@ -3,7 +3,9 @@ using AccurateArithmetic
 """
     compute(::ComputeTaskP, data::ParticleValue)
 
-Return the particle and value as is. 0 FLOP.
+Return the particle and value as is. 
+
+0 FLOP.
 """
 function compute(::ComputeTaskP, data::ParticleValue)
     return data
@@ -13,6 +15,8 @@ end
     compute(::ComputeTaskU, data::ParticleValue)
 
 Compute an outer edge. Return the particle value with the same particle and the value multiplied by an outer_edge factor.
+
+1 FLOP.
 """
 function compute(::ComputeTaskU, data::ParticleValue)
     return ParticleValue(data.p, data.v * outer_edge(data.p))
@@ -22,6 +26,8 @@ end
     compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
 
 Compute a vertex. Preserve momentum and particle types (AB->C etc.) to create resulting particle, multiply values together and times a vertex factor.
+
+6 FLOP.
 """
 function compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
     p3 = preserve_momentum(data1.p, data2.p)
@@ -35,6 +41,8 @@ end
 Compute a final inner edge (2 input particles, no output particle).
 
 For valid inputs, both input particles should have the same momenta at this point.
+
+12 FLOP.
 """
 function compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
     return data1.v * inner_edge(data1.p) * data2.v
@@ -44,6 +52,8 @@ end
     compute(::ComputeTaskS1, data::ParticleValue)
 
 Compute inner edge (1 input particle, 1 output particle).
+
+11 FLOP.
 """
 function compute(::ComputeTaskS1, data::ParticleValue)
     return ParticleValue(data.p, data.v * inner_edge(data.p))
@@ -53,6 +63,8 @@ end
     compute(::ComputeTaskSum, data::Vector{Float64})
 
 Compute a sum over the vector. Use an algorithm that accounts for accumulated errors in long sums with potentially large differences in magnitude of the summands.
+
+Linearly many FLOP with growing data.
 """
 function compute(::ComputeTaskSum, data::Vector{Float64})
     return sum_kbn(data)
diff --git a/src/models/abc/parse.jl b/src/models/abc/parse.jl
index 9afc23b..94e62f9 100644
--- a/src/models/abc/parse.jl
+++ b/src/models/abc/parse.jl
@@ -4,6 +4,9 @@ regex_c = r"^[A-C]\(([^']*),([^']*)\)$"     # Regex for the combinations of 2 pa
 regex_m = r"^M\(([^']*),([^']*),([^']*)\)$" # Regex for the combinations of 3 particles
 regex_plus = r"^\+$"                        # Regex for the sum
 
+const PARTICLE_VALUE_SIZE::Int = 48
+const FLOAT_SIZE::Int = 8
+
 """
     parse_nodes(input::AbstractString)
 
@@ -61,7 +64,8 @@ function parse_abc(filename::String, verbose::Bool = false)
     sizehint!(graph.nodes, estimate_no_nodes)
 
     sum_node = insert_node!(graph, make_node(ComputeTaskSum()), false, false)
-    global_data_out = insert_node!(graph, make_node(DataTask(10)), false, false)
+    global_data_out =
+        insert_node!(graph, make_node(DataTask(FLOAT_SIZE)), false, false)
     insert_edge!(graph, sum_node, global_data_out, false, false)
 
     # remember the data out nodes for connection
@@ -88,16 +92,26 @@ function parse_abc(filename::String, verbose::Bool = false)
             # add nodes and edges for the state reading to u(P(Particle))
             data_in = insert_node!(
                 graph,
-                make_node(DataTask(4), string(node)),
+                make_node(DataTask(PARTICLE_VALUE_SIZE), string(node)),
                 false,
                 false,
             ) # read particle data node
             compute_P =
                 insert_node!(graph, make_node(ComputeTaskP()), false, false) # compute P node
-            data_Pu = insert_node!(graph, make_node(DataTask(6)), false, false) # transfer data from P to u
+            data_Pu = insert_node!(
+                graph,
+                make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                false,
+                false,
+            ) # transfer data from P to u (one ParticleValue object)
             compute_u =
                 insert_node!(graph, make_node(ComputeTaskU()), false, false) # compute U node
-            data_out = insert_node!(graph, make_node(DataTask(3)), false, false) # transfer data out from u
+            data_out = insert_node!(
+                graph,
+                make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                false,
+                false,
+            ) # transfer data out from u (one ParticleValue object)
 
             insert_edge!(graph, data_in, compute_P, false, false)
             insert_edge!(graph, compute_P, data_Pu, false, false)
@@ -114,7 +128,12 @@ function parse_abc(filename::String, verbose::Bool = false)
 
             compute_v =
                 insert_node!(graph, make_node(ComputeTaskV()), false, false)
-            data_out = insert_node!(graph, make_node(DataTask(5)), false, false)
+            data_out = insert_node!(
+                graph,
+                make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                false,
+                false,
+            )
 
             if (occursin(regex_c, in1))
                 # put an S node after this input
@@ -124,8 +143,12 @@ function parse_abc(filename::String, verbose::Bool = false)
                     false,
                     false,
                 )
-                data_S_v =
-                    insert_node!(graph, make_node(DataTask(5)), false, false)
+                data_S_v = insert_node!(
+                    graph,
+                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                    false,
+                    false,
+                )
 
                 insert_edge!(graph, dataOutNodes[in1], compute_S, false, false)
                 insert_edge!(graph, compute_S, data_S_v, false, false)
@@ -144,8 +167,12 @@ function parse_abc(filename::String, verbose::Bool = false)
                     false,
                     false,
                 )
-                data_S_v =
-                    insert_node!(graph, make_node(DataTask(5)), false, false)
+                data_S_v = insert_node!(
+                    graph,
+                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                    false,
+                    false,
+                )
 
                 insert_edge!(graph, dataOutNodes[in2], compute_S, false, false)
                 insert_edge!(graph, compute_S, data_S_v, false, false)
@@ -168,7 +195,12 @@ function parse_abc(filename::String, verbose::Bool = false)
             # in2 + in3 with a v
             compute_v =
                 insert_node!(graph, make_node(ComputeTaskV()), false, false)
-            data_v = insert_node!(graph, make_node(DataTask(5)), false, false)
+            data_v = insert_node!(
+                graph,
+                make_node(DataTask(PARTICLE_VALUE_SIZE)),
+                false,
+                false,
+            )
 
             insert_edge!(graph, dataOutNodes[in2], compute_v, false, false)
             insert_edge!(graph, dataOutNodes[in3], compute_v, false, false)
@@ -177,8 +209,12 @@ function parse_abc(filename::String, verbose::Bool = false)
             # combine with the v of the combined other input
             compute_S2 =
                 insert_node!(graph, make_node(ComputeTaskS2()), false, false)
-            data_out =
-                insert_node!(graph, make_node(DataTask(10)), false, false)
+            data_out = insert_node!(
+                graph,
+                make_node(DataTask(FLOAT_SIZE)),
+                false,
+                false,
+            ) # output of a S2 task is only a float
 
             insert_edge!(graph, data_v, compute_S2, false, false)
             insert_edge!(graph, dataOutNodes[in1], compute_S2, false, false)
diff --git a/src/models/abc/particle.jl b/src/models/abc/particle.jl
index 7b125d3..b6807b5 100644
--- a/src/models/abc/particle.jl
+++ b/src/models/abc/particle.jl
@@ -17,6 +17,8 @@ const PARTICLE_MASSES =
     Particle
 
 A struct describing a particle of the ABC-Model. It has the 4 momentum parts P0...P3 and a [`ParticleType`](@ref).
+
+`sizeof(Particle())` = 40 Byte
 """
 struct Particle
     P0::Float64
@@ -31,6 +33,8 @@ end
     ParticleValue
 
 A struct describing a particle during a calculation of a Feynman Diagram, together with the value that's being calculated.
+
+`sizeof(ParticleValue())` = 48 Byte
 """
 struct ParticleValue
     p::Particle
@@ -64,6 +68,8 @@ end
     square(p::Particle)
 
 Return the square of the particle's momentum as a `Float` value.
+
+Takes 7 effective FLOP.
 """
 function square(p::Particle)
     return p.P0 * p.P0 - p.P1 * p.P1 - p.P2 * p.P2 - p.P3 * p.P3
@@ -73,6 +79,8 @@ end
     inner_edge(p::Particle)
 
 Return the factor of the inner edge with the given (virtual) particle.
+
+Takes 10 effective FLOP. (3 here + 10 in square(p))
 """
 function inner_edge(p::Particle)
     return 1.0 / (square(p) - mass(p.type) * mass(p.type))
@@ -82,6 +90,8 @@ end
     outer_edge(p::Particle)
 
 Return the factor of the outer edge with the given (real) particle.
+
+Takes 0 effective FLOP.
 """
 function outer_edge(p::Particle)
     return 1.0
@@ -91,6 +101,8 @@ end
     vertex()
 
 Return the factor of a vertex.
+
+Takes 0 effective FLOP since it's constant.
 """
 function vertex()
     i = 1.0
@@ -102,6 +114,8 @@ end
     preserve_momentum(p1::Particle, p2::Particle)
 
 Calculate and return a new particle from two given interacting ones at a vertex.
+
+Takes 4 effective FLOP.
 """
 function preserve_momentum(p1::Particle, p2::Particle)
     p3 = Particle(
diff --git a/src/models/abc/properties.jl b/src/models/abc/properties.jl
index 73b6367..8e08e97 100644
--- a/src/models/abc/properties.jl
+++ b/src/models/abc/properties.jl
@@ -3,35 +3,35 @@
 
 Return the compute effort of an S1 task.
 """
-compute_effort(t::ComputeTaskS1) = 10
+compute_effort(t::ComputeTaskS1) = 11
 
 """
     compute_effort(t::ComputeTaskS2)
 
 Return the compute effort of an S2 task.
 """
-compute_effort(t::ComputeTaskS2) = 10
+compute_effort(t::ComputeTaskS2) = 12
 
 """
     compute_effort(t::ComputeTaskU)
 
 Return the compute effort of a U task.
 """
-compute_effort(t::ComputeTaskU) = 6
+compute_effort(t::ComputeTaskU) = 1
 
 """
     compute_effort(t::ComputeTaskV)
 
 Return the compute effort of a V task.
 """
-compute_effort(t::ComputeTaskV) = 20
+compute_effort(t::ComputeTaskV) = 6
 
 """
     compute_effort(t::ComputeTaskP)
 
 Return the compute effort of a P task.
 """
-compute_effort(t::ComputeTaskP) = 15
+compute_effort(t::ComputeTaskP) = 0
 
 """
     compute_effort(t::ComputeTaskSum)
diff --git a/src/node/create.jl b/src/node/create.jl
index 1335eaa..1d28169 100644
--- a/src/node/create.jl
+++ b/src/node/create.jl
@@ -37,7 +37,7 @@ copy(n::DataTaskNode) = DataTaskNode(
     copy(n.nodeReduction),
     copy(n.nodeSplit),
     copy(n.nodeFusion),
-    copy(n.name),
+    n.name,
 )
 
 """
diff --git a/test/unit_tests_graph.jl b/test/unit_tests_graph.jl
index 77dd84c..9c835e6 100644
--- a/test/unit_tests_graph.jl
+++ b/test/unit_tests_graph.jl
@@ -143,9 +143,9 @@ import MetagraphOptimization.partners
     nf = first(operations.nodeFusions)
 
     properties = get_properties(graph)
-    @test properties.computeEffort == 134
+    @test properties.computeEffort == 28
     @test properties.data == 62
-    @test properties.computeIntensity ≈ 134 / 62
+    @test properties.computeIntensity ≈ 28 / 62
     @test properties.noNodes == 26
     @test properties.noEdges == 25
 
@@ -169,9 +169,9 @@ import MetagraphOptimization.partners
     @test length(graph.dirtyNodes) != 0
     @test properties.noNodes == 24
     @test properties.noEdges == 23
-    @test properties.computeEffort == 134
+    @test properties.computeEffort == 28
     @test properties.data < 62
-    @test properties.computeIntensity > 134 / 62
+    @test properties.computeIntensity > 28 / 62
 
     operations = get_operations(graph)
     @test length(graph.dirtyNodes) == 0
@@ -208,9 +208,9 @@ import MetagraphOptimization.partners
     properties = get_properties(graph)
     @test properties.noNodes == 26
     @test properties.noEdges == 25
-    @test properties.computeEffort == 134
+    @test properties.computeEffort == 28
     @test properties.data == 62
-    @test properties.computeIntensity ≈ 134 / 62
+    @test properties.computeIntensity ≈ 28 / 62
 
     operations = get_operations(graph)
     @test length(operations) ==
diff --git a/test/unit_tests_tasks.jl b/test/unit_tests_tasks.jl
index c039ece..1bc5bc8 100644
--- a/test/unit_tests_tasks.jl
+++ b/test/unit_tests_tasks.jl
@@ -10,11 +10,11 @@
     Data10 = MetagraphOptimization.DataTask(10)
     Data20 = MetagraphOptimization.DataTask(20)
 
-    @test MetagraphOptimization.compute_effort(S1) == 10
-    @test MetagraphOptimization.compute_effort(S2) == 10
-    @test MetagraphOptimization.compute_effort(U) == 6
-    @test MetagraphOptimization.compute_effort(V) == 20
-    @test MetagraphOptimization.compute_effort(P) == 15
+    @test MetagraphOptimization.compute_effort(S1) == 11
+    @test MetagraphOptimization.compute_effort(S2) == 12
+    @test MetagraphOptimization.compute_effort(U) == 1
+    @test MetagraphOptimization.compute_effort(V) == 6
+    @test MetagraphOptimization.compute_effort(P) == 0
     @test MetagraphOptimization.compute_effort(Sum) == 1
     @test MetagraphOptimization.compute_effort(Data10) == 0
     @test MetagraphOptimization.compute_effort(Data20) == 0

From 4c452dce981196d2914498d76656d8456b1b0c57 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@proton.me>
Date: Mon, 11 Sep 2023 11:04:50 +0200
Subject: [PATCH 7/8] Add execution test

---
 src/code_gen/main.jl         |  5 ++---
 src/models/abc/create.jl     | 38 +++++++++++++++++++++++++++++++++++-
 test/runtests.jl             |  1 +
 test/unit_tests_execution.jl | 31 +++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 test/unit_tests_execution.jl

diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index e22b590..0677d2d 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -57,14 +57,13 @@ Execute the given generated_code (as returned by [`gen_code`](@ref)) on the give
 """
 function execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = generated_code
-    @assert length(input) == length(inputSymbols)
 
     assignInputs = Vector{Expr}()
     for (name, symbol) in inputSymbols
         type = nothing
-        if startswith("A", name)
+        if startswith(name, "A")
             type = A
-        elseif startswith("B", name)
+        elseif startswith(name, "B")
             type = B
         else
             type = C
diff --git a/src/models/abc/create.jl b/src/models/abc/create.jl
index cf9e39d..865b18c 100644
--- a/src/models/abc/create.jl
+++ b/src/models/abc/create.jl
@@ -26,8 +26,44 @@ Note: This does not take into account the preservation of momenta required for a
 """
 function gen_particles(ns::Dict{ParticleType, Int})
     particles = Dict{ParticleType, Vector{Particle}}()
-
     rng = MersenneTwister(0)
+
+
+    if ns == Dict((A => 2), (B => 2))
+        rho = 1.0
+
+        omega = rand(rng, Float64)
+        theta = rand(rng, Float64) * π
+        phi = rand(rng, Float64) * π
+
+        particles[A] = Vector{Particle}()
+        particles[B] = Vector{Particle}()
+
+        push!(particles[A], Particle(omega, 0, 0, omega, A))
+        push!(particles[B], Particle(omega, 0, 0, -omega, B))
+        push!(
+            particles[A],
+            Particle(
+                omega,
+                rho * cos(theta) * cos(phi),
+                rho * cos(theta) * sin(phi),
+                rho * sin(theta),
+                A,
+            ),
+        )
+        push!(
+            particles[B],
+            Particle(
+                omega,
+                -rho * cos(theta) * cos(phi),
+                -rho * cos(theta) * sin(phi),
+                -rho * sin(theta),
+                B,
+            ),
+        )
+        return particles
+    end
+
     for (type, n) in ns
         particles[type] = Vector{Particle}()
         for i in 1:n
diff --git a/test/runtests.jl b/test/runtests.jl
index 2ea1a76..5827252 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,6 +8,7 @@ using Test
     include("unit_tests_properties.jl")
     include("node_reduction.jl")
     include("unit_tests_graph.jl")
+    include("unit_tests_execution.jl")
 
     include("known_graphs.jl")
 end
diff --git a/test/unit_tests_execution.jl b/test/unit_tests_execution.jl
new file mode 100644
index 0000000..d0c0f68
--- /dev/null
+++ b/test/unit_tests_execution.jl
@@ -0,0 +1,31 @@
+import MetagraphOptimization.A
+import MetagraphOptimization.B
+import MetagraphOptimization.ParticleType
+
+@testset "Unit Tests Graph" begin
+    particles = Dict{ParticleType, Vector{Particle}}(
+        (
+            A => [
+                Particle(0.823648, 0.0, 0.0, 0.823648, A),
+                Particle(0.823648, -0.835061, -0.474802, 0.277915, A),
+            ]
+        ),
+        (
+            B => [
+                Particle(0.823648, 0.0, 0.0, -0.823648, B),
+                Particle(0.823648, 0.835061, 0.474802, -0.277915, B),
+            ]
+        ),
+    )
+
+    expected_result = 5.5320567694746876e-5
+
+    for _ in 1:10   # test in a loop because graph layout should not change the result
+        graph = parse_abc(joinpath(@__DIR__, "..", "input", "AB->AB.txt"))
+        @test isapprox(execute(graph, particles), expected_result; rtol = 0.001)
+
+        code = MetagraphOptimization.gen_code(graph)
+        @test isapprox(execute(code, particles), expected_result; rtol = 0.001)
+    end
+end
+println("Execution Unit Tests Complete!")

From 62791ab422a2b745608627fb34651305df350ca9 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@proton.me>
Date: Sun, 17 Sep 2023 11:20:10 +0200
Subject: [PATCH 8/8] Fix docs

---
 docs/make.jl                         |  5 ++++-
 docs/src/lib/internals/code_gen.md   |  8 ++++++++
 docs/src/lib/internals/models.md     |  6 +++---
 docs/src/lib/internals/properties.md |  2 +-
 docs/src/lib/internals/utility.md    |  2 +-
 src/code_gen/main.jl                 | 10 +++++-----
 6 files changed, 22 insertions(+), 11 deletions(-)
 create mode 100644 docs/src/lib/internals/code_gen.md

diff --git a/docs/make.jl b/docs/make.jl
index 578cdb5..5c81178 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -7,10 +7,12 @@ makedocs(
     root = "docs",
     source = "src",
     build = "build",
+    warnonly = true,
     clean = true,
     doctest = true,
     modules = Module[MetagraphOptimization],
-    repo = "https://code.woubery.com/Rubydragon/MetagraphOptimization.jl/src/branch/{commit}{path}#L{line}",
+    #repo = "https://code.woubery.com/Rubydragon/MetagraphOptimization.jl/src/branch/{commit}{path}#L{line}",
+    remotes = nothing,
     sitename = "MetagraphOptimization.jl",
     pages = [
         "index.md",
@@ -24,6 +26,7 @@ makedocs(
             "Models" => "lib/internals/models.md",
             "Diff" => "lib/internals/diff.md",
             "Utility" => "lib/internals/utility.md",
+            "Code Generation" => "lib/internals/code_gen.md",
         ],
         "Contribution" => "contribution.md",
     ],
diff --git a/docs/src/lib/internals/code_gen.md b/docs/src/lib/internals/code_gen.md
new file mode 100644
index 0000000..a9f4246
--- /dev/null
+++ b/docs/src/lib/internals/code_gen.md
@@ -0,0 +1,8 @@
+# Code Generation
+
+## Main
+```@autodocs
+Modules = [MetagraphOptimization]
+Pages = ["code_gen/main.jl"]
+Order = [:function]
+```
\ No newline at end of file
diff --git a/docs/src/lib/internals/models.md b/docs/src/lib/internals/models.md
index 1f7877f..192b91c 100644
--- a/docs/src/lib/internals/models.md
+++ b/docs/src/lib/internals/models.md
@@ -13,7 +13,7 @@ Order   = [:type, :constant]
 ```@autodocs
 Modules = [MetagraphOptimization]
 Pages = ["models/abc/particle.jl"]
-Order   = [:type, :constant]
+Order   = [:type, :constant, :function]
 ```
 
 ### Parse
@@ -33,14 +33,14 @@ Order   = [:function]
 ### Create
 ```@autodocs
 Modules = [MetagraphOptimization]
-Pages = ["models/abc/create.jl]
+Pages = ["models/abc/create.jl"]
 Order = [:function]
 ```
 
 ### Compute
 ```@autodocs
 Modules = [MetagraphOptimization]
-Pages = ["models/abc/compute.jl]
+Pages = ["models/abc/compute.jl"]
 Order = [:function]
 ```
 
diff --git a/docs/src/lib/internals/properties.md b/docs/src/lib/internals/properties.md
index 5773020..1539e52 100644
--- a/docs/src/lib/internals/properties.md
+++ b/docs/src/lib/internals/properties.md
@@ -10,7 +10,7 @@ Order   = [:type]
 ## Create
 ```@autodocs
 Modules = [MetagraphOptimization]
-Pages = ["task/create.jl"]
+Pages = ["properties/create.jl"]
 Order   = [:function]
 ```
 
diff --git a/docs/src/lib/internals/utility.md b/docs/src/lib/internals/utility.md
index ca034d8..2549cc9 100644
--- a/docs/src/lib/internals/utility.md
+++ b/docs/src/lib/internals/utility.md
@@ -3,7 +3,7 @@
 ## Helper Functions
 ```@autodocs
 Modules = [MetagraphOptimization]
-Pages   = ["utility.jl"]
+Pages   = ["./utility.jl"]
 Order   = [:type, :function]
 ```
 
diff --git a/src/code_gen/main.jl b/src/code_gen/main.jl
index 0677d2d..3bbb5a3 100644
--- a/src/code_gen/main.jl
+++ b/src/code_gen/main.jl
@@ -5,9 +5,9 @@ using DataStructures
 
 Generate the code for a given graph. The return value is a tuple of:
 
-    - `code::Expr`: The julia expression containing the code for the whole graph.
-    - `inputSymbols::Dict{String, Symbol}`: A dictionary of symbols mapping the names of the input nodes of the graph to the symbols their inputs should be provided on.
-    - `outputSymbol::Symbol`: The symbol of the final calculated value
+- `code::Expr`: The julia expression containing the code for the whole graph.
+- `inputSymbols::Dict{String, Symbol}`: A dictionary of symbols mapping the names of the input nodes of the graph to the symbols their inputs should be provided on.
+- `outputSymbol::Symbol`: The symbol of the final calculated value
 
 See also: [`execute`](@ref)
 """
@@ -53,7 +53,7 @@ end
 """
     execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
 
-Execute the given generated_code (as returned by [`gen_code`](@ref)) on the given input particles.
+Execute the given `generated_code` (as returned by [`gen_code`](@ref)) on the given input particles.
 """
 function execute(generated_code, input::Dict{ParticleType, Vector{Particle}})
     (code, inputSymbols, outputSymbol) = generated_code
@@ -89,7 +89,7 @@ end
 """
     execute(graph::DAG, input::Dict{ParticleType, Vector{Particle}})
 
-Execute the given generated_code (as returned by [`gen_code`](@ref)) on the given input particles.
+Execute the given `generated_code` (as returned by [`gen_code`](@ref)) on the given input particles.
 The input particles should be sorted correctly into the dictionary to their according [`ParticleType`](@ref)s.
 
 See also: [`gen_particles`](@ref)