Add code gen documentation

2023-09-07 18:23:36 +02:00 · 2023-09-07 18:23:36 +02:00 · e59d24ebe5
commit e59d24ebe5
parent d1666de432
5 changed files with 240 additions and 34 deletions
--- 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)

--- 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
--- 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}}()
--- 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,
--- 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