Merge #1111

1111: cuda support for dss2 (serial) r=sriharshakandala a=sriharshakandala Co-authored-by: sriharshakandala <sriharsha.kvs@gmail.com>
CliMA · Jan 31, 2023 · ba71538 · ba71538
2 parents 95c84fc + de36bcb
commit ba71538
Show file tree

Hide file tree

Showing 12 changed files with 737 additions and 145 deletions.
diff --git a/examples/sphere/shallow_water_cuda.jl b/examples/sphere/shallow_water_cuda.jl
@@ -479,6 +479,10 @@ function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
     h_s = surface_topography(space, test)
     Y = set_initial_condition(space, test)
 
+    ghost_buffer = Spaces.create_dss_buffer(Y)
+    Spaces.weighted_dss_start2!(Y, ghost_buffer)
+    Spaces.weighted_dss_internal2!(Y, ghost_buffer)
+    Spaces.weighted_dss_ghost2!(Y, ghost_buffer)
     return nothing
 end
 

diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl
@@ -596,6 +596,9 @@ function IJF{S, Nij}(array::AbstractArray{T, 3}) where {S, Nij, T}
     IJF{S, Nij, typeof(array)}(array)
 end
 
+rebuild(data::IJF{S, Nij}, array::A) where {S, Nij, A <: AbstractArray} =
+    IJF{S, Nij}(array)
+
 function replace_basetype(data::IJF{S, Nij}, ::Type{T}) where {S, Nij, T}
     array = parent(data)
     S′ = replace_basetype(eltype(array), T, S)
@@ -1058,6 +1061,9 @@ function VIFH{S, Ni}(array::AbstractArray{T, 4}) where {S, Ni, T}
     VIFH{S, Ni, typeof(array)}(array)
 end
 
+rebuild(data::VIFH{S, Ni}, array::A) where {S, Ni, A <: AbstractArray} =
+    VIFH{S, Ni}(array)
+
 function replace_basetype(data::VIFH{S, Ni}, ::Type{T}) where {S, Ni, T}
     array = parent(data)
     S′ = replace_basetype(eltype(array), T, S)

diff --git a/src/Spaces/Spaces.jl b/src/Spaces/Spaces.jl
@@ -15,6 +15,8 @@
 module Spaces
 
 using ClimaComms
+using Adapt
+using CUDA
 
 import ..slab, ..column, ..level
 import ..Utilities: PlusHalf