test/GPU: Factor out scopes to test

Author: jpsamaroo

test/gpu.jl

@@ -1,64 +1,7 @@
 using Random
 using LinearAlgebra
 
-if USE_CUDA
-    using Pkg
-    Pkg.add("CUDA")
-end
-if USE_ROCM
-    using Pkg
-    Pkg.add("AMDGPU")
-end
-if USE_ONEAPI
-    using Pkg
-    Pkg.add("oneAPI")
-end
-if USE_METAL
-    using Pkg
-    Pkg.add("Metal")
-end
-if USE_OPENCL
-    using Pkg
-    Pkg.add("OpenCL")
-    Pkg.add("pocl_jll")
-end
-
 @everywhere begin
-    if $USE_CUDA
-        using CUDA
-    elseif !$IN_CI
-        try using CUDA
-        catch end
-    end
-
-    if $USE_ROCM
-        using AMDGPU
-    elseif !$IN_CI
-        try using AMDGPU
-        catch end
-    end
-
-    if $USE_ONEAPI
-        using oneAPI
-    elseif !$IN_CI
-        try using oneAPI
-        catch end
-    end
-
-    if $USE_METAL
-        using Metal
-    elseif !$IN_CI
-        try using Metal
-        catch end
-    end
-
-    if $USE_OPENCL
-        using pocl_jll, OpenCL
-    elseif !$IN_CI
-        try using pocl_jll, OpenCL
-        catch end
-    end
-
     using Distributed, Dagger
     import Dagger: Kernel
     using KernelAbstractions

test/runtests.jl

@@ -136,6 +136,11 @@ using Dagger
 using UUIDs
 import MemPool
 
+GPU_SCOPES = Pair{Symbol, Dagger.AbstractScope}[]
+if USE_GPU
+    include("setup_gpu.jl")
+end
+
 try
     for test in to_test
         test_title = tests[findfirst(x->x[2]==test * ".jl", tests)][1]

test/setup_gpu.jl

@@ -0,0 +1,89 @@
+if USE_CUDA
+    using Pkg
+    Pkg.add("CUDA")
+end
+if USE_ROCM
+    using Pkg
+    Pkg.add("AMDGPU")
+end
+if USE_ONEAPI
+    using Pkg
+    Pkg.add("oneAPI")
+end
+if USE_METAL
+    using Pkg
+    Pkg.add("Metal")
+end
+if USE_OPENCL
+    using Pkg
+    Pkg.add("OpenCL")
+    Pkg.add("pocl_jll")
+end
+
+@everywhere begin
+    if $USE_CUDA
+        using CUDA
+    elseif !$IN_CI
+        try using CUDA
+        catch end
+    end
+
+    if $USE_ROCM
+        using AMDGPU
+    elseif !$IN_CI
+        try using AMDGPU
+        catch end
+    end
+
+    if $USE_ONEAPI
+        using oneAPI
+    elseif !$IN_CI
+        try using oneAPI
+        catch end
+    end
+
+    if $USE_METAL
+        using Metal
+    elseif !$IN_CI
+        try using Metal
+        catch end
+    end
+
+    if $USE_OPENCL
+        using pocl_jll, OpenCL
+    elseif !$IN_CI
+        try using pocl_jll, OpenCL
+        catch end
+    end
+end
+
+if USE_CUDA
+    push!(GPU_SCOPES, :CUDA => Dagger.scope(;worker=1, cuda_gpu=1))
+    if length(CUDA.devices()) > 1
+        push!(GPU_SCOPES, :CUDA => Dagger.scope(;worker=1, cuda_gpu=2))
+    end
+end
+if USE_ROCM
+    push!(GPU_SCOPES, :ROCm => Dagger.scope(;worker=1, rocm_gpu=1))
+    if length(AMDGPU.devices()) > 1
+        push!(GPU_SCOPES, :ROCm => Dagger.scope(;worker=1, rocm_gpu=2))
+    end
+end
+if USE_ONEAPI
+    push!(GPU_SCOPES, :oneAPI => Dagger.scope(;worker=1, intel_gpu=1))
+    if length(oneAPI.devices()) > 1
+        push!(GPU_SCOPES, :oneAPI => Dagger.scope(;worker=1, intel_gpu=2))
+    end
+end
+if USE_METAL
+    push!(GPU_SCOPES, :Metal => Dagger.scope(;worker=1, metal_gpu=1))
+    if length(Metal.devices()) > 1
+        push!(GPU_SCOPES, :Metal => Dagger.scope(;worker=1, metal_gpu=2))
+    end
+end
+if USE_OPENCL
+    push!(GPU_SCOPES, :OpenCL => Dagger.scope(;worker=1, cl_device=1))
+    if length(cl.devices(cl.default_platform())) > 1
+        push!(GPU_SCOPES, :OpenCL => Dagger.scope(;worker=1, cl_device=2))
+    end
+end