From f1edce258a028db608653d64abca2548a9e154cc Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 31 Aug 2023 18:24:48 +0200 Subject: [PATCH] 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