Merge pull request #455 from vchuravy/vc/ka

Add KernelAbstractions.get_backend for VectorOfArray

Author: ChrisRackauckas

ext/RecursiveArrayToolsKernelAbstractions.jl

@@ -0,0 +1,16 @@
+module RecursiveArrayToolsKernelAbstractions
+
+import RecursiveArrayTools: VectorOfArray
+import KernelAbstractions
+
+
+function KernelAbstractions.get_backend(x::VectorOfArray)
+    u = parent(x)
+    if length(u) == 0
+        error("VectorOfArray is empty, cannot determine backend.")
+    end
+    # Use the backend of the first element in the parent array
+    return KernelAbstractions.get_backend(u[1])
+end
+
+end

test/gpu/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 
 [compat]
 CUDA = "3.12, 4, 5"

test/gpu/vectorofarray_gpu.jl

@@ -1,15 +1,19 @@
-using RecursiveArrayTools, CUDA, Test, Zygote, Adapt
+using RecursiveArrayTools, CUDA, Test, Zygote, Adapt, KernelAbstractions
 CUDA.allowscalar(false)
 
 # Test indexing with colon
 x = zeros(5)
 y = VectorOfArray([x, x, x])
 y[:, :]
 
+KernelAbstractions.get_backend(y) isa KernelAbstractions.CPU
+
 x = CUDA.zeros(5)
 y = VectorOfArray([x, x, x])
 y[:, :]
 
+KernelAbstractions.get_backend(y) isa CUDA.CUDABackend
+
 # Test indexing with boolean masks and colon
 nx, ny, nt = 3, 4, 5
 x = CUDA.rand(nx, ny, nt)