From bc28f147a1e1a279ed6d789d3a1f52aa6876a42e Mon Sep 17 00:00:00 2001
From: chengchingwen <chengchingwen214@gmail.com>
Date: Fri, 26 Apr 2024 13:59:24 +0800
Subject: [PATCH] update test for switching gpu backend

---
 test/collapseddims.jl |  2 +-
 test/functional.jl    | 12 +++++-----
 test/mask.jl          |  2 +-
 test/matmul.jl        |  4 ++--
 test/runtests.jl      | 53 +++++++++++++++++++++----------------------
 5 files changed, 36 insertions(+), 37 deletions(-)

diff --git a/test/collapseddims.jl b/test/collapseddims.jl
index 9cef2b8..03cc45f 100644
--- a/test/collapseddims.jl
+++ b/test/collapseddims.jl
@@ -1,4 +1,4 @@
-if !USE_CUDA
+if !USE_GPU
     @testset "CollapsedDim" begin
         using NeuralAttentionlib.Matmul
         using NeuralAttentionlib: collapseddims_fdim1, collapseddims_nonbatch, collapseddims_nonbatch_fdim1
diff --git a/test/functional.jl b/test/functional.jl
index 2be1488..edeac52 100644
--- a/test/functional.jl
+++ b/test/functional.jl
@@ -13,7 +13,7 @@
       layer_norm, rms_layer_norm, get_sincos_position_embeddings
 
     @testset "score" begin
-        if !USE_CUDA
+        if !USE_GPU
             @testset "AD" begin
                 test_rrule(dot_product_score, randn(5, 3, 2), randn(5, 4, 2); check_inferred = false)
                 test_rrule(dot_product_score, randn(5, 3, 2, 2), randn(5, 4, 2, 2))
@@ -84,7 +84,7 @@
             end
         end
 
-        if !USE_CUDA
+        if !USE_GPU
             @testset "AD" begin
                 test_rrule(
                     scalar_relative_position_embedding, t5_bucketed_position_id(8, 20), randn(3, 8),
@@ -182,7 +182,7 @@
         @test with_rotary_position_embedding(x) ≈ naive_rotary_pe(x)
         @test with_rotary_position_embedding(256, x) ≈ naive_rotary_pe_w_dim(256, x)
         @test with_rotary_position_embedding(256)(x) ≈ naive_rotary_pe_w_dim(256, x)
-        if !USE_CUDA
+        if !USE_GPU
             @testset "AD" begin
                 x = randn(512, 5, 3, 2)
                 @test Zygote.gradient(x->sum(sin.(with_rotary_position_embedding(x))), x)[1] ≈
@@ -226,7 +226,7 @@
             atol = 5e-1
         )
 
-        if !USE_CUDA
+        if !USE_GPU
             @testset "AD" begin
                 g = randn(20)
                 b = randn(20)
@@ -258,7 +258,7 @@
 
     @testset "attention" begin
         @testset "multihead_qkv_attention" begin
-            if !USE_CUDA
+            if !USE_GPU
                 @testset "AD" begin
                     for i = 1:3
                         a = randn(20, 3, 2)
@@ -296,7 +296,7 @@
             @test grad[2] ≈ ngrad[2]
             @test grad[3] ≈ ngrad[3]
 
-            if !USE_CUDA
+            if !USE_GPU
                 @testset "AD" begin
                     for i = 1:3
                         a = randn(30, 3, 2)
diff --git a/test/mask.jl b/test/mask.jl
index 47aca7a..921afee 100644
--- a/test/mask.jl
+++ b/test/mask.jl
@@ -285,7 +285,7 @@
         @test_throws DimensionMismatch drandn(5, 4) .* (GenericAttenMask(drand(Bool, 3, 4)) | SymLengthMask([2]))
     end
 
-    if !USE_CUDA
+    if !USE_GPU
         @testset "AD" begin
             m = (LocalMask(1) | CausalMask() & !(BandPartMask(5,5)) | BiLengthMask([2,3], [3, 7]))
 
diff --git a/test/matmul.jl b/test/matmul.jl
index 15091f7..283cd3e 100644
--- a/test/matmul.jl
+++ b/test/matmul.jl
@@ -18,7 +18,7 @@
     end
     uwcs(x) = size(unwrap_collapse(x))
 
-    if USE_CUDA
+    if USE_GPU
         eltype_list = (Float64, Float32, Float16, ComplexF64, ComplexF32)
     else
         eltype_list = (Float64, Float32, ComplexF64, ComplexF32)
@@ -178,7 +178,7 @@
         end
     end
 
-    if !USE_CUDA
+    if !USE_GPU
         @testset "AD" begin
             test_rrule(matmul, randn(7,6,5), randn(6, 2), randn())
             test_rrule(matmul, randn(7,6,5,4), randn(6), randn())
diff --git a/test/runtests.jl b/test/runtests.jl
index 5ee2300..b89f06d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -24,40 +24,39 @@ include("old_impl/old_impl.jl")
 using .Old_Impl
 using .Old_Impl: batched_triu!, batched_tril!
 
-function should_test_cuda()
-    e = get(ENV, "JL_PKG_TEST_CUDA", false)
-    e isa Bool && return e
+function testing_gpu()
+    e = get(ENV, "JL_PKG_TEST_GPU", nothing)
+    isnothing(e) && return nothing
     if e isa String
-        x = tryparse(Bool, e)
-        return isnothing(x) ? false : x
-    else
-        return false
+        x = lowercase(e)
+        if isempty(x)
+            return nothing
+        elseif x == "cuda"
+            return :cuda
+        elseif x == "amdgpu"
+            return :amdgpu
+        end
     end
+    error("Unknown value for `JL_PKG_TEST_GPU`: $x")
 end
 
-function should_test_amdgpu()
-    e = get(ENV, "JL_PKG_TEST_AMDGPU", false)
-    e isa Bool && return e
-    if e isa String
-        x = tryparse(Bool, e)
-        return isnothing(x) ? false : x
-    else
-        return false
+const GPUBACKEND = testing_gpu()
+if isnothing(GPUBACKEND)
+    const USE_GPU = false
+else
+    const USE_GPU = true
+    if GPUBACKEND == :cuda
+        using CUDA
+        CUDA.allowscalar(false)
+    elseif GPUBACKEND == :amdgpu
+        using AMDGPU
+        AMDGPU.allowscalar(false)
     end
 end
+@show GPUBACKEND
+@show USE_GPU
 
-const USE_CUDA = @show should_test_cuda()
-const USE_AMDGPU = @show should_test_amdgpu()
-
-if USE_CUDA
-    CUDA.allowscalar(false)
-end
-
-if USE_AMDGPU
-    AMDGPU.allowscalar(false)
-end
-
-device(x) = USE_CUDA || USE_AMDGPU ? gpu(x) : x
+device(x) = USE_GPU ? gpu(x) : x
 
 drandn(arg...) = randn(arg...) |> device
 drand(arg...) = rand(arg...) |> device