support SubArray on fill

CliMA · Jul 20, 2023 · dff77b0 · dff77b0
1 parent 6af8766
commit dff77b0
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/DataLayouts/cuda.jl b/src/DataLayouts/cuda.jl
@@ -97,7 +97,11 @@ end
 function Base.fill!(
     dest::IJFH{S, Nij, A},
     val,
-) where {S, Nij, A <: CUDA.CuArray}
+) where {
+    S,
+    Nij,
+    A <: Union{CUDA.CuArray, SubArray{<:Any, <:Any, <:CUDA.CuArray}},
+}
     _, _, _, _, Nh = size(dest)
     if Nh > 0
         CUDA.@cuda threads = (Nij, Nij) blocks = (Nh, 1) knl_fill!(dest, val)