add get_backend for StaticArrays (#621)

Author: vchuravy

Project.toml

@@ -1,7 +1,7 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 authors = ["Valentin Churavy <[email protected]> and contributors"]
-version = "0.9.35"
+version = "0.9.36"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/StaticArraysExt.jl

@@ -0,0 +1,9 @@
+module StaticArraysExt
+
+import KernelAbstractions: get_backend, CPU
+using StaticArrays: SizedArray, MArray
+
+get_backend(A::SizedArray) = get_backend(A.data)
+get_backend(::MArray) = CPU()
+
+end

src/KernelAbstractions.jl

@@ -865,4 +865,7 @@ if !isdefined(Base, :get_extension)
     include("../ext/SparseArraysExt.jl")
 end
 
+# Currently we have a direct dependency on StaticArrays
+include("../ext/StaticArraysExt.jl")
+
 end #module

test/runtests.jl

@@ -72,7 +72,6 @@ struct NewBackend <: KernelAbstractions.GPU end
     @test_throws MethodError kernel()
 end
 
-
 include("extensions/enzyme.jl")
 @static if VERSION >= v"1.7.0"
     @testset "Enzyme" begin

test/test.jl

@@ -4,6 +4,7 @@ using InteractiveUtils
 using LinearAlgebra
 using SparseArrays
 using Adapt
+using StaticArrays
 
 identity(x) = x
 
@@ -95,6 +96,18 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
         @test @inferred(KernelAbstractions.get_backend(sparse(A))) isa backendT
     end
 
+    @conditional_testset "StaticArrays" skip_tests begin
+        backend = Backend()
+        backendT = typeof(backend).name.wrapper # To look through CUDABackend{true, false}
+        @test backend isa backendT
+
+        @test KernelAbstractions.get_backend(@MMatrix [1.0]) isa CPU
+        @test_throws ArgumentError KernelAbstractions.get_backend(@SMatrix [1.0])
+
+        A = allocate(backend, Float32, 5, 5)
+        @test @inferred(KernelAbstractions.get_backend(SizedMatrix{5, 5}(A))) isa backendT
+    end
+
     @conditional_testset "adapt" skip_tests begin
         backend = Backend()
         x = allocate(backend, Float32, 5)