Try #269:

JuliaGPU · Oct 8, 2021 · 12fd5ad · 12fd5ad
2 parents 045fab2 + 4e49edf
commit 12fd5ad
Show file tree

Hide file tree

Showing 8 changed files with 46 additions and 10 deletions.
diff --git a/examples/matmul.jl b/examples/matmul.jl
@@ -23,11 +23,9 @@ function matmul!(a, b, c)
         println("Matrix size mismatch!")
         return nothing
     end
-    if isa(a, Array)
-        kernel! = matmul_kernel!(CPU(),4)
-    else
-        kernel! = matmul_kernel!(CUDADevice(),256)
-    end
+    device = KernelAbstractions.get_device(a)
+    n = device isa GPU ? 256 : 4
+    kernel! = matmul_kernel!(device, n)
     kernel!(a, b, c, ndrange=size(c)) 
 end
 

diff --git a/examples/naive_transpose.jl b/examples/naive_transpose.jl
@@ -16,11 +16,9 @@ function naive_transpose!(a, b)
         println("Matrix size mismatch!")
         return nothing
     end
-    if isa(a, Array)
-        kernel! = naive_transpose_kernel!(CPU(),4)
-    else
-        kernel! = naive_transpose_kernel!(CUDADevice(),256)
-    end
+    device = KernelAbstractions.get_device(a)
+    n = device isa GPU ? 256 : 4
+    kernel! = naive_transpose_kernel!(device, n)
     kernel!(a, b, ndrange=size(a))
 end
 

diff --git a/lib/CUDAKernels/src/CUDAKernels.jl b/lib/CUDAKernels/src/CUDAKernels.jl
@@ -10,6 +10,10 @@ import KernelAbstractions
 
 export CUDADevice
 
+KernelAbstractions.get_device(::Type{<:CUDA.CuArray}) = CUDADevice()
+KernelAbstractions.get_device(::Type{<:CUDA.CUSPARSE.AbstractCuSparseArray}) = CUDADevice()
+
+
 const FREE_STREAMS = CUDA.CuStream[]
 const STREAMS = CUDA.CuStream[]
 const STREAM_GC_THRESHOLD = Ref{Int}(16)

diff --git a/lib/CUDAKernels/test/runtests.jl b/lib/CUDAKernels/test/runtests.jl
@@ -8,6 +8,11 @@ using Test
 include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl"))
 include(joinpath(dirname(pathof(KernelGradients)), "..", "test", "testsuite.jl"))
 
+@testset "get_device" begin
+    @test @inferred(KernelAbstractions.get_device(CUDA.CuArray{Float32,3})) == CUDADevice()
+    @test @inferred(KernelAbstractions.get_device(CUDA.CUSPARSE.CuSparseMatrixCSC{Float32})) == CUDADevice()
+end
+
 if parse(Bool, get(ENV, "CI", "false"))
     default = "CPU"
 else

diff --git a/lib/ROCKernels/src/ROCKernels.jl b/lib/ROCKernels/src/ROCKernels.jl
@@ -11,6 +11,9 @@ import KernelAbstractions
 
 export ROCDevice
 
+KernelAbstractions.get_device(::Type{<:AMDGPU.ROCArray}) = ROCDevice()
+
+
 const FREE_QUEUES = HSAQueue[]
 const QUEUES = HSAQueue[]
 const QUEUE_GC_THRESHOLD = Ref{Int}(16)

diff --git a/lib/ROCKernels/test/runtests.jl b/lib/ROCKernels/test/runtests.jl
@@ -8,6 +8,10 @@ using Test
 include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl"))
 include(joinpath(dirname(pathof(KernelGradients)), "..", "test", "testsuite.jl"))
 
+@test "get_device" begin
+    @test @inferred(KernelAbstractions.get_device(AMDGPU.ROCArray{Float32,3})) == ROCDevice()
+end
+
 if parse(Bool, get(ENV, "CI", "false"))
     default = "CPU"
 else

diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl
@@ -337,6 +337,24 @@ abstract type GPU <: Device end
 
 struct CPU <: Device end
 
+
+"""
+    KernelAbstractions.get_device(A::AbstractArray)::KernelAbstractions.Device
+    KernelAbstractions.get_device(TA::Type{<:AbstractArray})::KernelAbstractions.Device
+
+Get a `KernelAbstractions.Device` instance suitable for array `A` resp. array
+type `TA`.
+"""
+function get_device end
+
+get_device(A::AbstractArray) = get_device(typeof(A))
+
+get_device(::Type{<:AbstractArray}) = CPU()
+
+# Would require dependency on GPUArrays:
+# get_device(TA::Type{<:GPUArrays.AbstractGPUArray}) = throw(ArgumentError("NoKernelAbstractions.Device type defined for arrays of type $(TA.name.name)"))
+
+
 include("nditeration.jl")
 using .NDIteration
 import .NDIteration: get

diff --git a/test/test.jl b/test/test.jl
@@ -64,6 +64,12 @@ end
        A[I] = i
 end
 
+@testset "get_device" begin
+    A = rand(5)
+    @test @inferred(KernelAbstractions.get_device(typeof(A))) == CPU()
+    @test @inferred(KernelAbstractions.get_device(A)) == KernelAbstractions.get_device(typeof(A))
+end
+
 @testset "indextest" begin
     # TODO: add test for _group and _local_cartesian
     A = ArrayT{Int}(undef, 16, 16)