From d33b2e30c9e7cb8f62833468b8320fcf281173a9 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@iitk.ac.in>
Date: Sat, 24 Feb 2024 12:47:27 -0500
Subject: [PATCH] Add tests

---
 Project.toml   |  2 +-
 test/amdgpu.jl | 23 +++++++++++++++++++++++
 test/cuda.jl   | 23 +++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 8e83cce..f78a118 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxDeviceUtils"
 uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.1.15"
+version = "0.1.16"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
diff --git a/test/amdgpu.jl b/test/amdgpu.jl
index 68e8db0..3675a0e 100644
--- a/test/amdgpu.jl
+++ b/test/amdgpu.jl
@@ -72,3 +72,26 @@ using FillArrays, Zygote  # Extensions
         @test ps_cpu.farray isa Fill
     end
 end
+
+if LuxAMDGPU.functional()
+    ps = (; weight=rand(Float32, 10), bias=rand(Float32, 10))
+    ps_cpu = deepcopy(ps)
+    cdev = cpu_device()
+    for idx in 1:length(AMDGPU.devices())
+        amdgpu_device = gpu_device(idx)
+        @test typeof(amdgpu_device.device) <: AMDGPU.HIPDevice
+        @test AMDGPU.device_id(amdgpu_device.device) == idx
+
+        ps = ps |> amdgpu_device
+        @test ps.weight isa ROCArray
+        @test ps.bias isa ROCArray
+        @test AMDGPU.device_id(AMDGPU.device(ps.weight)) == idx
+        @test AMDGPU.device_id(AMDGPU.device(ps.bias)) == idx
+        @test isequal(cdev(ps.weight), ps_cpu.weight)
+        @test isequal(cdev(ps.bias), ps_cpu.bias)
+    end
+
+    ps = ps |> cdev
+    @test ps.weight isa Array
+    @test ps.bias isa Array
+end
diff --git a/test/cuda.jl b/test/cuda.jl
index 613f132..9a7c2c3 100644
--- a/test/cuda.jl
+++ b/test/cuda.jl
@@ -72,3 +72,26 @@ using FillArrays, Zygote  # Extensions
         @test ps_cpu.farray isa Fill
     end
 end
+
+if LuxCUDA.functional()
+    ps = (; weight=rand(Float32, 10), bias=rand(Float32, 10))
+    ps_cpu = deepcopy(ps)
+    cdev = cpu_device()
+    for idx in 1:length(CUDA.devices())
+        cuda_device = gpu_device(idx)
+        @test typeof(cuda_device.device) <: CUDA.CuDevice
+        @test cuda_device.device.handle == (idx - 1)
+
+        ps = ps |> cuda_device
+        @test ps.weight isa CuArray
+        @test ps.bias isa CuArray
+        @test CUDA.device(ps.weight).handle == idx - 1
+        @test CUDA.device(ps.bias).handle == idx - 1
+        @test isequal(cdev(ps.weight), ps_cpu.weight)
+        @test isequal(cdev(ps.bias), ps_cpu.bias)
+    end
+
+    ps = ps |> cdev
+    @test ps.weight isa Array
+    @test ps.bias isa Array
+end