Added unrolled implementation of recursivefill! which works on GPUs and avoids recomputing global indices for each setindex!

Author: anicusan

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)