Try #269:

JuliaGPU · Oct 7, 2021 · 264a7ba · 264a7ba
2 parents 045fab2 + c555ebd
commit 264a7ba
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 0 deletions.
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)