GPU support

.gitlab-ci.yml

@@ -8,6 +8,8 @@ generate_pipeline:
   variables:
     CI_GIT_CI_TOOLS_URL: https://github.com/QEDjl-project/QuantumElectrodynamics.jl.git
     CI_GIT_CI_TOOLS_BRANCH: dev
+    CI_ENABLE_CUDA_TESTS: "ON"
+    CI_ENABLE_AMDGPU_TESTS: "ON"
   script:
     - apt update && apt install -y git
     - git clone --depth 1 -b $CI_GIT_CI_TOOLS_BRANCH $CI_GIT_CI_TOOLS_URL /generator

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.3.0"
 
 [deps]
@@ -23,9 +17,10 @@ StaticArrays = "1"
 julia = "1.10"
 
 [extras]
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Random", "SafeTestsets", "Test"]
+test = ["Random", "SafeTestsets", "Pkg", "Test"]

src/processes/one_photon_compton/perturbative/cross_section.jl

@@ -75,8 +75,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,
@@ -106,11 +106,8 @@ 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,
+    matrix_elements::NTuple{length(base_states_comb),ComplexF64} = (
+        (
             _pert_compton_matrix_element_single(
                 in_electron_mom,
                 in_el,
@@ -120,9 +117,9 @@ function _pert_compton_matrix_element(
                 out_el,
                 out_photon_mom,
                 out_ph,
-            ),
-        )
-    end
+            ) 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,Perturba
     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,170 @@
+using QEDprocesses
+using QEDbase
+using QEDcore
+
+using Random
+using SafeTestsets
+
+DEF_POLS = (PolX(), PolY())
+DEF_SPINS = (SpinUp(), SpinDown())
+
+PROC_DEF_TUPLES = [
+    (
+        Compton(),
+        PerturbativeQED(),
+        PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()),
+    ),
+    [
+        (
+            Compton(s1, p1, s2, p2),
+            PerturbativeQED(),
+            PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()),
+        ) for
+        (s1, p1, s2, p2) in Iterators.product(DEF_SPINS, DEF_POLS, DEF_SPINS, DEF_POLS)
+    ]...,
+]
+
+RNG = Random.MersenneTwister(573)
+
+@testset "Testing with $GPU_MODULE" for (GPU_MODULE, VECTOR_TYPE) in GPUS
+    @testset "$proc $model $ps_def" for (proc, model, ps_def) in PROC_DEF_TUPLES
+        N = 100
+
+        @info "$proc $model $ps_def"
+        flush(stdout)
+
+        psps = [
+            PhaseSpacePoint(
+                proc, model, ps_def, _rand_coordinates(RNG, proc, model, ps_def)...
+            ) for _ in 1:N
+        ]
+        procs = [proc for _ in 1:N]
+
+        @info "GPU allocation"
+        flush(stdout)
+
+        gpupsps = VECTOR_TYPE(psps)
+        gpuprocs = VECTOR_TYPE(procs)
+
+        @testset "PSP interface" begin
+            @info "GPU momenta.()"
+            flush(stdout)
+
+            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
+            @info "GPU _averaging_norm.()"
+            flush(stdout)
+
+            @test all(
+                isapprox.(
+                    Vector(QEDbase._averaging_norm.(gpuprocs)),
+                    QEDbase._averaging_norm.(procs),
+                ),
+            )
+        end
+
+        @testset "Public Process Functions" begin
+            @info "GPU public process functions"
+            flush(stdout)
+
+            @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
+            @info "GPU private PSP functions"
+            flush(stdout)
+
+            @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);
+                    rtol=sqrt(eps(Float64)),
+                ),
+            )
+
+            @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),
+                ),
+            )
+
+            # this currently throws an exception because QuadGK does not work on the GPU
+            @test all(
+                isapprox.(
+                    Vector(QEDprocesses._total_probability.(gpupsps)),
+                    QEDprocesses._total_probability.(psps),
+                ),
+            ) broken = true
+        end
+
+        @testset "Public PhaseSpacePoint Interface" begin
+            @info "GPU public PSP functions"
+            flush(stdout)
+
+            @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),
+                ),
+            )
+
+            # as above, this currently throws an exception because QuadGK does not work on the GPU
+            @test all(
+                isapprox.(Vector(total_cross_section.(gpupsps)), total_cross_section.(psps))
+            ) broken = true
+        end
+    end
+end

test/gpu/runtests.jl

@@ -0,0 +1,94 @@
+"""
+This file sets up GPU testing. By default, it will check if GPU libraries are installed and
+functional, and execute the unit tests then. Additionally, if an environment variable is set
+("TEST_<GPU> = 1"), the tests will fail if the library is not functional.
+"""
+
+GPUS = Vector{Tuple{Module,Type}}()
+
+# check if we test with AMDGPU
+amdgpu_tests = tryparse(Bool, get(ENV, "TEST_AMDGPU", "0"))
+if amdgpu_tests
+    try
+        using Pkg
+        Pkg.add("AMDGPU")
+
+        using AMDGPU
+        AMDGPU.functional() || throw(
+            "trying to test with AMDGPU.jl but it is not functional (AMDGPU.functional() == false)",
+        )
+        push!(GPUS, (AMDGPU, ROCVector))
+        @info "Testing with AMDGPU.jl"
+    catch e
+        @error "failed to run GPU tests, make sure the required libraries are installed\n$(e)"
+        @test false
+    end
+end
+
+# check if we test with CUDA
+cuda_tests = tryparse(Bool, get(ENV, "TEST_CUDA", "0"))
+if cuda_tests
+    try
+        using Pkg
+        Pkg.add("CUDA")
+
+        using CUDA
+        CUDA.functional() || throw(
+            "trying to test with CUDA.jl but it is not functional (CUDA.functional() == false)",
+        )
+        push!(GPUS, (CUDA, CuVector))
+        @info "Testing with CUDA.jl"
+    catch e
+        @error "failed to run GPU tests, make sure the required libraries are installed\n$(e)"
+        @test false
+    end
+end
+
+# check if we test with oneAPI
+oneapi_tests = tryparse(Bool, get(ENV, "TEST_ONEAPI", "0"))
+if oneapi_tests
+    try
+        using Pkg
+        Pkg.add("oneAPI")
+
+        using oneAPI
+        oneAPI.functional() || throw(
+            "trying to test with oneAPI.jl but it is not functional (oneAPI.functional() == false)",
+        )
+        push!(GPUS, (oneAPI, oneVector))
+        @info "Testing with oneAPI.jl"
+    catch e
+        @error "failed to run GPU tests, make sure the required libraries are installed\n$(e)"
+        @test false
+    end
+end
+
+# check if we test with Metal
+metal_tests = tryparse(Bool, get(ENV, "TEST_METAL", "0"))
+if metal_tests
+    try
+        using Pkg
+        Pkg.add("Metal")
+
+        using Metal
+        Metal.functional() || throw(
+            "trying to test with Metal.jl but it is not functional (Metal.functional() == false)",
+        )
+        push!(GPUS, (Metal, MtlVector))
+        @info "Testing with Metal.jl"
+    catch e
+        @error "failed to run GPU tests, make sure the required libraries are installed\n$(e)"
+        @test false
+    end
+end
+
+if isempty(GPUS)
+    @info """No GPU tests are enabled, skipping tests...
+    To test GPU functionality, please use 'TEST_<GPU> = 1 julia ...' for one of GPU=[CUDA, AMDGPU, METAL, ONEAPI]"""
+    return nothing
+end
+
+include("../test_implementation/random_coordinates.jl")
+
+# from here on, we cannot use safe test sets or we would unload the GPU libraries again
+include("process_interface.jl")

test/runtests.jl

@@ -1,7 +1,18 @@
 using Test
 using SafeTestsets
 
-begin
+# check if we run CPU tests (yes by default)
+cpu_tests = tryparse(Bool, get(ENV, "TEST_CPU", "1"))
+
+if cpu_tests
     # scattering processes
     include("processes/run_process_test.jl")
+else
+    @info "Skipping CPU tests"
+end
+
+begin
+    @time @safetestset "GPU testing" begin
+        include("gpu/runtests.jl")
+    end
 end

test/test_implementation/random_coordinates.jl

@@ -0,0 +1,21 @@
+using QEDcore, QEDprocesses
+using Random
+
+"""
+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_iaspprox(::Tuple{}, ::Tuple{Vararg}; atol=0.0, rtol=eps()) = false
+tuple_iaspprox(::Tuple{Vararg}, ::Tuple{}; atol=0.0, rtol=eps()) = false
+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