Merge pull request #25 from JuliaGPU/jps/ka-function

Add Kernel helper for KA kernels
JuliaGPU · Jul 17, 2023 · 5e50b4b · 5e50b4b
2 parents 1e0cc54 + 5990d37
commit 5e50b4b
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/src/DaggerGPU.jl b/src/DaggerGPU.jl
@@ -8,6 +8,17 @@ import Dagger: Chunk
 
 const CPUProc = Union{OSProc, Dagger.ThreadProc}
 
+struct Kernel{F} end
+Kernel(f) = Kernel{f}()
+
+function (::Kernel{F})(args...; ndrange) where F
+    @nospecialize args
+    dev = kernel_backend()
+    kern = F(dev)
+    kern(args...; ndrange)
+    KernelAbstractions.synchronize(dev)
+end
+
 macro gpuproc(PROC, T)
     PROC = esc(PROC)
     T = esc(T)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -13,14 +13,15 @@ addprocs(2, exeflags="--project")
     catch end
 
     using Distributed, Dagger, DaggerGPU
+    import DaggerGPU: Kernel
     using KernelAbstractions
 end
 @everywhere begin
     function isongpu(X)
         return !(X isa Array)
     end
 
-    KernelAbstractions.@kernel function fill_kernel(A, x)
+    @kernel function fill_kernel(A, x)
         idx = @index(Global, Linear)
         A[idx] = x
     end
@@ -32,6 +33,11 @@ end
         return A, typeof(A)
     end
 
+    @kernel function copy_kernel(B, A)
+        idx = @index(Global, Linear)
+        B[idx] = A[idx]
+    end
+
     # Create a function to perform an in-place operation.
     function addarray!(x)
         x .= x .+ 1.0f0
@@ -53,6 +59,13 @@ end
         DA, T = fetch(Dagger.@spawn fill_thunk(A, 2.3f0))
         @test all(DA .== 2.3f0)
         @test T <: Array
+
+        A = rand(Float64, 128)
+        B = zeros(Float64, 128)
+        Dagger.with_options(scope=Dagger.scope(worker=1,thread=1)) do
+            fetch(Dagger.@spawn Kernel(copy_kernel)(B, A; ndrange=length(A)))
+        end
+        @test all(B .== A)
     end
 end
 
@@ -81,6 +94,13 @@ end
             end
             @test all(DA .== 2.3f0)
             @test T <: CuArray
+
+            A = CUDA.rand(128)
+            B = CUDA.zeros(128)
+            Dagger.with_options(;scope=Dagger.scope(worker=1,cuda_gpu=1)) do
+                fetch(Dagger.@spawn Kernel(copy_kernel)(B, A; ndrange=length(A)))
+            end
+            @test all(B .== A)
         end
     end
 end
@@ -110,6 +130,13 @@ end
             end
             @test all(DA .== 2.3f0)
             @test T <: ROCArray
+
+            A = AMDGPU.rand(128)
+            B = AMDGPU.zeros(128)
+            Dagger.with_options(;scope=Dagger.scope(worker=1,rocm_gpu=1)) do
+                fetch(Dagger.@spawn Kernel(copy_kernel)(B, A; ndrange=length(A)))
+            end
+            @test all(B .== A)
         end
     end
 end
@@ -139,6 +166,13 @@ end
             end
             @test all(DA .== 2.3f0)
             @test T <: MtlArray
+
+            A = Metal.rand(128)
+            B = Metal.zeros(128)
+            Dagger.with_options(;scope=Dagger.scope(worker=1,metal_gpu=1)) do
+                fetch(Dagger.@spawn Kernel(copy_kernel)(B, A; ndrange=length(A)))
+            end
+            @test all(B .== A)
         end
 
         @testset "In-place operations" begin