Miscellaneous updates to get ConformalCubedSphereGrid simulations run…

…ning on the GPU (#1608)

* ConformalCubedSphereGrid and friends running on the GPU...

* Uncomment exchange bcs test

* Updates hydrostatic benchmarks for GPU

* Removes vestigial on_architecture

* Some cleanup and changes all xnode, ynode, znode to use instantiated types

* Goodbye xC yC zC

* Use znode rather than zF in nonlinear equation of state

* Adapt FunctionField grid

* Dont fill halo regions after broadcasting (for now)

* Updates test_field.jl

* Fixes geopotential height function for nonlinear equation of state

* Adds manual fill halo regions in broadcast test

* Allow scalar operations for broadcast of ==

* Imports node into continuous_forcing

* Use AbstractCPUArchitecture etc in zeros instead of CPU

* Use interior(field) in implicit free surface solver test

* Special set! for fields on the same architecture

* Dont test GPU cubed sphere in CPU tests

* Fill eta halos before computing free surface linear operation

* Fixes node for reduced locations

* Free surfaces are two-dimensional

* Update benchmark_hydrostatic_model.jl

* Update fill_halo_regions.jl

* Update cubed_sphere_halo_filling.jl

* Update cubed_sphere_set!.jl

* When StackOverflow doesnt help

* Decompose arguments to fill halo regions into faces

benchmark/benchmark_hydrostatic_model.jl

@@ -19,22 +19,27 @@ DataDeps.register(dd)
 
 # All grids have 6 * 32^2 = 6144 grid points.
 grids = Dict(
-    :RegularRectilinearGrid => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
-    :RegularLatitudeLongitudeGrid => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
-    :ConformalCubedSphereFaceGrid => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0)),
-    :ConformalCubedSphereGrid => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0))
+     (CPU, :RegularRectilinearGrid)       => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
+     (CPU, :RegularLatitudeLongitudeGrid) => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
+     (CPU, :ConformalCubedSphereFaceGrid) => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0)),
+     (CPU, :ConformalCubedSphereGrid)     => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0)),
+     (GPU, :RegularRectilinearGrid)       => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
+     (GPU, :RegularLatitudeLongitudeGrid) => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
+     # Uncomment when ConformalCubedSphereFaceGrids of any size can be built natively without loading from file:
+     # (GPU, :ConformalCubedSphereFaceGrid) => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0), architecture=GPU()),
+     (GPU, :ConformalCubedSphereGrid)     => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0), architecture=GPU()),
 )
 
 free_surfaces = Dict(
     :ExplicitFreeSurface => ExplicitFreeSurface(),
-    :ImplicitFreeSurface => ImplicitFreeSurface(maximum_iterations=1, tolerance=-Inf)  # Force it to take exactly 1 iteration.
+    :ImplicitFreeSurface => ImplicitFreeSurface(maximum_iterations=1, tolerance=-Inf) # Force it to take exactly 1 iteration.
 )
 
 function benchmark_hydrostatic_model(Arch, grid_type, free_surface_type)
 
     model = HydrostaticFreeSurfaceModel(
               architecture = Arch(),
-                      grid = grids[grid_type],
+                      grid = grids[(Arch, grid_type)],
         momentum_advection = VectorInvariant(),
               free_surface = free_surfaces[free_surface_type]
     )
@@ -55,13 +60,15 @@ Architectures = has_cuda() ? [CPU, GPU] : [CPU]
 grid_types = [
     :RegularRectilinearGrid,
     :RegularLatitudeLongitudeGrid,
-    :ConformalCubedSphereFaceGrid,
+    # Uncomment when ConformalCubedSphereFaceGrids of any size can be built natively without loading from file:
+    # :ConformalCubedSphereFaceGrid,
     :ConformalCubedSphereGrid
 ]
 
 free_surface_types = [
     :ExplicitFreeSurface,
-    :ImplicitFreeSurface
+    # ImplicitFreeSurface doesn't yet work on MultiRegionGrids like the ConformalCubedSphereGrid:
+    # :ImplicitFreeSurface
 ]
 
 # Run and summarize benchmarks

examples/geostrophic_adjustment.jl

@@ -57,13 +57,13 @@ g = model.free_surface.gravitational_acceleration
 
 η₀ = coriolis.f * U * L / g # geostrohpic free surface amplitude
 
-ηᵍ(x, y, z) = η₀ * Gaussian(x - x₀, L)
+ηᵍ(x) = η₀ * Gaussian(x - x₀, L)
 
 # We use an initial height field that's twice the geostrophic solution,
 # thus superimposing a geostrophic and ageostrophic component in the free
 # surface displacement field:
 
-ηⁱ(x, y, z) = 2 * ηᵍ(x, y, z)
+ηⁱ(x, y) = 2 * ηᵍ(x)
 
 # We set the initial condition to ``vᵍ`` and ``ηⁱ``,
 

src/Architectures.jl

@@ -1,12 +1,13 @@
 module Architectures
 
-export
-    AbstractArchitecture, AbstractCPUArchitecture, AbstractGPUArchitecture, CPU, GPU,
-    device, architecture, array_type, arch_array
+export AbstractArchitecture, AbstractCPUArchitecture, AbstractGPUArchitecture
+export CPU, GPU
+export device, architecture, array_type, arch_array
 
 using CUDA
-
 using KernelAbstractions
+using Adapt
+using OffsetArrays
 
 """
     AbstractArchitecture
@@ -60,4 +61,13 @@ arch_array(::AbstractCPUArchitecture, A::CuArray) = Array(A)
 arch_array(::AbstractGPUArchitecture, A::Array) = CuArray(A)
 arch_array(::AbstractGPUArchitecture, A::CuArray) = A
 
+const OffsetCPUArray = OffsetArray{FT, N, <:Array} where {FT, N}
+const OffsetGPUArray = OffsetArray{FT, N, <:CuArray} where {FT, N}
+
+Adapt.adapt_structure(::CPU, a::OffsetCPUArray) = a
+Adapt.adapt_structure(::GPU, a::OffsetGPUArray) = a
+
+Adapt.adapt_structure(::GPU, a::OffsetCPUArray) = OffsetArray(CuArray(a.parent), a.offsets...)
+Adapt.adapt_structure(::CPU, a::OffsetGPUArray) = OffsetArray(Array(a.parent), a.offsets...)
+
 end

src/BoundaryConditions/continuous_boundary_function.jl

@@ -76,14 +76,14 @@ The regularization of `bc.condition::ContinuousBoundaryFunction` requries
    of the boundary.
 """
 function regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoundaryFunction},
-                                       X, Y, Z, I, model_field_names) where C
+                                       LX, LY, LZ, I, model_field_names) where C
     boundary_func = bc.condition
 
-    indices, interps = index_and_interp_dependencies(X, Y, Z,
+    indices, interps = index_and_interp_dependencies(LX, LY, LZ,
                                                      boundary_func.field_dependencies,
                                                      model_field_names)
 
-    regularized_boundary_func = ContinuousBoundaryFunction{X, Y, Z, I}(boundary_func.func,
+    regularized_boundary_func = ContinuousBoundaryFunction{LX, LY, LZ, I}(boundary_func.func,
                                                                        boundary_func.parameters,
                                                                        boundary_func.field_dependencies,
                                                                        indices, interps)
@@ -95,19 +95,19 @@ end
 ##### Kernel functions
 #####
 
-@inline function (bc::ContinuousBoundaryFunction{Nothing, Y, Z, i})(j, k, grid, clock, model_fields) where {Y, Z, i}
+@inline function (bc::ContinuousBoundaryFunction{Nothing, LY, LZ, i})(j, k, grid, clock, model_fields) where {LY, LZ, i}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(ynode(Y, j, grid), znode(Z, k, grid), clock.time, args...)
+    return bc.func(ynode(LY(), j, grid), znode(LZ(), k, grid), clock.time, args...)
 end
 
-@inline function (bc::ContinuousBoundaryFunction{X, Nothing, Z, j})(i, k, grid, clock, model_fields) where {X, Z, j}
+@inline function (bc::ContinuousBoundaryFunction{LX, Nothing, LZ, j})(i, k, grid, clock, model_fields) where {LX, LZ, j}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(xnode(X, i, grid), znode(Z, k, grid), clock.time, args...)
+    return bc.func(xnode(LX(), i, grid), znode(LZ(), k, grid), clock.time, args...)
 end
 
-@inline function (bc::ContinuousBoundaryFunction{X, Y, Nothing, k})(i, j, grid, clock, model_fields) where {X, Y, k}
+@inline function (bc::ContinuousBoundaryFunction{LX, LY, Nothing, k})(i, j, grid, clock, model_fields) where {LX, LY, k}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(xnode(X, i, grid), ynode(Y, j, grid), clock.time, args...)
+    return bc.func(xnode(LX(), i, grid), ynode(LY(), j, grid), clock.time, args...)
 end
 
 # Don't re-convert ContinuousBoundaryFunctions passed to BoundaryCondition constructor
@@ -117,8 +117,8 @@ BoundaryCondition(TBC, condition::ContinuousBoundaryFunction) =
 Adapt.adapt_structure(to, bc::BoundaryCondition{C, <:ContinuousBoundaryFunction}) where C =
     BoundaryCondition(C, Adapt.adapt(to, bc.condition))
 
-Adapt.adapt_structure(to, bf::ContinuousBoundaryFunction{X, Y, Z, I}) where {X, Y, Z, I} =
-    ContinuousBoundaryFunction{X, Y, Z, I}(Adapt.adapt(to, bf.func),
+Adapt.adapt_structure(to, bf::ContinuousBoundaryFunction{LX, LY, LZ, I}) where {LX, LY, LZ, I} =
+    ContinuousBoundaryFunction{LX, LY, LZ, I}(Adapt.adapt(to, bf.func),
                                            Adapt.adapt(to, bf.parameters),
                                            nothing,
                                            Adapt.adapt(to, bf.field_dependencies_indices),

src/BoundaryConditions/fill_halo_regions.jl

@@ -21,6 +21,9 @@ function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)
     return nothing
 end
 
+# Some fields have `nothing` boundary conditions, such as `FunctionField` and `ZeroField`.
+fill_halo_regions!(c::OffsetArray, ::Nothing, args...; kwargs...) = nothing
+
 "Fill halo regions in x, y, and z for a given field's data."
 function fill_halo_regions!(c::OffsetArray, fieldbcs, arch, grid, args...; kwargs...)
 

src/BuoyancyModels/buoyancy_field.jl

@@ -1,9 +1,11 @@
 using Adapt
 using KernelAbstractions
-using Oceananigans.Fields: AbstractDataField, FieldStatus, validate_field_data, new_data, conditional_compute!
+
+using Oceananigans.Fields: AbstractDataField, FieldStatus, validate_field_data, conditional_compute!
 using Oceananigans.Fields: architecture, tracernames
 using Oceananigans.Architectures: device
-using Oceananigans.Utils: work_layout
+using Oceananigans.Utils: work_layout, new_data
+using Oceananigans.Grids: new_data
 
 import Oceananigans.Fields: compute!, compute_at!
 

src/BuoyancyModels/nonlinear_equation_of_state.jl

@@ -1,26 +1,26 @@
 using Oceananigans.Fields: AbstractField
-using Oceananigans.Grids: zF, zC
+using Oceananigans.Grids: znode
 using Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ
 
 """ Return the geopotential height at `i, j, k` at cell centers. """
 @inline function Zᵃᵃᶜ(i, j, k, grid::AbstractGrid{FT}) where FT
     if k < 1
-        return zC(1, grid) + (1 - k) * Δzᵃᵃᶠ(i, j, 1, grid)
+        return znode(Center(), 1, grid) + (1 - k) * Δzᵃᵃᶠ(i, j, 1, grid)
     elseif k > grid.Nz
-        return zC(grid.Nz, grid) - (k - grid.Nz) * Δzᵃᵃᶠ(i, j, grid.Nz, grid)
+        return znode(Center(), grid.Nz, grid) - (k - grid.Nz) * Δzᵃᵃᶠ(i, j, grid.Nz, grid)
     else
-        return zC(k, grid)
+        return znode(Center(), k, grid)
     end
 end
 
 """ Return the geopotential height at `i, j, k` at cell z-interfaces. """
 @inline function Zᵃᵃᶠ(i, j, k, grid::AbstractGrid{FT}) where FT
     if k < 1
-        return zF(1, grid) + (1 - k) * Δzᵃᵃᶜ(i, j, 1, grid)
+        return znode(Face(), 1, grid) + (1 - k) * Δzᵃᵃᶜ(i, j, 1, grid)
     elseif k > grid.Nz + 1
-        return zF(grid.Nz + 1, grid) - (k - grid.Nz + 1) * Δzᵃᵃᶜ(i, j, k, grid)
+        return znode(Face(), grid.Nz + 1, grid) - (k - grid.Nz + 1) * Δzᵃᵃᶜ(i, j, k, grid)
     else
-        return zF(k, grid)
+        return znode(Face(), k, grid)
     end
 end
 

src/CubedSpheres/CubedSpheres.jl

@@ -35,10 +35,15 @@ validate_vertical_velocity_boundary_conditions(w::AbstractCubedSphereField) =
 
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: apply_flux_bcs!
 
-apply_flux_bcs!(Gcⁿ::AbstractCubedSphereField, events, c::AbstractCubedSphereField, arch, barrier, clock, model_fields) = [
-    apply_flux_bcs!(get_face(Gcⁿ, face_index), events, get_face(c, face_index), arch, barrier, clock, model_fields)
-    for face_index in 1:length(Gcⁿ.data.faces)
-]
+function apply_flux_bcs!(Gcⁿ::AbstractCubedSphereField, events, c::AbstractCubedSphereField, arch, barrier, clock, model_fields)
+
+    for (face_index, Gcⁿ_face) in enumerate(faces(Gcⁿ))
+        apply_flux_bcs!(Gcⁿ_face, events, get_face(c, face_index), arch, barrier,
+                        clock, get_face(model_fields, face_index))
+    end
+
+    return nothing
+end
 
 #####
 ##### NaN checker for cubed sphere fields

src/CubedSpheres/conformal_cubed_sphere_grid.jl

@@ -156,11 +156,11 @@ function ConformalCubedSphereGrid(FT=Float64; face_size, z, radius=R_Earth)
     return ConformalCubedSphereGrid{FT, typeof(faces), typeof(face_connectivity)}(faces, face_connectivity)
 end
 
-function ConformalCubedSphereGrid(filepath::AbstractString, FT=Float64; Nz, z, radius = R_Earth, halo = (1, 1, 1))
+function ConformalCubedSphereGrid(filepath::AbstractString, FT=Float64; Nz, z, architecture = CPU(), radius = R_Earth, halo = (1, 1, 1))
     @warn "ConformalCubedSphereGrid is experimental: use with caution!"
 
     face_topo = (Connected, Connected, Bounded)
-    face_kwargs = (Nz=Nz, z=z, topology=face_topo, radius=radius, halo=halo)
+    face_kwargs = (Nz=Nz, z=z, topology=face_topo, radius=radius, halo=halo, architecture=architecture)
 
     faces = Tuple(ConformalCubedSphereFaceGrid(filepath, FT; face=n, face_kwargs...) for n in 1:6)
 

src/CubedSpheres/cubed_sphere_exchange_bcs.jl

@@ -1,3 +1,5 @@
+using Adapt
+
 using Oceananigans.BoundaryConditions
 using Oceananigans.BoundaryConditions: BCType
 
@@ -68,3 +70,6 @@ function inject_cubed_sphere_exchange_boundary_conditions(field_bcs, face_number
 
     return FieldBoundaryConditions(x_bcs, y_bcs, field_bcs.z)
 end
+
+Adapt.adapt_structure(to, ::CubedSphereExchangeInformation) = nothing
+Adapt.adapt_structure(to, ::CubedSphereExchangeBC) = nothing

src/CubedSpheres/cubed_sphere_faces.jl

@@ -6,7 +6,8 @@ using OffsetArrays: OffsetArray
 import Base: getindex, size, show, minimum, maximum
 import Statistics: mean
 
-import Oceananigans.Fields: AbstractField, AbstractDataField, AbstractReducedField, Field, new_data, minimum, maximum, mean, location
+import Oceananigans.Fields: AbstractField, AbstractDataField, AbstractReducedField, Field, minimum, maximum, mean, location
+import Oceananigans.Grids: new_data
 import Oceananigans.BoundaryConditions: FieldBoundaryConditions
 
 struct CubedSphereFaces{E, F}

src/CubedSpheres/cubed_sphere_halo_filling.jl

@@ -1,5 +1,5 @@
 import Oceananigans.BoundaryConditions:
-    fill_halo_regions!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!
+    fill_halo_regions!, fill_top_halo!, fill_bottom_halo!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!
 
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: fill_horizontal_velocity_halos!
 
@@ -11,9 +11,9 @@ fill_north_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing
 
 function fill_halo_regions!(field::AbstractCubedSphereField, arch, args...; kwargs...)
 
-    for face_field in faces(field)
+    for (i, face_field) in enumerate(faces(field))
         # Fill the top and bottom halos the usual way.
-        fill_halo_regions!(face_field, arch, args...; kwargs...)
+        fill_halo_regions!(face_field, arch, get_face(args, i)...; kwargs...)
 
         # Deal with halo exchanges.
         fill_west_halo!(face_field, field)

src/CubedSpheres/cubed_sphere_set!.jl

@@ -4,6 +4,7 @@ using Oceananigans.Fields: AbstractField, ReducedField
 import Oceananigans.Fields: set!
 
 const CubedSphereCPUField = CubedSphereField{X, Y, Z, <:AbstractCPUArchitecture} where {X, Y, Z}
+const CubedSphereGPUField = CubedSphereField{X, Y, Z, <:AbstractGPUArchitecture} where {X, Y, Z}
 
 function set!(u::CubedSphereCPUField , v::CubedSphereCPUField)
     for (u_face, v_face) in zip(faces(u), faces(v))
@@ -14,31 +15,7 @@ end
 
 set!(field::CubedSphereCPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
 set!(field::CubedSphereCPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]
+set!(field::CubedSphereGPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
+set!(field::CubedSphereGPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]
 
 set_face_field!(field, a) = set!(field, a)
-
-function set_face_field!(field, f::Function)
-    LX, LY, LZ = location(field)
-    grid = field.grid
-
-    for i in 1:grid.Nx, j in 1:grid.Ny, k in 1:grid.Nz
-        λ = λnode(LX(), LY(), LZ(), i, j, k, grid)
-        φ = φnode(LX(), LY(), LZ(), i, j, k, grid)
-        z = znode(LX(), LY(), LZ(), i, j, k, grid)
-        field[i, j, k] = f(λ, φ, z)
-    end
-
-    return nothing
-end
-
-function set_face_field!(field::ReducedField{LX, LY, Nothing}, f::Function) where {LX, LY}
-    grid = field.grid
-
-    for i in 1:grid.Nx, j in 1:grid.Ny
-        λ = λnode(LX(), LY(), nothing, i, j, 1, grid)
-        φ = φnode(LX(), LY(), nothing, i, j, 1, grid)
-        field[i, j, 1] = f(λ, φ)
-    end
-
-    return nothing
-end

src/Fields/Fields.jl

@@ -15,13 +15,6 @@ using Oceananigans.Architectures
 using Oceananigans.Grids
 using Oceananigans.BoundaryConditions
 
-import Base: zeros
-
-zeros(FT, ::CPU, Nx, Ny, Nz) = zeros(FT, Nx, Ny, Nz)
-zeros(FT, ::GPU, Nx, Ny, Nz) = zeros(FT, Nx, Ny, Nz) |> CuArray
-zeros(arch, grid, Nx, Ny, Nz) = zeros(eltype(grid), arch, Nx, Ny, Nz)
-
-include("new_data.jl")
 include("abstract_field.jl")
 include("field.jl")
 include("zero_field.jl")

src/Fields/abstract_field.jl

@@ -209,9 +209,9 @@ end
 ##### Coordinates of fields
 #####
 
-@propagate_inbounds xnode(i, ψ::AbstractField{X, Y, Z}) where {X, Y, Z} = xnode(X, i, ψ.grid)
-@propagate_inbounds ynode(j, ψ::AbstractField{X, Y, Z}) where {X, Y, Z} = ynode(Y, j, ψ.grid)
-@propagate_inbounds znode(k, ψ::AbstractField{X, Y, Z}) where {X, Y, Z} = znode(Z, k, ψ.grid)
+@propagate_inbounds xnode(i, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = xnode(LX(), i, ψ.grid)
+@propagate_inbounds ynode(j, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = ynode(LY(), j, ψ.grid)
+@propagate_inbounds znode(k, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = znode(LZ(), k, ψ.grid)
 
 @propagate_inbounds Base.lastindex(f::AbstractDataField) = lastindex(f.data)
 @propagate_inbounds Base.lastindex(f::AbstractDataField, dim) = lastindex(f.data, dim)

src/Fields/broadcasting_abstract_fields.jl

@@ -82,7 +82,5 @@ broadcasted_to_abstract_operation(loc, grid, a) = a
 
     wait(device(arch), event)
 
-    fill_halo_regions!(dest, arch)
-
     return nothing
 end