Define CuArray on FieldArrays

CliMA · Aug 31, 2024 · 6748636 · 6748636
1 parent 5df8fbb
commit 6748636
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/ext/cuda/data_layouts.jl b/ext/cuda/data_layouts.jl
@@ -77,3 +77,12 @@ function Adapt.adapt_structure(
         Adapt.adapt(to, bc.axes),
     )
 end
+
+import ClimaCore.DataLayouts as DL
+import CUDA
+function CUDA.CuArray(fa::DL.FieldArray{FD}) where {FD}
+    arrays = ntuple(Val(DL.ncomponents(fa))) do f
+        CUDA.CuArray(fa.arrays[f])
+    end
+    return DL.FieldArray{FD}(arrays)
+end