Make ArrayPartitions of GPU arrays work in DifferentialEquations solvers

src/array_partition.jl

@@ -209,6 +209,13 @@ function Base.copyto!(A::ArrayPartition, src::ArrayPartition)
     A
 end
 
+function recursivefill!(b::ArrayPartition, a::T2) where {T2 <: Union{Number, Bool}}
+    unrolled_foreach!(b.x) do x
+        fill!(x, a)
+    end
+end
+
+
 ## indexing
 
 # Interface for the linear indexing. This is just a view of the underlying nested structure

src/utils.jl

@@ -1,3 +1,7 @@
+unrolled_foreach!(f, t::Tuple) = (f(t[1]); unrolled_foreach!(f, Base.tail(t)))
+unrolled_foreach!(f, ::Tuple{}) = nothing
+
+
 """
 ```julia
 recursivecopy(a::Union{AbstractArray{T, N}, AbstractVectorOfArray{T, N}})
@@ -127,6 +131,7 @@ function recursivefill!(bs::AbstractVectorOfArray{T, N},
     end
 end
 
+
 for type in [AbstractArray, AbstractVectorOfArray]
     @eval function recursivefill!(b::$type{T, N}, a::T2) where {T <: Enum, T2 <: Enum, N}
         fill!(b, a)

test/gpu/arraypartition_gpu.jl

@@ -0,0 +1,16 @@
+using RecursiveArrayTools, CUDA, Test
+CUDA.allowscalar(false)
+
+
+# Test indexing with colon
+a = (CUDA.zeros(5), CUDA.zeros(5))
+pA = ArrayPartition(a)
+pA[:, :]
+
+# Indexing with boolean masks does not work yet
+mask = pA .> 0
+# pA[mask]
+
+# Test recursive filling is done using GPU kernels and not scalar indexing
+RecursiveArrayTools.recursivefill!(pA, true)
+@test all(pA .== true)

test/runtests.jl

@@ -54,5 +54,6 @@ end
     if GROUP == "GPU"
         activate_gpu_env()
         @time @safetestset "VectorOfArray GPU" include("gpu/vectorofarray_gpu.jl")
+        @time @safetestset "ArrayPartition GPU" include("gpu/arraypartition_gpu.jl")
     end
 end