From e59d24ebe5262f26784e55190048d55fcea09574 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 7 Sep 2023 18:23:36 +0200 Subject: [PATCH] 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