add get_backend for StaticArrays

Author: vchuravy

ext/StaticArraysExt.jl

@@ -0,0 +1,11 @@
+module StaticArraysExt
+
+import KernelAbstractions: get_backend, CPU
+using StaticArrays: SizedArray, MArray, SArray
+
+get_backend(A::SizedArray) = get_backend(A.data)
+get_backend(::MArray) = CPU()
+# TODO: It makes sense to pass SArray to the GPU backend, so we can't make a determination
+# get_backend(::SArray) = 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)