Run on GPU

examples/congruent_in_ph.jl

@@ -1,11 +1,18 @@
+ENV["UCX_ERROR_SIGNALS"] = "SIGILL,SIGBUS,SIGFPE"
+
 using MetagraphOptimization
 using QEDbase
 using QEDcore
 using QEDprocesses
 using Random
 using UUIDs
+
+using CUDA
+
+using NamedDims
 using CSV
 using JLD2
+using FlexiMaps
 
 RNG = Random.MersenneTwister(123)
 
@@ -26,15 +33,12 @@ end
 
 function congruent_input_momenta(processDescription::GenericQEDProcess, omega::Number)
     # generate an input sample for given e + nk -> e' + k' process, where the nk are equal
-    massSum = 0
     inputMasses = Vector{Float64}()
     for particle in incoming_particles(processDescription)
-        massSum += mass(particle)
         push!(inputMasses, mass(particle))
     end
     outputMasses = Vector{Float64}()
     for particle in outgoing_particles(processDescription)
-        massSum += mass(particle)
         push!(outputMasses, mass(particle))
     end
 
@@ -49,7 +53,7 @@ function congruent_input_momenta(processDescription::GenericQEDProcess, omega::N
     return (tuple(initial_momenta...), tuple(final_momenta...))
 end
 
-# cos_theta ∈ [-1, 1] and phi ∈ [0, 2π]
+# theta ∈ [0, 2π] and phi ∈ [0, 2π]
 function congruent_input_momenta_scenario_2(
     processDescription::GenericQEDProcess,
     omega::Number,
@@ -60,15 +64,12 @@ function congruent_input_momenta_scenario_2(
     # same as above
 
     # generate an input sample for given e + nk -> e' + k' process, where the nk are equal
-    massSum = 0
     inputMasses = Vector{Float64}()
     for particle in incoming_particles(processDescription)
-        massSum += mass(particle)
         push!(inputMasses, mass(particle))
     end
     outputMasses = Vector{Float64}()
     for particle in outgoing_particles(processDescription)
-        massSum += mass(particle)
         push!(outputMasses, mass(particle))
     end
 
@@ -83,7 +84,7 @@ function congruent_input_momenta_scenario_2(
     # ----------
 
     # now calculate the final_momenta from omega, cos_theta and phi
-    n = number_particles(processDescription, ParticleStateful{Incoming, Photon, SFourMomentum})
+    n = number_particles(processDescription, Incoming(), Photon())
 
     cos_theta = cos(theta)
     omega_prime = (n * omega) / (1 + n * omega * (1 - cos_theta))
@@ -108,106 +109,86 @@ function build_psp(processDescription::GenericQEDProcess, momenta)
     )
 end
 
-return 0
+# hack to fix stacksize for threading
+with_stacksize(f, n) = fetch(schedule(Task(f, n)))
 
 # scenario 2
-N = 1000
-M = 1000
+N = 1024    # thetas
+M = 1024    # phis
+K = 64      # omegas
 
 thetas = collect(LinRange(0, 2π, N))
 phis = collect(LinRange(0, 2π, M))
+omegas = collect(maprange(log, 2e-2, 2e-7, K))
 
-for photons in [6]
+for photons in 1:5
     # temp process to generate momenta
-    for omega in [2e-3, 2e-6]
-        println("Generating $(N*M) inputs for $photons photons (Scenario 2 grid walk)...")
-        temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp())
+    println("Generating $(K*N*M) inputs for $photons photons (Scenario 2 grid walk)...")
+    temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp())
 
-        input_momenta = [
-            congruent_input_momenta_scenario_2(temp_process, omega, theta, phi) for
-            (theta, phi) in Iterators.product(thetas, phis)
-        ]
-        results = [0.0 for _ in 1:length(input_momenta)]
+    input_momenta =
+        Array{typeof(congruent_input_momenta_scenario_2(temp_process, omegas[1], thetas[1], phis[1]))}(undef, (K, N, M))
 
-        i = 1
-        for (in_pol, in_spin, out_pol, out_spin) in
-            Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()])
-
-            print(
-                "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ",
-            )
-            process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin)
-            inputs = build_psp.(Ref(process), input_momenta)
-
-            print("Preparing graph... ")
-            # prepare function
-            graph = gen_graph(process)
-            optimize_to_fixpoint!(ReductionOptimizer(), graph)
-            print("Preparing function... ")
-            func = get_compute_function(graph, process, mock_machine())
-
-            print("Calculating... ")
-            Threads.@threads for i in 1:length(inputs)
-                results[i] += abs2(func(inputs[i]))
+    Threads.@threads for k in 1:K
+        Threads.@threads for i in 1:N
+            Threads.@threads for j in 1:M
+                input_momenta[k, i, j] = congruent_input_momenta_scenario_2(temp_process, omegas[k], thetas[i], phis[j])
             end
-            println("Done.")
-            i += 1
         end
-
-        println("Writing results")
-
-        out_ph_moms = getindex.(getindex.(input_momenta, 2), 1)
-        out_el_moms = getindex.(getindex.(input_momenta, 2), 2)
-
-        @save "$(photons)_congruent_photons_omega_$(omega)_grid.jld2" out_ph_moms out_el_moms results
     end
-end
 
 
-n = 1000000
+    cu_results = CuArray{Float64}(undef, size(input_momenta))
+    fill!(cu_results, 0.0)
 
-# n is the number of incoming photons
-# omega is the number
+    i = 1
+    for (in_pol, in_spin, out_pol, out_spin) in
+        Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()])
 
-for photons in [6]
-    # temp process to generate momenta
-    for omega in [2e-3, 2e-6]
-        println("Generating $n inputs for $photons photons...")
-        temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp())
+        print(
+            "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ",
+        )
+        process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin)
 
-        input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n]
-        results = [0.0 for _ in 1:length(input_momenta)]
-
-        i = 1
-        for (in_pol, in_spin, out_pol, out_spin) in
-            Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()])
-
-            print(
-                "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ",
-            )
-            process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin)
-            inputs = build_psp.(Ref(process), input_momenta)
-
-            print("Preparing graph... ")
-            # prepare function
-            graph = gen_graph(process)
-            optimize_to_fixpoint!(ReductionOptimizer(), graph)
-            print("Preparing function... ")
-            func = get_compute_function(graph, process, mock_machine())
-
-            print("Calculating... ")
-            Threads.@threads for i in 1:n
-                results[i] += abs2(func(inputs[i]))
+        inputs = Array{typeof(build_psp(process, input_momenta[1, 1, 1]))}(undef, (K, N, M))
+        #println("input_momenta: $input_momenta")
+        Threads.@threads for k in 1:K
+            Threads.@threads for i in 1:N
+                Threads.@threads for j in 1:M
+                    inputs[k, i, j] = build_psp(process, input_momenta[k, i, j])
+                end
             end
-            println("Done.")
-            i += 1
         end
+        cu_inputs = CuArray(inputs)
 
-        println("Writing results")
+        print("Preparing graph... ")
+        graph = gen_graph(process)
+        optimize_to_fixpoint!(ReductionOptimizer(), graph)
+        print("Preparing function... ")
+        kernel! = get_cuda_kernel(graph, process, mock_machine())
+        #func = get_compute_function(graph, process, mock_machine())
 
-        out_ph_moms = getindex.(getindex.(input_momenta, 2), 1)
-        out_el_moms = getindex.(getindex.(input_momenta, 2), 2)
+        print("Calculating... ")
+        ts = 32
+        bs = Int64(length(cu_inputs) / 32)
 
-        @save "$(photons)_congruent_photons_omega_$(omega).jld2" out_ph_moms out_el_moms results
+        outputs = CuArray{ComplexF64}(undef, size(cu_inputs))
+
+        @cuda threads = ts blocks = bs always_inline = true kernel!(cu_inputs, outputs, length(cu_inputs))
+        CUDA.device_synchronize()
+        cu_results += abs2.(outputs)
+
+        println("Done.")
+        i += 1
     end
+
+    println("Writing results")
+
+    out_ph_moms = getindex.(getindex.(input_momenta, 2), 1)
+    out_el_moms = getindex.(getindex.(input_momenta, 2), 2)
+
+    results = NamedDimsArray{(:omegas, :thetas, :phis)}(Array(cu_results))
+    println("Named results array: $(typeof(results))")
+
+    @save "$(photons)_congruent_photons_grid.jld2" omegas thetas phis results
 end

src/MetagraphOptimization.jl

@@ -100,9 +100,6 @@ export ==, in, show, isempty, delete!, length
 
 export bytes_to_human_readable
 
-# TODO: this is probably not good
-import QEDprocesses.compute
-
 import Base.length
 import Base.show
 import Base.==

src/QEDprocesses_patch.jl

@@ -1,25 +1,5 @@
 # patch QEDprocesses
 # see issue https://github.com/QEDjl-project/QEDprocesses.jl/issues/77
-@inline function QEDprocesses.number_particles(
-    proc_def::QEDbase.AbstractProcessDefinition,
-    dir::DIR,
-    ::PT,
-) where {DIR <: QEDbase.ParticleDirection, PT <: QEDbase.AbstractParticleType}
-    return count(x -> x isa PT, particles(proc_def, dir))
-end
-
-@inline function QEDprocesses.number_particles(
-    proc_def::QEDbase.AbstractProcessDefinition,
-    ::PS,
-) where {
-    DIR <: QEDbase.ParticleDirection,
-    PT <: QEDbase.AbstractParticleType,
-    EL <: AbstractFourMomentum,
-    PS <: ParticleStateful{DIR, PT, EL},
-}
-    return QEDprocesses.number_particles(proc_def, DIR(), PT())
-end
-
 @inline function QEDprocesses.number_particles(
     proc_def::QEDbase.AbstractProcessDefinition,
     ::Type{PS},
@@ -43,29 +23,3 @@ end
 ) where {DIR <: ParticleDirection, SPECIES <: AbstractParticleType, EL <: AbstractFourMomentum}
     return ParticleStateful(DIR(), SPECIES(), mom)
 end
-
-@inline function QEDbase.momentum(
-    psp::AbstractPhaseSpacePoint{MODEL, PROC, PS_DEF, INT, OUTT},
-    dir::ParticleDirection,
-    species::AbstractParticleType,
-    n::Int,
-) where {MODEL, PROC, PS_DEF, INT, OUTT}
-    # TODO: can be done through fancy template recursion too with 0 overhead
-    i = 0
-    c = n
-    for p in particles(psp, dir)
-        i += 1
-        if particle_species(p) isa typeof(species)
-            c -= 1
-        end
-        if c == 0
-            break
-        end
-    end
-
-    if c != 0 || n <= 0
-        throw(InvalidInputError("could not get $n-th momentum of $dir $species, does not exist"))
-    end
-
-    return momenta(psp, dir)[i]
-end

src/models/physics_models/qed/compute.jl

@@ -17,7 +17,7 @@ function input_expr(instance::GenericQEDProcess, name::String, psp_symbol::Symbo
 
     return Meta.parse(
         "ParticleValueSP(
-    $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), $index)),
+    $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), Val($index))),
     0.0im,
     $(construction_string(spin_or_pol(instance, type, index))),
 )",

src/models/physics_models/qed/diagrams.jl

@@ -504,18 +504,15 @@ function gen_compton_diagram_from_order(order::Vector{Int}, inFerm, outFerm, n::
     return new_diagram
 end
 
-#=
 """
     gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, outFerm, n::Int, m::Int)
 
 Helper function for [`gen_compton_diagrams`](@Ref). Generates a single diagram for the given order and n input and m output photons.
 """
-function gen_compton_diagram_from_order_one_side(
-    order::Vector{Int}, inFerm, outFerm, n::Int, m::Int
-)
+function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, outFerm, n::Int, m::Int)
     photons = vcat(
-        [FeynmanParticle(ParticleStateful{Incoming, Photon}, i) for i in 1:n],
-        [FeynmanParticle(ParticleStateful{Outgoing, Photon}, i) for i in 1:m],
+        [FeynmanParticle(ParticleStateful{Incoming, Photon, SFourMomentum}, i) for i in 1:n],
+        [FeynmanParticle(ParticleStateful{Outgoing, Photon, SFourMomentum}, i) for i in 1:m],
     )
 
     new_diagram = FeynmanDiagram(
@@ -523,10 +520,10 @@ function gen_compton_diagram_from_order_one_side(
         missing,
         [inFerm, outFerm, photons...],
         Dict{Type, Int64}(
-            ParticleStateful{Incoming, Electron} => 1,
-            ParticleStateful{Outgoing, Electron} => 1,
-            ParticleStateful{Incoming, Photon} => n,
-            ParticleStateful{Outgoing, Photon} => m,
+            ParticleStateful{Incoming, Electron, SFourMomentum} => 1,
+            ParticleStateful{Outgoing, Electron, SFourMomentum} => 1,
+            ParticleStateful{Incoming, Photon, SFourMomentum} => n,
+            ParticleStateful{Outgoing, Photon, SFourMomentum} => m,
         ),
     )
 
@@ -538,9 +535,9 @@ function gen_compton_diagram_from_order_one_side(
     while left_index <= right_index
         # left side
         v_left = FeynmanVertex(
-            FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations),
+            FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations),
             photons[order[left_index]],
-            FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations + 1),
+            FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations + 1),
         )
         left_index += 1
         add_vertex!(new_diagram, v_left)
@@ -553,9 +550,9 @@ function gen_compton_diagram_from_order_one_side(
         if (iterations == 1)
             # right side
             v_right = FeynmanVertex(
-                FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations),
+                FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations),
                 photons[order[right_index]],
-                FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations + 1),
+                FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations + 1),
             )
             right_index -= 1
             add_vertex!(new_diagram, v_right)
@@ -569,7 +566,6 @@ function gen_compton_diagram_from_order_one_side(
     add_tie!(new_diagram, FeynmanTie(ps[1], ps[2]))
     return new_diagram
 end
-=#
 
 """
     gen_compton_diagrams(n::Int, m::Int)