Add support for copyto! and fill! for DataF

(Used for Point spaces)
CliMA · Nov 15, 2023 · e5a0157 · e5a0157
1 parent c76e128
commit e5a0157
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/DataLayouts/cuda.jl b/src/DataLayouts/cuda.jl
@@ -144,3 +144,15 @@ function Base.fill!(dest::VF{S, A}, val) where {S, A <: CUDA.CuArray}
     end
     return dest
 end
+
+function Base.copyto!(
+    dest::DataF{S},
+    bc::Union{DataF{S, A}, Base.Broadcast.Broadcasted{DataFStyle{A}}},
+) where {S, A <: CUDA.CuArray}
+    CUDA.@cuda threads = (1, 1) blocks = (1, 1) knl_copyto!(dest, bc)
+    return dest
+end
+function Base.fill!(dest::DataF{S, A}, val) where {S, A <: CUDA.CuArray}
+    CUDA.@cuda threads = (1, 1) blocks = (1, 1) knl_fill!(dest, val)
+    return dest
+end
diff --git a/test/DataLayouts/cuda.jl b/test/DataLayouts/cuda.jl
@@ -72,4 +72,8 @@ end
     @test Array(parent(data)) == FT[
         f == 1 ? 1 : 2 for v in 1:Nv, i in 1:4, j in 1:4, f in 1:2, h in 1:3
     ]
+
+    data = DataF{S}(CuArray{FT})
+    data .= Complex(1.0, 2.0)
+    @test Array(parent(data)) == FT[f == 1 ? 1 : 2 for f in 1:2]
 end