Fix the types, add some profiling examples (#15)

Reviewed-on: Rubydragon/MetagraphOptimization.jl#15
Co-authored-by: Anton Reinhard <anton.reinhard@proton.me>
Co-committed-by: Anton Reinhard <anton.reinhard@proton.me>
This commit is contained in:
2023-11-13 12:55:02 +01:00
committed by Anton Reinhard
parent 5a30f57e1f
commit 2709eeb3dc
6 changed files with 497 additions and 457 deletions

View File

@@ -7,7 +7,7 @@ Return the particle and value as is.
0 FLOP.
"""
function compute(::ComputeTaskP, data::ParticleValue)
function compute(::ComputeTaskP, data::ParticleValue{P})::ParticleValue{P} where {P <: ABCParticle}
return data
end
@@ -18,7 +18,7 @@ Compute an outer edge. Return the particle value with the same particle and the
1 FLOP.
"""
function compute(::ComputeTaskU, data::ParticleValue)
function compute(::ComputeTaskU, data::ParticleValue{P})::ParticleValue{P} where {P <: ABCParticle}
return ParticleValue(data.p, data.v * outer_edge(data.p))
end
@@ -29,7 +29,11 @@ Compute a vertex. Preserve momentum and particle types (AB->C etc.) to create re
6 FLOP.
"""
function compute(::ComputeTaskV, data1::ParticleValue, data2::ParticleValue)
function compute(
::ComputeTaskV,
data1::ParticleValue{P1},
data2::ParticleValue{P2},
)::ParticleValue where {P1 <: ABCParticle, P2 <: ABCParticle}
p3 = preserve_momentum(data1.p, data2.p)
dataOut = ParticleValue(p3, data1.v * vertex() * data2.v)
return dataOut
@@ -44,7 +48,7 @@ For valid inputs, both input particles should have the same momenta at this poin
12 FLOP.
"""
function compute(::ComputeTaskS2, data1::ParticleValue, data2::ParticleValue)
function compute(::ComputeTaskS2, data1::ParticleValue{P}, data2::ParticleValue{P})::Float64 where {P <: ABCParticle}
#=
@assert isapprox(abs(data1.p.momentum.E), abs(data2.p.momentum.E), rtol = 0.001, atol = sqrt(eps())) "E: $(data1.p.momentum.E) vs. $(data2.p.momentum.E)"
@assert isapprox(data1.p.momentum.px, -data2.p.momentum.px, rtol = 0.001, atol = sqrt(eps())) "px: $(data1.p.momentum.px) vs. $(data2.p.momentum.px)"
@@ -61,7 +65,7 @@ Compute inner edge (1 input particle, 1 output particle).
11 FLOP.
"""
function compute(::ComputeTaskS1, data::ParticleValue)
function compute(::ComputeTaskS1, data::ParticleValue{P})::ParticleValue{P} where {P <: ABCParticle}
return ParticleValue(data.p, data.v * inner_edge(data.p))
end
@@ -72,7 +76,7 @@ Compute a sum over the vector. Use an algorithm that accounts for accumulated er
Linearly many FLOP with growing data.
"""
function compute(::ComputeTaskSum, data::Vector{Float64})
function compute(::ComputeTaskSum, data::Vector{Float64})::Float64
return sum_kbn(data)
end