diff --git a/ext/RecursiveArrayToolsKernelAbstractions.jl b/ext/RecursiveArrayToolsKernelAbstractions.jl new file mode 100644 index 00000000..94f79884 --- /dev/null +++ b/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 diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index 18ce094b..dca24d21 100644 --- a/test/gpu/Project.toml +++ b/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" diff --git a/test/gpu/vectorofarray_gpu.jl b/test/gpu/vectorofarray_gpu.jl index 7422a6c1..68ea44d7 100644 --- a/test/gpu/vectorofarray_gpu.jl +++ b/test/gpu/vectorofarray_gpu.jl @@ -1,4 +1,4 @@ -using RecursiveArrayTools, CUDA, Test, Zygote, Adapt +using RecursiveArrayTools, CUDA, Test, Zygote, Adapt, KernelAbstractions CUDA.allowscalar(false) # Test indexing with colon @@ -6,10 +6,14 @@ 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)