Try #269:

Author: bors[bot]

Project.toml

@@ -7,7 +7,9 @@ version = "0.7.0"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

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
 

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
 

lib/CUDAKernels/src/CUDAKernels.jl

@@ -10,6 +10,9 @@ import KernelAbstractions
 
 export CUDADevice
 
+KernelAbstractions.get_device(::CUDA.CuArray) = CUDADevice()
+KernelAbstractions.get_device(::CUDA.CUSPARSE.AbstractCuSparseArray) = CUDADevice()
+
 const FREE_STREAMS = CUDA.CuStream[]
 const STREAMS = CUDA.CuStream[]
 const STREAM_GC_THRESHOLD = Ref{Int}(16)

lib/CUDAKernels/test/Project.toml

@@ -4,8 +4,10 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KernelGradients = "e5faadeb-7f6c-408e-9747-a7a26e81c66a"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

lib/CUDAKernels/test/runtests.jl

@@ -4,10 +4,16 @@ using Enzyme
 using CUDA
 using CUDAKernels
 using Test
+using SparseArrays
 
 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(CuArray(rand(Float32, 3,3)))) == CUDADevice()
+    @test @inferred(KernelAbstractions.get_device(CuArray(sparse(rand(Float32, 3,3))))) == CUDADevice()
+end
+
 if parse(Bool, get(ENV, "CI", "false"))
     default = "CPU"
 else

lib/ROCKernels/src/ROCKernels.jl

@@ -11,6 +11,9 @@ import KernelAbstractions
 
 export ROCDevice
 
+KernelAbstractions.get_device(::AMDGPU.ROCArray) = ROCDevice()
+
+
 const FREE_QUEUES = HSAQueue[]
 const QUEUES = HSAQueue[]
 const QUEUE_GC_THRESHOLD = Ref{Int}(16)

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(ROCArray(rand(Float32, 3, 3)))) == ROCDevice()
+end
+
 if parse(Bool, get(ENV, "CI", "false"))
     default = "CPU"
 else

src/KernelAbstractions.jl

@@ -6,7 +6,9 @@ export Device, GPU, CPU, Event, MultiEvent, NoneEvent
 export async_copy!
 
 
+using LinearAlgebra
 using MacroTools
+using SparseArrays
 using StaticArrays
 using Cassette
 using Adapt
@@ -337,6 +339,23 @@ abstract type GPU <: Device end
 
 struct CPU <: Device end
 
+
+"""
+    KernelAbstractions.get_device(A::AbstractArray)::KernelAbstractions.Device
+
+Get a `KernelAbstractions.Device` instance suitable for array `A`.
+"""
+function get_device end
+
+# Should cover SubArray, ReshapedArray, ReinterpretArray, Hermitian, AbstractTriangular, etc.:
+get_device(A::AbstractArray) = get_device(parent(A))
+
+get_device(A::AbstractSparseArray) = get_device(rowvals(A))
+get_device(A::Diagonal) = get_device(A.diag)
+get_device(A::Tridiagonal) = get_device(A.d)
+
+get_device(::Array) = CPU()
+
 include("nditeration.jl")
 using .NDIteration
 import .NDIteration: get

test/test.jl

@@ -1,6 +1,8 @@
 using KernelAbstractions
 using KernelAbstractions.NDIteration
 using InteractiveUtils
+using LinearAlgebra
+using SparseArrays
 import SpecialFunctions
 
 identity(x) = x
@@ -64,6 +66,16 @@ end
        A[I] = i
 end
 
+@testset "get_device" begin
+    x = rand(5)
+    A = rand(5,5)
+    @test @inferred(KernelAbstractions.get_device(A)) == CPU()
+    @test @inferred(KernelAbstractions.get_device(view(A, 2:4, 1:3))) == CPU()
+    @test @inferred(KernelAbstractions.get_device(sparse(A))) == CPU()
+    @test @inferred(KernelAbstractions.get_device(Diagonal(x))) == CPU()
+    @test @inferred(KernelAbstractions.get_device(Tridiagonal(A))) == CPU()
+end
+
 @testset "indextest" begin
     # TODO: add test for _group and _local_cartesian
     A = ArrayT{Int}(undef, 16, 16)