CI: implement GPU test support (#96)

Implement support for running unit tests on Nvidia and AMD GPU runner.
Implement unit tests for all unit test generator functions.

.ci/CI/src/Bootloader.jl

@@ -186,10 +186,45 @@ function main()
         test_package,
         unit_test_julia_versions,
         target_branch,
+        CPU,
         tools_git_repo,
         get_unit_test_nightly_baseimage(),
     )
 
+    if haskey(ENV, "CI_ENABLE_CUDA_TESTS") || haskey(ENV, "CI_ENABLE_AMDGPU_TESTS")
+        for version in ["rc", "nightly"]
+            if version in unit_test_julia_versions
+                @info "Remove unit test version $(version) for GPU tests"
+                filter!(v -> v != version, unit_test_julia_versions)
+            end
+        end
+    end
+
+    if haskey(ENV, "CI_ENABLE_CUDA_TESTS")
+        @info "Generate CUDA unit tests"
+        add_unit_test_job_yaml!(
+            job_yaml,
+            test_package,
+            unit_test_julia_versions,
+            target_branch,
+            CUDA,
+            tools_git_repo,
+        )
+    end
+
+    if haskey(ENV, "CI_ENABLE_AMDGPU_TESTS")
+        @info "Generate AMDGPU unit tests"
+
+        add_unit_test_job_yaml!(
+            job_yaml,
+            test_package,
+            unit_test_julia_versions,
+            target_branch,
+            AMDGPU,
+            tools_git_repo,
+        )
+    end
+
     add_integration_test_job_yaml!(job_yaml, test_package, target_branch, tools_git_repo)
 
     add_unit_test_verify_job_yaml!(job_yaml, target_branch, tools_git_repo)

.ci/CI/src/CI.jl

@@ -1,6 +1,7 @@
 
 module CI
 include("./modules/Utils.jl")
+include("./modules/UnitTest.jl")
 include("./modules/IntegTest.jl")
 include("./modules/GitLabTargetBranch.jl")
 include("./SetupDevEnv.jl")

.ci/CI/src/modules/UnitTest.jl

@@ -1,10 +1,15 @@
+"""
+Specify target processor for the tests.
+"""
+@enum TestPlatform CPU CUDA AMDGPU ONEAPI METAL
 
 """
     add_unit_test_job_yaml!(
         job_dict::Dict,
         test_package::TestPackage,
         julia_versions::Vector{String},
         target_branch::AbstractString,
+        test_platform::TestPlatform=CPU,
         tools_git_repo::ToolsGitRepo=ToolsGitRepo(
             "https://github.com/QEDjl-project/QuantumElectrodynamics.jl.git", "dev"
         ),
@@ -19,6 +24,7 @@ to be directly translated to GitLab CI yaml.
 - `test_package::TestPackage`: Properties of the package to be tested, such as name and version.
 - `julia_versions::Vector{String}`: Julia version used for the tests.
 - `target_branch::AbstractString`: A different job code is generated depending on the target branch.
+- `test_platform::TestPlatform`: Set target platform test, e.g. CPU, Nvidia GPU or AMD GPU.
 - `tools_git_repo::ToolsGitRepo`: URL and branch of the Git repository from which the CI tools are
     cloned in unit test job.
 - `nightly_base_image::AbstractString`: Name of the job base image if the Julia version is nightly.
@@ -28,25 +34,41 @@ function add_unit_test_job_yaml!(
     test_package::TestPackage,
     julia_versions::Vector{String},
     target_branch::AbstractString,
+    test_platform::TestPlatform=CPU,
     tools_git_repo::ToolsGitRepo=ToolsGitRepo(
         "https://github.com/QEDjl-project/QuantumElectrodynamics.jl.git", "dev"
     ),
     nightly_base_image::AbstractString="debian:bookworm-slim",
 )
+    if test_platform in [ONEAPI, METAL]
+        throw(ArgumentError("argument test_platform not implemented for $(test_platform)"))
+    end
+
     if !haskey(job_dict, "stages")
         job_dict["stages"] = []
     end
 
     push!(job_dict["stages"], "unit-test")
 
     for version in julia_versions
+        test_platform_name = lowercase(string(test_platform))
         if version != "nightly"
-            job_dict["unit_test_julia_$(replace(version, "." => "_"))"] = _get_normal_unit_test(
-                version, test_package, target_branch, tools_git_repo
-            )
+            if test_platform == AMDGPU
+                job_dict["unit_test_julia_$(test_platform_name)_$(replace(version, "." => "_"))"] = _get_amdgpu_unit_test(
+                    version, test_package, target_branch, tools_git_repo
+                )
+            else
+                job_dict["unit_test_julia_$(test_platform_name)_$(replace(version, "." => "_"))"] = _get_normal_unit_test(
+                    version, test_package, target_branch, test_platform, tools_git_repo
+                )
+            end
         else
-            job_dict["unit_test_julia_nightly"] = _get_nightly_unit_test(
-                test_package, target_branch, tools_git_repo, nightly_base_image
+            job_dict["unit_test_julia_$(test_platform_name)_nightly"] = _get_nightly_unit_test(
+                test_package,
+                target_branch,
+                test_platform,
+                tools_git_repo,
+                nightly_base_image,
             )
         end
     end
@@ -79,6 +101,9 @@ function add_unit_test_verify_job_yaml!(
 )
     # verification script that no custom URLs are used in unit tests
     if target_branch != "main"
+        if !haskey(job_dict, "stages")
+            job_dict["stages"] = []
+        end
         push!(job_dict["stages"], "verify-unit-test-deps")
         job_dict["verify-unit-test-deps"] = Dict(
             "image" => "julia:1.10",
@@ -99,6 +124,7 @@ end
         version::AbstractString,
         test_package::TestPackage,
         target_branch::AbstractString,
+        test_platform::TestPlatform,
         tools_git_repo::ToolsGitRepo,
     )
 
@@ -108,6 +134,7 @@ Creates a normal unit test job for a specific Julia version.
 - `version::AbstractString`: Julia version used for the tests.
 - `test_package::TestPackage`: Properties of the package to be tested, such as name and version.
 - `target_branch::AbstractString`: A different job code is generated depending on the target branch.
+- `test_platform::TestPlatform`: Set target platform test, e.g. CPU, Nvidia GPU or AMD GPU.
 - `tools_git_repo::ToolsGitRepo`: URL and branch of the Git repository from which the CI tools are
     cloned in unit test job.
 
@@ -119,6 +146,7 @@ function _get_normal_unit_test(
     version::AbstractString,
     test_package::TestPackage,
     target_branch::AbstractString,
+    test_platform::TestPlatform,
     tools_git_repo::ToolsGitRepo,
 )::Dict
     job_yaml = Dict()
@@ -130,6 +158,14 @@ function _get_normal_unit_test(
     )
     job_yaml["image"] = "julia:$(version)"
 
+    if !haskey(job_yaml, "variables")
+        job_yaml["variables"] = Dict()
+    end
+
+    for tp in instances(TestPlatform)
+        job_yaml["variables"]["TEST_$(tp)"] = (tp == test_platform) ? "1" : "0"
+    end
+
     script = [
         "apt update && apt install -y git",
         "git clone --depth 1 -b $(tools_git_repo.branch) $(tools_git_repo.url) /tmp/integration_test_tools/",
@@ -157,7 +193,20 @@ function _get_normal_unit_test(
     job_yaml["script"] = script
 
     job_yaml["interruptible"] = true
-    job_yaml["tags"] = ["cpuonly"]
+
+    if test_platform == CPU
+        job_yaml["tags"] = ["cpuonly"]
+    elseif test_platform == CUDA
+        job_yaml["tags"] = ["cuda", "x86_64"]
+    elseif test_platform == AMDGPU
+        job_yaml["tags"] = ["rocm", "x86_64"]
+    else
+        throw(
+            ArgumentError(
+                "test_platform argument with value $(test_platform) not supported"
+            ),
+        )
+    end
 
     return job_yaml
 end
@@ -166,13 +215,16 @@ end
     _get_nightly_unit_test(
         test_package::TestPackage,
         target_branch::AbstractString,
+        test_platform::TestPlatform,
         tools_git_repo::ToolsGitRepo,
+        nightly_base_image::AbstractString,
     )
 
 Creates a unit test job which uses the Julia nightly version.
 
 # Args
 - `test_package::TestPackage`: Properties of the package to be tested, such as name and version.
+- `test_platform::TestPlatform`: Set target platform test, e.g. CPU, Nvidia GPU or AMD GPU.
 - `target_branch::AbstractString`: A different job code is generated depending on the target branch.
 - `tools_git_repo::ToolsGitRepo`: URL and branch of the Git repository from which the CI tools are
     cloned in unit test job.
@@ -185,10 +237,13 @@ Returns a dict containing the unit test, which can be output directly as GitLab
 function _get_nightly_unit_test(
     test_package::TestPackage,
     target_branch::AbstractString,
+    test_platform::TestPlatform,
     tools_git_repo::ToolsGitRepo,
     nightly_base_image::AbstractString,
 )
-    job_yaml = _get_normal_unit_test("1", test_package, target_branch, tools_git_repo)
+    job_yaml = _get_normal_unit_test(
+        "1", test_package, target_branch, test_platform, tools_git_repo
+    )
     job_yaml["image"] = nightly_base_image
 
     if !haskey(job_yaml, "variables")
@@ -218,8 +273,50 @@ fi",
         "cp -r \$JULIA_EXTRACT_FOLDER/* /usr",
     ]
 
-    job_yaml["image"] = "debian:bookworm-slim"
-
     job_yaml["allow_failure"] = true
     return job_yaml
 end
+
+"""
+    _get_amdgpu_unit_test(
+        version::AbstractString,
+        test_package::TestPackage,
+        target_branch::AbstractString,
+        tools_git_repo::ToolsGitRepo,
+    )
+
+Creates a amdgpu unit test job for a specific Julia version. Use a different base image and install 
+Julia in it.
+
+# Args
+- `version::AbstractString`: Julia version used for the tests.
+- `test_package::TestPackage`: Properties of the package to be tested, such as name and version.
+- `target_branch::AbstractString`: A different job code is generated depending on the target branch.
+- `tools_git_repo::ToolsGitRepo`: URL and branch of the Git repository from which the CI tools are
+    cloned in unit test job.
+
+Return
+
+Returns a dict containing the unit test, which can be output directly as GitLab CI yaml.
+"""
+function _get_amdgpu_unit_test(
+    version::AbstractString,
+    test_package::TestPackage,
+    target_branch::AbstractString,
+    tools_git_repo::ToolsGitRepo,
+)::Dict
+    job_yaml = _get_normal_unit_test(
+        version, test_package, target_branch, AMDGPU, tools_git_repo
+    )
+    job_yaml["image"] = "rocm/dev-ubuntu-24.04:6.2.4-complete"
+    job_yaml["before_script"] = [
+        "curl -fsSL https://install.julialang.org | sh -s -- -y -p /julia",
+        "export PATH=/julia/bin:\$PATH",
+        "echo \$PATH",
+        "juliaup add $(version)",
+        "juliaup default $(version)",
+    ]
+    job_yaml["tags"] = ["rocm", "x86_64"]
+
+    return job_yaml
+end

.ci/CI/test/UnitTest/amdgpu_unit_test.jl

@@ -0,0 +1,74 @@
+@testset "test _get_amdgpu_unit_test" begin
+    test_package = CI.TestPackage("QEDfoo", "/path/to/project", "42.0")
+    git_url = "http://github.com/name/repo"
+    git_branch = "branch"
+
+    expected_job = get_generic_unit_job("1.11", test_package)
+
+    expected_job["script"] = get_main_unit_job_script_section(git_url, git_branch)
+
+    expected_job["image"] = "rocm/dev-ubuntu-24.04:6.2.4-complete"
+    expected_job["before_script"] = get_amdgpu_before_script("1.11")
+    expected_job["variables"]["TEST_AMDGPU"] = "1"
+    expected_job["tags"] = ["rocm", "x86_64"]
+
+    job_yaml = CI._get_amdgpu_unit_test(
+        "1.11", test_package, "main", CI.ToolsGitRepo(git_url, git_branch)
+    )
+
+    @test keys(expected_job) == keys(job_yaml)
+
+    for k in keys(expected_job)
+        @test (
+            @assert job_yaml[k] == expected_job[k] (
+                "\nkey: \"$(k)\"\n:" * yaml_diff(job_yaml[k], expected_job[k])
+            );
+            true
+        )
+    end
+
+    @test (@assert job_yaml == expected_job yaml_diff(job_yaml, expected_job);
+    true)
+
+    @testset "test public interface" begin
+        julia_versions = Vector{String}(["1.9", "1.10", "1.11"])
+
+        job_dict = Dict()
+        CI.add_unit_test_job_yaml!(
+            job_dict,
+            test_package,
+            julia_versions,
+            "main",
+            CI.AMDGPU,
+            CI.ToolsGitRepo(git_url, git_branch),
+        )
+
+        for julia_version in julia_versions
+            job_name = "unit_test_julia_amdgpu_$(replace(julia_version, "." => "_"))"
+            expected_job["before_script"] = get_amdgpu_before_script(julia_version)
+
+            @test haskey(job_dict, job_name)
+            @test job_dict["stages"] == ["unit-test"]
+
+            unit_test_job = job_dict[job_name]
+
+            @test keys(expected_job) == keys(unit_test_job)
+
+            for k in keys(expected_job)
+                @test (
+                    @assert unit_test_job[k] == expected_job[k] (
+                        "\nkey: \"$(k)\"\n:" * yaml_diff(unit_test_job[k], expected_job[k])
+                    );
+                    true
+                )
+            end
+
+            @test (
+                @assert unit_test_job == expected_job yaml_diff(
+                    unit_test_job, expected_job
+                );
+                true
+            )
+        end
+    end
+end