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>

examples/ab5.jl

@@ -0,0 +1,34 @@
+using MetagraphOptimization
+using BenchmarkTools
+
+println("Getting machine info")
+@time machine = get_machine_info()
+
+println("Making model")
+@time model = ABCModel()
+
+println("Making process")
+process_str = "AB->ABBBBB"
+@time process = parse_process(process_str, model)
+
+println("Parsing DAG")
+@time graph = parse_dag("input/$process_str.txt", model)
+
+println("Generating input data")
+@time input_data = [gen_process_input(process) for _ in 1:1000]
+
+include("profiling_utilities.jl")
+println("Reducing graph")
+@time reduce_all!(graph)
+
+println("Generating compute function")
+@time compute_func = get_compute_function(graph, process, machine)
+
+println("First run, single argument")
+@time compute_func(input_data[1])
+
+println("\nBenchmarking function, 1 input")
+display(@benchmark compute_func($(input_data[1])))
+
+println("\nBenchmarking function, 1000 inputs")
+display(@benchmark compute_func.($input_data))

examples/ab7.jl

@@ -0,0 +1,34 @@
+using MetagraphOptimization
+using BenchmarkTools
+
+println("Getting machine info")
+@time machine = get_machine_info()
+
+println("Making model")
+@time model = ABCModel()
+
+println("Making process")
+process_str = "AB->ABBBBBBB"
+@time process = parse_process(process_str, model)
+
+println("Parsing DAG")
+@time graph = parse_dag("input/$process_str.txt", model)
+
+println("Generating input data")
+@time input_data = [gen_process_input(process) for _ in 1:1000]
+
+include("profiling_utilities.jl")
+println("Reducing graph")
+@time reduce_all!(graph)
+
+println("Generating compute function")
+@time compute_func = get_compute_function(graph, process, machine)
+
+println("First run, single argument")
+@time compute_func(input_data[1])
+
+println("\nBenchmarking function, 1 input")
+display(@benchmark compute_func($(input_data[1])))
+
+println("\nBenchmarking function, 1000 inputs")
+display(@benchmark compute_func.($input_data))

src/code_gen/main.jl

@@ -79,7 +79,7 @@ function gen_input_assignment_code(
             # TODO: how to get the "default" cpu device?
             device = entry_device(machine)
             evalExpr = eval(gen_access_expr(device, symbol))
-            push!(assignInputs, Meta.parse("$(evalExpr) = ParticleValue($p, 1.0)"))
+            push!(assignInputs, Meta.parse("$(evalExpr)::ParticleValue{$type} = ParticleValue($p, 1.0)"))
         end
     end
 
@@ -102,6 +102,7 @@ function get_compute_function(graph::DAG, process::AbstractProcessDescription, m
     expr = Meta.parse(
         "function compute_$(functionId)(input::AbstractProcessInput) $initCaches; $assignInputs; $code; return $resSym; end",
     )
+
     func = eval(expr)
 
     return func

src/models/abc/compute.jl

@@ -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
 

src/optimization/greedy.jl

@@ -0,0 +1,8 @@
+"""
+    GreedyOptimizer
+
+An implementation of the greedy optimization algorithm, simply choosing the best next option evaluated with the given estimator.
+"""
+struct GreedyOptimizer
+    estimator::AbstractEstimator
+end