GPU support

Project.toml

@@ -1,12 +1,6 @@
 name = "QEDprocesses"
 uuid = "46de9c38-1bb3-4547-a1ec-da24d767fdad"
-authors = [
-    "Uwe Hernandez Acosta <[email protected]>",
-    "Simeon Ehrig",
-    "Klaus Steiniger",
-    "Tom Jungnickel",
-    "Anton Reinhard",
-]
+authors = ["Uwe Hernandez Acosta <[email protected]>", "Simeon Ehrig", "Klaus Steiniger", "Tom Jungnickel", "Anton Reinhard"]
 version = "0.2.0"
 
 [deps]
@@ -15,17 +9,31 @@ QEDcore = "35dc0263-cb5f-4c33-a114-1d7f54ab753e"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+
+[extensions]
+AMDGPUExt = "AMDGPU"
+CUDAExt = "CUDA"
+
 [compat]
 QEDbase = "0.2.2"
 QEDcore = "0.1"
 QuadGK = "2"
 StaticArrays = "1"
+AMDGPU = "1"
+CUDA = "5"
 julia = "1.6"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+# these need to be here to support Julia versions < 1.9
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+
 [targets]
-test = ["Random", "SafeTestsets", "Test"]
+test = ["Random", "SafeTestsets", "Test", "CUDA", "AMDGPU"]

ext/AMDGPUExt.jl

@@ -0,0 +1,7 @@
+module AMDGPUExt
+
+using QEDprocesses, AMDGPU
+
+# include specialized AMDGPU functions here
+
+end

ext/CUDAExt.jl

@@ -0,0 +1,7 @@
+module CUDAExt
+
+using QEDprocesses, CUDA
+
+# include specialized CUDA functions here
+
+end

src/processes/one_photon_compton/perturbative/cross_section.jl

@@ -78,8 +78,8 @@ end
     in_photon_state = base_state(Photon(), Incoming(), in_photon_mom, proc.in_pol)
 
     out_electron_state = base_state(Electron(), Outgoing(), out_electron_mom, proc.out_spin)
-
     out_photon_state = base_state(Photon(), Outgoing(), out_photon_mom, proc.out_pol)
+
     return _pert_compton_matrix_element(
         in_electron_mom,
         in_electron_state,
@@ -109,23 +109,18 @@ function _pert_compton_matrix_element(
         QEDbase._as_svec(out_photon_state),
     )
 
-    matrix_elements = Vector{ComplexF64}()
-    sizehint!(matrix_elements, length(base_states_comb))
-    for (in_el, in_ph, out_el, out_ph) in base_states_comb
-        push!(
-            matrix_elements,
-            _pert_compton_matrix_element_single(
-                in_electron_mom,
-                in_el,
-                in_photon_mom,
-                in_ph,
-                out_electron_mom,
-                out_el,
-                out_photon_mom,
-                out_ph,
-            ),
-        )
-    end
+    matrix_elements = NTuple{length(base_states_comb),ComplexF64}(
+        _pert_compton_matrix_element_single(
+            in_electron_mom,
+            in_el,
+            in_photon_mom,
+            in_ph,
+            out_electron_mom,
+            out_el,
+            out_photon_mom,
+            out_ph,
+        ) for (in_el, in_ph, out_el, out_ph) in base_states_comb
+    )
 
     return matrix_elements
 end

src/processes/one_photon_compton/perturbative/total_probability.jl

@@ -8,7 +8,6 @@ function QEDbase._total_probability(in_psp::InPhaseSpacePoint{<:Compton,<:Pertur
     end
 
     tot_prob, _ = quadgk(func, -1, 1; rtol=sqrt(eps(omega)))
-
     tot_prob *= 2 * pi # phi integration is trivial
     return tot_prob
 end

test/gpu/process_interface.jl

@@ -0,0 +1,161 @@
+using QEDprocesses
+using QEDbase
+using QEDcore
+
+using AMDGPU
+using CUDA
+using Random
+using SafeTestsets
+
+include("../test_implementation/random_coordinates.jl")
+
+GPU_VECTOR_TYPES = Vector{Type}()
+if AMDGPU.functional()
+    println("Testing with AMDGPU.jl")
+    push!(GPU_VECTOR_TYPES, ROCVector)
+end
+if CUDA.functional()
+    println("Testing with CUDA.jl")
+    push!(GPU_VECTOR_TYPES, CuVector)
+end
+
+if isempty(GPU_VECTOR_TYPES)
+    @warn "No functional GPUs found for testing!"
+    return nothing
+end
+
+PROC_DEF_TUPLES = [
+    (
+        Compton(),
+        PerturbativeQED(),
+        PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()),
+    ),
+    (
+        Compton(SpinUp(), PolX(), SpinDown(), PolY()),
+        PerturbativeQED(),
+        PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()),
+    ),
+]
+
+RNG = Random.MersenneTwister(573)
+
+@testset "Testing with $VECTOR_TYPE" for VECTOR_TYPE in GPU_VECTOR_TYPES
+    @testset "$proc $model $ps_def" for (proc, model, ps_def) in PROC_DEF_TUPLES
+        N = 10_000
+
+        psps = [
+            PhaseSpacePoint(
+                proc, model, ps_def, _rand_coordinates(RNG, proc, model, ps_def)...
+            ) for _ in 1:N
+        ]
+        procs = [proc for _ in 1:N]
+
+        gpupsps = VECTOR_TYPE(psps)
+        gpuprocs = VECTOR_TYPE(procs)
+
+        @testset "PSP interface" begin
+            in_moms_gpu = Vector(momenta.(gpupsps, Incoming()))
+            out_moms_gpu = Vector(momenta.(gpupsps, Outgoing()))
+            in_moms = momenta.(psps, Incoming())
+            out_moms = momenta.(psps, Outgoing())
+
+            @test getindex.(in_moms_gpu, Ref(1)) == getindex.(in_moms, Ref(1))
+            @test getindex.(in_moms_gpu, Ref(2)) == getindex.(in_moms, Ref(2))
+            @test getindex.(out_moms_gpu, Ref(1)) == getindex.(out_moms, Ref(1))
+            @test getindex.(out_moms_gpu, Ref(2)) == getindex.(out_moms, Ref(2))
+        end
+
+        @testset "Private Process Functions" begin
+            @test all(
+                isapprox.(
+                    Vector(QEDbase._averaging_norm.(gpuprocs)),
+                    QEDbase._averaging_norm.(procs),
+                ),
+            )
+        end
+
+        @testset "Public Process Functions" begin
+            @test Vector(incoming_particles.(gpuprocs)) == incoming_particles.(procs)
+            @test Vector(outgoing_particles.(gpuprocs)) == outgoing_particles.(procs)
+
+            @test Vector(particles.(gpuprocs, Incoming())) == particles.(procs, Incoming())
+            @test Vector(particles.(gpuprocs, Outgoing())) == particles.(procs, Outgoing())
+
+            @test Vector(number_incoming_particles.(gpuprocs)) ==
+                number_incoming_particles.(procs)
+            @test Vector(number_outgoing_particles.(gpuprocs)) ==
+                number_outgoing_particles.(procs)
+
+            @test Vector(number_particles.(gpuprocs, Incoming())) ==
+                number_particles.(procs, Incoming())
+            @test Vector(number_particles.(gpuprocs, Outgoing())) ==
+                number_particles.(procs, Outgoing())
+
+            @test Vector(QEDbase.in_phase_space_dimension.(gpuprocs, model)) ==
+                QEDbase.in_phase_space_dimension.(procs, model)
+            @test Vector(QEDbase.out_phase_space_dimension.(gpuprocs, model)) ==
+                QEDbase.out_phase_space_dimension.(procs, model)
+        end
+
+        @testset "Private PhaseSpacePoint Interface" begin
+            @test all(
+                isapprox.(
+                    Vector(QEDbase._incident_flux.(gpupsps)), QEDbase._incident_flux.(psps)
+                ),
+            )
+
+            @test all(
+                tuple_isapprox.(
+                    Vector(QEDbase._matrix_element.(gpupsps)),
+                    QEDbase._matrix_element.(psps),
+                ),
+            )
+
+            @test Vector(QEDbase._is_in_phasespace.(gpupsps)) ==
+                QEDbase._is_in_phasespace.(psps)
+
+            @test all(
+                isapprox.(
+                    Vector(QEDbase._phase_space_factor.(gpupsps)),
+                    QEDbase._phase_space_factor.(psps),
+                ),
+            )
+
+            #= TODO: this is currently broken
+            @test all(
+                isapprox.(
+                    Vector(QEDprocesses._total_probability.(gpupsps)),
+                    QEDprocesses._total_probability.(psps),
+                ),
+            )=#
+        end
+
+        @testset "Public PhaseSpacePoint Interface" begin
+            @test all(
+                isapprox.(
+                    Vector(differential_probability.(gpupsps)),
+                    differential_probability.(psps),
+                ),
+            )
+
+            @test all(
+                isapprox.(
+                    Vector(QEDbase._is_in_phasespace.(gpupsps)),
+                    QEDbase._is_in_phasespace.(psps),
+                ),
+            )
+
+            @test all(
+                isapprox.(
+                    Vector(differential_cross_section.(gpupsps)),
+                    differential_cross_section.(psps),
+                ),
+            )
+
+            #= TODO: this is currently broken
+            @test all(
+                isapprox.(Vector(total_cross_section.(gpupsps)), total_cross_section.(psps))
+            )=#
+        end
+    end
+end

test/runtests.jl

@@ -4,4 +4,6 @@ using SafeTestsets
 begin
     # scattering processes
     include("processes/run_process_test.jl")
+
+    include("gpu/process_interface.jl")
 end

test/test_implementation/random_coordinates.jl

@@ -0,0 +1,16 @@
+"""
+Return a tuple of tuples of incoming and outgoing coordinates for a given process, model and ps_def that make up a physical phase space point.
+"""
+function _rand_coordinates(
+    rng::AbstractRNG, ::PROCESS, ::MODEL, ::PS_DEF
+) where {PROCESS<:Compton,MODEL<:PerturbativeQED,PS_DEF<:PhasespaceDefinition}
+    return ((rand(rng, Float64),), (rand(rng, Float64), rand(rng, Float64)))
+end
+
+tuple_isapprox(::Tuple{}, ::Tuple{}; atol=0.0, rtol=eps()) = true
+function tuple_isapprox(
+    a::Tuple{<:Number,Vararg}, b::Tuple{<:Number,Vararg}; atol=0.0, rtol=eps()
+)
+    return isapprox(a[1], b[1]; atol=atol, rtol=rtol) &&
+           tuple_isapprox(a[2:end], b[2:end]; atol=atol, rtol=rtol)
+end