diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 7fdc1b51f8..84f3382639 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -773,17 +773,6 @@ steps: agents: slurm_gpus: 1 - - label: "Unit: FD operator (shmem)" - key: unit_fd_operator_shmem - retry: *retry_policy - command: - - "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/finitedifference/unit_fd_ops_shared_memory.jl" - - "julia --color=yes --project=.buildkite test/Operators/finitedifference/benchmark_fd_ops_shared_memory.jl" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_gpus: 1 - - label: "Unit: gpu columnwise" key: unit_gpu_columnwise retry: *retry_policy diff --git a/docs/make.jl b/docs/make.jl index ea10d7ed2e..b4742b34e3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -106,7 +106,6 @@ withenv("GKSwstype" => "nul") do "Developer docs" => [ "Developer Guides" => "dev_guides.md", "Performance tips" => "performance_tips.md", - "Shared memory design" => "shmem_design.md", ], "Tutorials" => [ joinpath("tutorials", tutorial * ".md") for diff --git a/docs/src/shmem_design.md b/docs/src/shmem_design.md deleted file mode 100644 index 42250ddd99..0000000000 --- a/docs/src/shmem_design.md +++ /dev/null @@ -1,74 +0,0 @@ -# Shared memory design - -ClimaCore stencil operators support staggered (or collocated) finite difference -and interpolation operations. For example, the `DivergenceF2C` operator takes -an argument that lives on the cell faces and the resulting divergence -calculation lives on the cell centers. Such operations are effectively -matrix-vector multiplication and are often a significant portion of the runtime -cost for users. - -Here, we outline an optimization, shared memory (or, "shmem" for short), that we -use to improve the performance of these operations. - -## Motivation - -A naive and simplified implementation of this operation looks like `div[i] = (f -[i+1] - f[i]) / dz[i]`. Such a calculation on the gpu (or cpu) requires `f[i]` -be read from global memory to compute the result of `div[i]` and `div[i-1]`. Not -to mention, if `f` is a `Broadcasted` object (`Broadcasted` objects behave like -arrays, and support `f[i]` behavior), then `f[i]` may require several reads and -or computations. - -Reading data from global memory is often the main bottleneck for -bandwidth-limited cuda kernels. As such, we use shmem to reduce the number of global memory reads (and compute) in our kernels. - -## High-level design - -The high-level view of the design is: - - - The `bc::StencilBroadcasted` type has a `work` field, which is used to store - shmem for the `bc.op` operator. The element type of the `work` - (or parts of `work` if there are multiple parts) is the type returned by the - `bc.op`'s `Operator.return_eltype`. - - Recursively reconstruct the broadcasted object, allocating shmem for - each `StencilBroadcasted` along the way that supports shmem - (different operators require different arguments, and therefore different - types and amounts of shmem). - - Recursively fill the shmem for all `StencilBroadcasted`. This is done - by reading the argument data from `getidx`. See the section discussion below for more details. - - The destination field is filled with the result of `getidx` (as it is without - shmem), except that we overload `getidx` (for supported `StencilBroadcasted` - types) to retrieve the result of `getidx` via `fd_operator_evaluate`, which - retrieves the result from the shmem, instead of global memory. - -### Populating shared memory, and memory access safety - -We use tail-recursion when filling shared memory of the broadcast expressions. -That is, we visit leaves of the broadcast expression, then work our way up. -It's important to note that the `StencilBroadcasted` and `Broadcasted` can be -interleaved. - -Let's take `DivergenceF2C()(f*GradientC2F()(a*b)))` as an example (depicted in -the image below). - -Recursion must go through the entire expression in order to ensure that we've -reached all of the leaves of the `StencilBroadcasted` objects (otherwise, we -could introduce race conditions with memory access). The leaves of the -`StencilBroadcasted` will call `getidx`, below which there are (by definition) -no more `StencilBroadcasted`, and those `getidx` calls will read from global -memory. All subsequent reads will be from shmem(as they will be caught by the -`getidx(parent_space, bc::StencilBroadcasted -{CUDAWithShmemColumnStencilStyle}, idx, hidx)` defined in the -`ClimaCoreCUDAExt` module). - -In the diagram below, we traverse and fill the yellow highlighted sections -(bottom first and top last). The algorithmic impact of using shared memory is -that the duplicate global memory reads (highlighted in red circles) become one -global memory read (performed in `fd_operator_fill_shmem!`). - -Finally, its important to note that threads must by syncrhonized after each node -in the tree is filled, to avoid race conditions for subsequent `getidx -(parent_space, bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, idx, -hidx)` calls (which are retrieved via shmem). - -![](shmem_diagram_example.png) diff --git a/docs/src/shmem_diagram_example.png b/docs/src/shmem_diagram_example.png deleted file mode 100644 index b079cbd2f4..0000000000 Binary files a/docs/src/shmem_diagram_example.png and /dev/null differ diff --git a/ext/ClimaCoreCUDAExt.jl b/ext/ClimaCoreCUDAExt.jl index 89583eabd4..807880c6d6 100644 --- a/ext/ClimaCoreCUDAExt.jl +++ b/ext/ClimaCoreCUDAExt.jl @@ -30,8 +30,6 @@ include(joinpath("cuda", "operators_integral.jl")) include(joinpath("cuda", "remapping_interpolate_array.jl")) include(joinpath("cuda", "limiters.jl")) include(joinpath("cuda", "operators_sem_shmem.jl")) -include(joinpath("cuda", "operators_fd_shmem_common.jl")) -include(joinpath("cuda", "operators_fd_shmem.jl")) include(joinpath("cuda", "operators_columnwise.jl")) include(joinpath("cuda", "matrix_fields_single_field_solve.jl")) include(joinpath("cuda", "matrix_fields_multiple_field_solve.jl")) diff --git a/ext/cuda/operators_fd_shmem.jl b/ext/cuda/operators_fd_shmem.jl deleted file mode 100644 index 7fb8f3f317..0000000000 --- a/ext/cuda/operators_fd_shmem.jl +++ /dev/null @@ -1,311 +0,0 @@ -import ClimaCore: DataLayouts, Spaces, Geometry, DataLayouts -import CUDA -import ClimaCore.Operators: return_eltype, get_local_geometry -import ClimaCore.Geometry: ⊗ - -Base.@propagate_inbounds function fd_operator_shmem( - space, - shmem_params, - op::Operators.DivergenceF2C, - args..., -) - # allocate temp output - RT = return_eltype(op, args...) - Ju³ = CUDA.CuStaticSharedArray(RT, interior_size(shmem_params)) - lJu³ = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) - rJu³ = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) - return (Ju³, lJu³, rJu³) -end - -Base.@propagate_inbounds function fd_operator_fill_shmem!( - op::Operators.DivergenceF2C, - (Ju³, lJu³, rJu³), - bc_bds, - arg_space, - space, - idx::Utilities.PlusHalf, - hidx, - arg, -) - @inbounds begin - vt = threadIdx().x - lg = Geometry.LocalGeometry(space, idx, hidx) - if !on_boundary(idx, space, op) - u³ = Operators.getidx(space, arg, idx, hidx) - Ju³[vt] = Geometry.Jcontravariant3(u³, lg) - elseif on_left_boundary(idx, space, op) - bloc = Operators.left_boundary_window(space) - bc = Operators.get_boundary(op, bloc) - ub = Operators.getidx(space, bc.val, nothing, hidx) - bJu³ = on_left_boundary(idx, space) ? lJu³ : rJu³ - if bc isa Operators.SetValue - bJu³[1] = Geometry.Jcontravariant3(ub, lg) - elseif bc isa Operators.SetDivergence - bJu³[1] = ub - elseif bc isa Operators.Extrapolate # no shmem needed - end - elseif on_right_boundary(idx, space, op) - bloc = Operators.right_boundary_window(space) - bc = Operators.get_boundary(op, bloc) - ub = Operators.getidx(space, bc.val, nothing, hidx) - bJu³ = on_left_boundary(idx, space) ? lJu³ : rJu³ - if bc isa Operators.SetValue - bJu³[1] = Geometry.Jcontravariant3(ub, lg) - elseif bc isa Operators.SetDivergence - bJu³[1] = ub - elseif bc isa Operators.Extrapolate # no shmem needed - end - end - end - return nothing -end - -Base.@propagate_inbounds function fd_operator_evaluate( - op::Operators.DivergenceF2C, - (Ju³, lJu³, rJu³), - space, - idx::Integer, - hidx, - arg, -) - @inbounds begin - vt = threadIdx().x - lg = Geometry.LocalGeometry(space, idx, hidx) - if !on_boundary(idx, space, op) - Ju³₋ = Ju³[vt] # corresponds to idx - half - Ju³₊ = Ju³[vt + 1] # corresponds to idx + half - return (Ju³₊ - Ju³₋) * lg.invJ - else - bloc = - on_left_boundary(idx, space, op) ? - Operators.left_boundary_window(space) : - Operators.right_boundary_window(space) - bc = Operators.get_boundary(op, bloc) - @assert bc isa Operators.SetValue || bc isa Operators.SetDivergence - if on_left_boundary(idx, space) - if bc isa Operators.SetValue - Ju³₋ = lJu³[1] # corresponds to idx - half - Ju³₊ = Ju³[vt + 1] # corresponds to idx + half - return (Ju³₊ - Ju³₋) * lg.invJ - else - # @assert bc isa Operators.SetDivergence - return lJu³[1] - end - else - @assert on_right_boundary(idx, space) - if bc isa Operators.SetValue - Ju³₋ = Ju³[vt] # corresponds to idx - half - Ju³₊ = rJu³[1] # corresponds to idx + half - return (Ju³₊ - Ju³₋) * lg.invJ - else - @assert bc isa Operators.SetDivergence - return rJu³[1] - end - end - end - end -end - -Base.@propagate_inbounds function fd_operator_shmem( - space, - shmem_params, - op::Operators.GradientC2F, - args..., -) - # allocate temp output - RT = return_eltype(op, args...) - u = CUDA.CuStaticSharedArray(RT, interior_size(shmem_params)) # cell centers - lb = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) # left boundary - rb = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) # right boundary - return (u, lb, rb) -end - -Base.@propagate_inbounds function fd_operator_fill_shmem!( - op::Operators.GradientC2F, - (u, lb, rb), - bc_bds, - arg_space, - space, - idx::Integer, - hidx, - arg, -) - @inbounds begin - is_out_of_bounds(idx, space) && return nothing - vt = threadIdx().x - cov3 = Geometry.Covariant3Vector(1) - if in_domain(idx, arg_space) - u[vt] = cov3 ⊗ Operators.getidx(space, arg, idx, hidx) - end - if on_any_boundary(idx, space, op) - lloc = Operators.left_boundary_window(space) - rloc = Operators.right_boundary_window(space) - bloc = on_left_boundary(idx, space, op) ? lloc : rloc - @assert bloc isa typeof(lloc) && on_left_boundary(idx, space, op) || - bloc isa typeof(rloc) && on_right_boundary(idx, space, op) - bc = Operators.get_boundary(op, bloc) - @assert bc isa Operators.SetValue || bc isa Operators.SetGradient - ub = Operators.getidx(space, bc.val, nothing, hidx) - bu = on_left_boundary(idx, space) ? lb : rb - if bc isa Operators.SetValue - bu[1] = cov3 ⊗ ub - elseif bc isa Operators.SetGradient - lg = Geometry.LocalGeometry(space, idx, hidx) - bu[1] = Geometry.project(Geometry.Covariant3Axis(), ub, lg) - elseif bc isa Operators.Extrapolate # no shmem needed - end - end - end - return nothing -end - -Base.@propagate_inbounds function fd_operator_evaluate( - op::Operators.GradientC2F, - (u, lb, rb), - space, - idx::PlusHalf, - hidx, - args..., -) - @inbounds begin - vt = threadIdx().x - lg = Geometry.LocalGeometry(space, idx, hidx) - if !on_boundary(idx, space, op) - u₋ = u[vt - 1] # corresponds to idx - half - u₊ = u[vt] # corresponds to idx + half - return u₊ - u₋ - else - bloc = - on_left_boundary(idx, space, op) ? - Operators.left_boundary_window(space) : - Operators.right_boundary_window(space) - bc = Operators.get_boundary(op, bloc) - @assert bc isa Operators.SetValue - if on_left_boundary(idx, space) - if bc isa Operators.SetValue - u₋ = 2 * lb[1] # corresponds to idx - half - u₊ = 2 * u[vt] # corresponds to idx + half - return u₊ - u₋ - end - else - @assert on_right_boundary(idx, space) - if bc isa Operators.SetValue - u₋ = 2 * u[vt - 1] # corresponds to idx - half - u₊ = 2 * rb[1] # corresponds to idx + half - return u₊ - u₋ - end - end - end - end -end - -Base.@propagate_inbounds function fd_operator_shmem( - space, - shmem_params, - op::Operators.InterpolateC2F, - args..., -) - # allocate temp output - RT = return_eltype(op, args...) - u = CUDA.CuStaticSharedArray(RT, interior_size(shmem_params)) # cell centers - lb = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) # left boundary - rb = CUDA.CuStaticSharedArray(RT, boundary_size(shmem_params)) # right boundary - return (u, lb, rb) -end - -Base.@propagate_inbounds function fd_operator_fill_shmem!( - op::Operators.InterpolateC2F, - (u, lb, rb), - bc_bds, - arg_space, - space, - idx::Integer, - hidx, - arg, -) - @inbounds begin - is_out_of_bounds(idx, space) && return nothing - ᶜidx = get_cent_idx(idx) - if in_domain(idx, arg_space) - u[idx] = Operators.getidx(space, arg, idx, hidx) - else - lloc = Operators.left_boundary_window(space) - rloc = Operators.right_boundary_window(space) - bloc = on_left_boundary(idx, space, op) ? lloc : rloc - @assert bloc isa typeof(lloc) && on_left_boundary(idx, space, op) || - bloc isa typeof(rloc) && on_right_boundary(idx, space, op) - bc = Operators.get_boundary(op, bloc) - @assert bc isa Operators.SetValue || - bc isa Operators.SetGradient || - bc isa Operators.Extrapolate || - bc isa Operators.NullBoundaryCondition - if bc isa Operators.NullBoundaryCondition || - bc isa Operators.Extrapolate - u[idx] = Operators.getidx(space, arg, idx, hidx) - return nothing - end - bu = on_left_boundary(idx, space) ? lb : rb - ub = Operators.getidx(space, bc.val, nothing, hidx) - if bc isa Operators.SetValue - bu[1] = ub - elseif bc isa Operators.SetGradient - lg = Geometry.LocalGeometry(space, idx, hidx) - bu[1] = Geometry.covariant3(ub, lg) - end - end - end - return nothing -end - -Base.@propagate_inbounds function fd_operator_evaluate( - op::Operators.InterpolateC2F, - (u, lb, rb), - space, - idx::PlusHalf, - hidx, - args..., -) - @inbounds begin - vt = threadIdx().x - lg = Geometry.LocalGeometry(space, idx, hidx) - ᶜidx = get_cent_idx(idx) - if !on_boundary(idx, space, op) - u₋ = u[ᶜidx - 1] # corresponds to idx - half - u₊ = u[ᶜidx] # corresponds to idx + half - return (u₊ + u₋) / 2 - else - bloc = - on_left_boundary(idx, space, op) ? - Operators.left_boundary_window(space) : - Operators.right_boundary_window(space) - bc = Operators.get_boundary(op, bloc) - @assert bc isa Operators.SetValue || - bc isa Operators.SetGradient || - bc isa Operators.Extrapolate - if on_left_boundary(idx, space) - if bc isa Operators.SetValue - return lb[1] - elseif bc isa Operators.SetGradient - u₋ = lb[1] # corresponds to idx - half - u₊ = u[ᶜidx] # corresponds to idx + half - return u₊ - u₋ / 2 - else - @assert bc isa Operators.Extrapolate - return u[ᶜidx] - end - else - @assert on_right_boundary(idx, space) - if bc isa Operators.SetValue - return rb[1] - elseif bc isa Operators.SetGradient - u₋ = u[ᶜidx - 1] # corresponds to idx - half - u₊ = rb[1] # corresponds to idx + half - return u₋ + u₊ / 2 - else - @assert bc isa Operators.Extrapolate - return u[ᶜidx - 1] - end - end - end - end -end diff --git a/ext/cuda/operators_fd_shmem_common.jl b/ext/cuda/operators_fd_shmem_common.jl deleted file mode 100644 index cb978a084e..0000000000 --- a/ext/cuda/operators_fd_shmem_common.jl +++ /dev/null @@ -1,400 +0,0 @@ -import ClimaCore: DataLayouts, Spaces, Geometry, DataLayouts -import CUDA -import ClimaCore.Operators: return_eltype, get_local_geometry -import ClimaCore.Operators: getidx -import ClimaCore.Utilities: PlusHalf -import ClimaCore.Utilities - -##### -##### Boundary helpers -##### - -@inline on_boundary(idx, space, op) = - on_left_boundary(idx, space, op) || on_right_boundary(idx, space, op) - -@inline on_left_boundary(idx, space, op) = on_left_boundary(idx, space) -@inline on_right_boundary(idx, space, op) = on_right_boundary(idx, space) - -@inline on_boundary(idx::PlusHalf, space) = - idx == Operators.left_face_boundary_idx(space) || - idx == Operators.right_face_boundary_idx(space) -@inline on_boundary(idx::Integer, space) = - idx == Operators.left_center_boundary_idx(space) || - idx == Operators.right_center_boundary_idx(space) - -@inline on_left_boundary(idx::PlusHalf, space) = - idx == Operators.left_face_boundary_idx(space) -@inline on_left_boundary(idx::Integer, space) = - idx == Operators.left_center_boundary_idx(space) - -@inline on_right_boundary(idx::PlusHalf, space) = - idx == Operators.right_face_boundary_idx(space) -@inline on_right_boundary(idx::Integer, space) = - idx == Operators.right_center_boundary_idx(space) - -@inline on_any_boundary(idx, space, op) = - on_left_boundary(idx, space) || on_right_boundary(idx, space) - -@inline function is_out_of_bounds(idx::Integer, space) - ᶜspace = Spaces.center_space(space) - return idx == Spaces.nlevels(ᶜspace) + 1 -end - -##### -##### range window helpers (faces) -##### - -@inline function in_interior_range(idx::PlusHalf, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - get_face_idx(bc_lw) ≤ idx ≤ get_face_idx(bc_rw) + 1 -end -@inline function in_left_boundary_window_range(idx::PlusHalf, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - return get_face_idx(bc_li) - 1 ≤ idx ≤ get_face_idx(bc_lw + half) -end -@inline function in_right_boundary_window_range(idx::PlusHalf, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - return get_face_idx(bc_rw) ≤ idx ≤ get_face_idx(bc_ri) + 1 -end - -##### -##### range window helpers (centers) -##### - -@inline function in_interior_range(idx::Integer, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - get_cent_idx(bc_lw) ≤ idx ≤ get_cent_idx(bc_rw) -end -@inline function in_left_boundary_window_range(idx::Integer, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - return get_cent_idx(bc_li) ≤ idx < get_cent_idx(bc_lw) -end -@inline function in_right_boundary_window_range(idx::Integer, bc_bds) - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - return get_cent_idx(bc_rw) < idx ≤ get_cent_idx(bc_ri) -end - -##### -##### window helpers (faces) -##### - -@inline function in_left_boundary_window( - idx::PlusHalf, - space, - bc_bds, - op, - args..., -) - Operators.should_call_left_boundary(idx, space, op, args...) || - in_left_boundary_window_range(idx, bc_bds) -end - -@inline function in_right_boundary_window( - idx::PlusHalf, - space, - bc_bds, - op, - args..., -) - Operators.should_call_right_boundary(idx, space, op, args...) || - in_right_boundary_window_range(idx, bc_bds) -end - -@inline function in_interior(idx::PlusHalf, space, bc_bds, op, args...) - # TODO: simplify this function / logic / arithmetic - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - IP = Topologies.isperiodic(Spaces.vertical_topology(space)) - crb = in_right_boundary_window_range(idx, bc_bds) - clb = in_left_boundary_window_range(idx, bc_bds) - return IP || in_interior_range(idx, bc_bds) && !(crb || clb) -end - -##### -##### window helpers (centers) -##### - -@inline function in_interior(idx::Integer, space, bc_bds, op, args...) - # TODO: simplify this function / logic / arithmetic - # TODO: use the (commented) range methods instead, as it would be much simpler. - (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds - IP = Topologies.isperiodic(Spaces.vertical_topology(space)) - # crb = in_right_boundary_window_range(idx, bc_bds) - crb = in_right_boundary_window(idx, space, bc_bds, op, args...) - # clb = in_left_boundary_window_range(idx, bc_bds) - clb = in_left_boundary_window(idx, space, bc_bds, op, args...) - return IP || in_interior_range(idx, bc_bds) && !(crb || clb) -end - -@inline function in_domain(idx::Integer, space) - ᶜspace = Spaces.center_space(space) - return 1 ≤ idx ≤ Spaces.nlevels(ᶜspace) -end - -@inline function in_left_boundary_window( - idx::Integer, - space, - bc_bds, - op, - args..., -) - Operators.should_call_left_boundary(idx, space, op, args...) || - in_left_boundary_window_range(idx, bc_bds) -end - -@inline function in_right_boundary_window( - idx::Integer, - space, - bc_bds, - op, - args..., -) - ᶜspace = Spaces.center_space(space) - idx > Spaces.nlevels(ᶜspace) && return false # short-circuit if - Operators.should_call_right_boundary(idx, space, op, args...) || - in_right_boundary_window_range(idx, bc_bds) -end - -##### -##### -##### - -Base.@propagate_inbounds function getidx( - parent_space, - bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, - idx, - hidx, -) - space = axes(bc) - if Operators.fd_shmem_is_supported(bc) - return fd_operator_evaluate( - bc.op, - bc.work, - space, - idx, - hidx, - bc.args..., - ) - end - op = bc.op - if Operators.should_call_left_boundary(idx, space, bc.op, bc.args...) - Operators.stencil_left_boundary( - op, - Operators.get_boundary(op, Operators.left_boundary_window(space)), - space, - idx, - hidx, - bc.args..., - ) - elseif Operators.should_call_right_boundary(idx, space, bc.op, bc.args...) - Operators.stencil_right_boundary( - op, - Operators.get_boundary(op, Operators.right_boundary_window(space)), - space, - idx, - hidx, - bc.args..., - ) - else - # fallback to interior stencil - Operators.stencil_interior(op, space, idx, hidx, bc.args...) - end -end - -""" - fd_allocate_shmem(shmem_params, b) - -Create a new broadcasted object with necessary share memory allocated, -using `params` nodal points per block. -""" -@inline function fd_allocate_shmem(::ShmemParams, obj) - obj -end -@inline function fd_allocate_shmem( - shmem_params::ShmemParams, - bc::Broadcasted{Style}, -) where {Style} - Broadcasted{Style}( - bc.f, - _fd_allocate_shmem(shmem_params, bc.args...), - bc.axes, - ) -end - -######### MatrixFields -# MatrixField operators are not yet supported, and we must stop recursing because -# we can have something of the form -# MatrixFields.LazyOneArgFDOperatorMatrix{DivergenceF2C{@NamedTuple{}}}(DivergenceF2C{@NamedTuple{}}(NamedTuple())) -# which `fd_shmem_is_supported` will return `true` for. - -@inline fd_allocate_shmem(_, bc::MatrixFields.LazyOperatorBroadcasted) = bc -@inline fd_allocate_shmem(_, bc::MatrixFields.FDOperatorMatrix) = bc -@inline fd_allocate_shmem(_, bc::MatrixFields.LazyOneArgFDOperatorMatrix) = bc -######### - -@inline function fd_allocate_shmem( - shmem_params::ShmemParams, - sbc::StencilBroadcasted{Style}, -) where {Style} - args = _fd_allocate_shmem(shmem_params, sbc.args...) - work = if Operators.fd_shmem_is_supported(sbc) - fd_operator_shmem(sbc.axes, shmem_params, sbc.op, args...) - else - nothing - end - StencilBroadcasted{Style}(sbc.op, args, sbc.axes, work) -end - -@inline _fd_allocate_shmem(::ShmemParams) = () -@inline _fd_allocate_shmem(shmem_params::ShmemParams, arg, xargs...) = ( - fd_allocate_shmem(shmem_params, arg), - _fd_allocate_shmem(shmem_params, xargs...)..., -) - -""" - fd_shmem_needed_per_column(::Base.Broadcast.Broadcasted) - fd_shmem_needed_per_column(::StencilBroadcasted) - -Return the total number of shared memory (in bytes) for the given -broadcast expression. -""" -@inline fd_shmem_needed_per_column(bc) = fd_shmem_needed_per_column(0, bc) -@inline fd_shmem_needed_per_column(shmem_bytes, obj) = shmem_bytes -@inline fd_shmem_needed_per_column( - shmem_bytes, - bc::Broadcasted{Style}, -) where {Style} = - shmem_bytes + _fd_shmem_needed_per_column(shmem_bytes, bc.args) - -@inline function fd_shmem_needed_per_column( - shmem_bytes, - sbc::StencilBroadcasted{Style}, -) where {Style} - shmem_bytes₀ = _fd_shmem_needed_per_column(shmem_bytes, sbc.args) - return if Operators.fd_shmem_is_supported(sbc) - sizeof(return_eltype(sbc.op, sbc.args...)) + shmem_bytes₀ - else - shmem_bytes₀ - end -end - -@inline _fd_shmem_needed_per_column(shmem_bytes::Integer, ::Tuple{}) = - shmem_bytes -@inline _fd_shmem_needed_per_column(shmem_bytes::Integer, args::Tuple{Any}) = - shmem_bytes + fd_shmem_needed_per_column(shmem_bytes::Integer, args[1]) -@inline _fd_shmem_needed_per_column(shmem_bytes::Integer, args::Tuple) = - shmem_bytes + - fd_shmem_needed_per_column(shmem_bytes::Integer, args[1]) + - _fd_shmem_needed_per_column(shmem_bytes::Integer, Base.tail(args)) - - -get_arg_space(bc::StencilBroadcasted, args::Tuple{}) = axes(bc) -get_arg_space(bc::StencilBroadcasted, args::Tuple) = axes(args[1]) - -get_cent_idx(idx::Integer) = idx # center when traversing centers (trivial) -get_face_idx(idx::PlusHalf) = idx # face when traversing faces (trivial) - -get_cent_idx(idx::PlusHalf) = idx + half # center when traversing faces. Convention: use center right of face -get_face_idx(idx::Integer) = idx - half # face when traversing centers. Convention: use face left of center - -""" - fd_resolve_shmem!( - sbc::StencilBroadcasted, - idx, - hidx, - bds - ) - -Recursively stores the arguments to all operators into shared memory, at the -given indices (if they are valid). -""" -Base.@propagate_inbounds function fd_resolve_shmem!( - sbc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, - idx, # top-level index - hidx, - bds, -) - (li, lw, rw, ri) = bds - space = axes(sbc) - arg_space = get_arg_space(sbc, sbc.args) - ᶜidx = get_cent_idx(idx) - ᶠidx = get_face_idx(idx) - - # Here, we use tail-recursion. We visit leaves of the broadcast expression, - # then work our way up. The StencilBroadcasted and Broadcasted can be - # interleaved (e.g., `DivergenceF2C()(f*GradientC2F()(a*b)))`. The leaves of - # the StencilBroadcasted will call `getidx`, below which there are - # (by definition) no more `StencilBroadcasted`, and those `getidx` calls - # will read from global memory. Immediately above those reads, all - # subsequent reads will be from shmem (as they will be caught by the - # `getidx` defined above). - _fd_resolve_shmem!(idx, hidx, bds, sbc.args...) - - # Once we're about ready to fill the shmem, check if shmem is supported for - # this operator - Operators.fd_shmem_is_supported(sbc) || return nothing - - # There are `Nf` threads, where `Nf` is the number of face levels. So, - # each thread is responsible for filling shared memory at its cell center - # (if the broadcasted argument lives on cell centers) - # or cell face (if the broadcasted argument lives on cell faces) index. - # We use `get_face_idx` and `get_cent_idx` to grab the nearest in-bounds - # index, and `get_arg_space` to get the space of the first broadcasted argument - # (the space of all broadcasted arguments must all match, so using the first is valid). - - bc_bds = Operators.window_bounds(space, sbc) - ᵃidx = arg_space isa Operators.AllFaceFiniteDifferenceSpace ? ᶠidx : ᶜidx - - fd_operator_fill_shmem!( - sbc.op, - sbc.work, - bc_bds, - arg_space, - space, - ᵃidx, - hidx, - sbc.args..., - ) - CUDA.sync_threads() - return nothing -end - -Base.@propagate_inbounds function fd_resolve_shmem!( - sbc::StencilBroadcasted, - idx, # top-level index - hidx, - bds, -) - _fd_resolve_shmem!(idx, hidx, bds, sbc.args...) -end - -Base.@propagate_inbounds _fd_resolve_shmem!(idx, hidx, bds) = nothing -Base.@propagate_inbounds function _fd_resolve_shmem!( - idx, - hidx, - bds, - arg, - xargs..., -) - fd_resolve_shmem!(arg, idx, hidx, bds) - _fd_resolve_shmem!(idx, hidx, bds, xargs...) -end - -Base.@propagate_inbounds fd_resolve_shmem!(bc::Broadcasted, idx, hidx, bds) = - _fd_resolve_shmem!(idx, hidx, bds, bc.args...) -@inline fd_resolve_shmem!(obj, idx, hidx, bds) = nothing - -if hasfield(Method, :recursion_relation) - dont_limit = (args...) -> true - for m in methods(fd_resolve_shmem!) - m.recursion_relation = dont_limit - end - for m in methods(_fd_resolve_shmem!) - m.recursion_relation = dont_limit - end - for m in methods(_fd_allocate_shmem) - m.recursion_relation = dont_limit - end - for m in methods(fd_allocate_shmem) - m.recursion_relation = dont_limit - end -end diff --git a/ext/cuda/operators_fd_shmem_is_supported.jl b/ext/cuda/operators_fd_shmem_is_supported.jl deleted file mode 100644 index adf0310de2..0000000000 --- a/ext/cuda/operators_fd_shmem_is_supported.jl +++ /dev/null @@ -1,199 +0,0 @@ -import ClimaCore.MatrixFields -import ClimaCore.Operators: any_fd_shmem_supported - -""" - any_fd_shmem_supported(x) - -Returns a Bool indicating if any broadcasted object can support a shmem style -""" -function any_fd_shmem_supported end - -# Main entry point: call recursive function with 2 args -@inline any_fd_shmem_supported(bc) = any_fd_shmem_supported(false, bc) - -@inline _any_fd_shmem_supported_args(falsesofar, args::Tuple) = - falsesofar || - any_fd_shmem_supported(falsesofar, args[1]) || - _any_fd_shmem_supported_args(falsesofar, Base.tail(args)) - -@inline _any_fd_shmem_supported_args(falsesofar, args::Tuple{Any}) = - falsesofar || any_fd_shmem_supported(falsesofar, args[1]) - -@inline _any_fd_shmem_supported_args(falsesofar, args::Tuple{}) = falsesofar - -@inline function _any_fd_shmem_supported( - falsesofar, - bc::Base.Broadcast.Broadcasted, -) - return falsesofar || _any_fd_shmem_supported_args(falsesofar, bc.args) -end - -@inline _any_fd_shmem_supported(falsesofar, bc) = falsesofar - -@inline _any_fd_shmem_supported(falsesofar, _, x) = falsesofar - -@inline any_fd_shmem_supported(falsesofar, x) = falsesofar # generic fallback -@inline any_fd_shmem_supported(falsesofar, bc::StencilBroadcasted) = - falsesofar || - Operators.fd_shmem_is_supported(bc) || - _any_fd_shmem_supported_args(falsesofar, bc.args) - -@inline any_fd_shmem_supported(falsesofar, bc::Base.Broadcast.Broadcasted) = - falsesofar || _any_fd_shmem_supported_args(falsesofar, bc.args) - -@inline any_fd_shmem_supported(bc::Base.Broadcast.Broadcasted) = - _any_fd_shmem_supported_args(false, bc.args) - - -# Fallback is false: -@inline Operators.fd_shmem_is_supported(bc::StencilBroadcasted) = - Operators.fd_shmem_is_supported(bc.op) - - -""" - any_fd_shmem_style() - -Returns a Bool indicating if any broadcasted object has a shmem style -""" -function any_fd_shmem_style end - -@inline any_fd_shmem_style(bc) = any_fd_shmem_style(false, bc) - -@inline _any_fd_shmem_style_args(falsesofar, args::Tuple) = - falsesofar || - any_fd_shmem_style(falsesofar, args[1]) || - _any_fd_shmem_style_args(falsesofar, Base.tail(args)) - -@inline _any_fd_shmem_style_args(falsesofar, args::Tuple{Any}) = - falsesofar || any_fd_shmem_style(falsesofar, args[1]) - -@inline _any_fd_shmem_style_args(falsesofar, args::Tuple{}) = falsesofar - -@inline function _any_fd_shmem_style(falsesofar, bc::Base.Broadcast.Broadcasted) - return falsesofar || _any_fd_shmem_style_args(falsesofar, bc.args) -end - -@inline _any_fd_shmem_style(falsesofar, bc) = falsesofar - -@inline _any_fd_shmem_style(falsesofar, _, x) = falsesofar - -@inline any_fd_shmem_style(falsesofar, x) = falsesofar # generic fallback -@inline any_fd_shmem_style( - falsesofar, - bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, -) = falsesofar || _any_fd_shmem_style_args(falsesofar, bc.args) - -@inline any_fd_shmem_style(falsesofar, bc::StencilBroadcasted) = - falsesofar || _any_fd_shmem_style_args(falsesofar, bc.args) - -""" - disable_shmem_style() - -For high resolution cases, shmem may not work, so `disable_shmem` transforms a -the boradcast style from -`CUDAWithShmemColumnStencilStyle` to `CUDAColumnStencilStyle`. -""" -function disable_shmem_style end - -@inline disable_shmem_style_args(args::Tuple) = - (disable_shmem_style(args[1]), disable_shmem_style_args(Base.tail(args))...) -@inline disable_shmem_style_args(args::Tuple{Any}) = - (disable_shmem_style(args[1]),) -@inline disable_shmem_style_args(args::Tuple{}) = () - -@inline function disable_shmem_style( - bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, -) - StencilBroadcasted{CUDAColumnStencilStyle}( - bc.op, - disable_shmem_style_args(bc.args), - bc.axes, - nothing, - ) -end - -@inline function disable_shmem_style( - bc::Base.Broadcast.Broadcasted{CUDAWithShmemColumnStencilStyle}, -) - Base.Broadcast.Broadcasted{CUDAColumnStencilStyle}( - bc.f, - disable_shmem_style_args(bc.args), - bc.axes, - ) -end -@inline disable_shmem_style(x) = x - -##### MatrixFields -@inline Operators.fd_shmem_is_supported( - bc::MatrixFields.LazyOperatorBroadcasted, -) = false - -@inline Operators.fd_shmem_is_supported(op::MatrixFields.FDOperatorMatrix) = - false - -@inline Operators.fd_shmem_is_supported( - op::MatrixFields.LazyOneArgFDOperatorMatrix, -) = false -##### - -@inline Operators.fd_shmem_is_supported(op::Operators.AbstractOperator) = - Operators.fd_shmem_is_supported(op, op.bcs) - -@inline Operators.fd_shmem_is_supported( - op::MatrixFields.MultiplyColumnwiseBandMatrixField, -) = false - -@inline Operators.fd_shmem_is_supported( - op::Operators.AbstractOperator, - bcs::NamedTuple, -) = false - -# Add cases here where shmem is supported: - -##### DivergenceF2C -@inline Operators.fd_shmem_is_supported(op::Operators.DivergenceF2C) = - Operators.fd_shmem_is_supported(op, op.bcs) -@inline Operators.fd_shmem_is_supported( - op::Operators.DivergenceF2C, - ::@NamedTuple{}, -) = true -@inline Operators.fd_shmem_is_supported( - op::Operators.DivergenceF2C, - bcs::NamedTuple, -) = - all(values(bcs)) do bc - any(supported_bc -> bc isa supported_bc, (Operators.SetValue,)) - end - -##### GradientC2F -@inline Operators.fd_shmem_is_supported(op::Operators.GradientC2F) = - Operators.fd_shmem_is_supported(op, op.bcs) -@inline Operators.fd_shmem_is_supported( - op::Operators.GradientC2F, - ::@NamedTuple{}, -) = false -@inline Operators.fd_shmem_is_supported( - op::Operators.GradientC2F, - bcs::NamedTuple, -) = - all(values(bcs)) do bc - any(supported_bc -> bc isa supported_bc, (Operators.SetValue,)) - end - -##### InterpolateC2F -@inline Operators.fd_shmem_is_supported(op::Operators.InterpolateC2F) = - Operators.fd_shmem_is_supported(op, op.bcs) -@inline Operators.fd_shmem_is_supported( - op::Operators.InterpolateC2F, - ::@NamedTuple{}, -) = true -@inline Operators.fd_shmem_is_supported( - op::Operators.InterpolateC2F, - bcs::NamedTuple, -) = - all(values(bcs)) do bc - any( - supported_bc -> bc isa supported_bc, - (Operators.SetValue, Operators.SetGradient, Operators.Extrapolate), - ) - end diff --git a/ext/cuda/operators_finite_difference.jl b/ext/cuda/operators_finite_difference.jl index d8f49ffedd..b72cd9515b 100644 --- a/ext/cuda/operators_finite_difference.jl +++ b/ext/cuda/operators_finite_difference.jl @@ -10,7 +10,6 @@ import ClimaCore.Operators: StencilBroadcasted import ClimaCore.Operators: LeftBoundaryWindow, RightBoundaryWindow, Interior struct CUDAColumnStencilStyle <: AbstractStencilStyle end -struct CUDAWithShmemColumnStencilStyle <: AbstractStencilStyle end AbstractStencilStyle(bc, ::ClimaComms.CUDADevice) = CUDAColumnStencilStyle @@ -18,20 +17,13 @@ Base.Broadcast.BroadcastStyle( x::Operators.ColumnStencilStyle, y::CUDAColumnStencilStyle, ) = y - -include("operators_fd_shmem_is_supported.jl") include("operators_fd_eager.jl") -struct ShmemParams{Nv} end -interior_size(::ShmemParams{Nv}) where {Nv} = (Nv,) -boundary_size(::ShmemParams{Nv}) where {Nv} = (1,) function Base.copyto!( out::Field, bc::Union{ StencilBroadcasted{CUDAColumnStencilStyle}, - StencilBroadcasted{CUDAWithShmemColumnStencilStyle}, Broadcasted{CUDAColumnStencilStyle}, - Broadcasted{CUDAWithShmemColumnStencilStyle}, }, mask = Spaces.get_mask(axes(out)), ) @@ -42,99 +34,63 @@ function Base.copyto!( fspace = Spaces.face_space(space) n_face_levels = Spaces.nlevels(fspace) - high_resolution = !(n_face_levels ≤ 256) - # https://github.com/JuliaGPU/CUDA.jl/issues/2672 - # max_shmem = 166912 # CUDA.limit(CUDA.LIMIT_SHMEM_SIZE) # - max_shmem = CUDA.attribute( - device(), - CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, - ) - total_shmem = fd_shmem_needed_per_column(bc) - enough_shmem = total_shmem ≤ max_shmem - # TODO: Use CUDA.limit(CUDA.LIMIT_SHMEM_SIZE) to determine how much shmem should be used - # TODO: add shmem support for masked operations - if Operators.any_fd_shmem_supported(bc) && - !high_resolution && - mask isa NoMask && - enough_shmem && - Operators.use_fd_shmem() - shmem_params = ShmemParams{n_face_levels}() - p = fd_shmem_stencil_partition(us, n_face_levels) + (Ni, Nj, _, Nv, Nh) = DataLayouts.universal_size(out_fv) + # This uses block and grid indices instead of computing cartesian indices from a + # linear index. The launch configuration is optimized for common use case of 64 face + # levels and Ni = Nj = 4. Periodic toppologies and masks are not currently supported + # `eager_copyto_stencil_kernel!` requires a block size of (n_face_levels, Ni, 1) + # this block config is better for VIJFH. It is only used when the total number of + # threads in a block is between 32 and 256 to avoid underutilization of the GPU and + # errors due to too many registers used when the block size is too large. + if !Topologies.isperiodic(space) && mask isa NoMask && + 32 <= n_face_levels * Ni <= 256 + op_matrix_bc = replace_fd_ops(bc) args = ( strip_space(out, space), - strip_space(bc, space), + strip_space(op_matrix_bc, space), axes(out), - bounds, - us, - mask, - shmem_params, ) auto_launch!( - copyto_stencil_kernel_shmem!, + eager_copyto_stencil_kernel!, args; - threads_s = p.threads, - blocks_s = p.blocks, + threads_s = (n_face_levels, Ni, 1), + blocks_s = (1, Nj, Nh), + always_inline = true, + shmem = n_face_levels * Ni * 9 * 4, # see `check_if_fits_in_shmem` for how this is calculated ) + call_post_op_callback() && post_op_callback(out, out, bc) + return out + end + cart_inds = if mask isa NoMask + cartesian_indices(us) else - bc′ = disable_shmem_style(bc) - (Ni, Nj, _, Nv, Nh) = DataLayouts.universal_size(out_fv) - # This uses block and grid indices instead of computing cartesian indices from a - # linear index. The launch configuration is optimized for common use case of 64 face - # levels and Ni = Nj = 4. Periodic toppologies and masks are not currently supported - # `eager_copyto_stencil_kernel!` requires a block size of (n_face_levels, Ni, 1) - # this block config is better for VIJFH. It is only used when the total number of - # threads in a block is between 32 and 256 to avoid underutilization of the GPU and - # errors due to too many registers used when the block size is too large. - if !Topologies.isperiodic(space) && mask isa NoMask && - 32 <= n_face_levels * Ni <= 256 - op_matrix_bc = replace_fd_ops(bc′) - args = ( - strip_space(out, space), - strip_space(op_matrix_bc, space), - axes(out), - ) - auto_launch!( - eager_copyto_stencil_kernel!, - args; - threads_s = (n_face_levels, Ni, 1), - blocks_s = (1, Nj, Nh), - always_inline = true, - shmem = n_face_levels * Ni * 9 * 4, # see `check_if_fits_in_shmem` for how this is calculated - ) - return out - end - @assert !any_fd_shmem_style(bc′) - cart_inds = if mask isa NoMask - cartesian_indices(us) - else - cartesian_indices_mask(us, mask) - end + cartesian_indices_mask(us, mask) + end - args = cudaconvert(( - strip_space(out, space), - strip_space(bc′, space), - axes(out), - bounds, - us, - mask, - cart_inds, - )) + args = cudaconvert(( + strip_space(out, space), + strip_space(bc, space), + axes(out), + bounds, + us, + mask, + cart_inds, + )) - threads = threads_via_occupancy(copyto_stencil_kernel!, args) - n_max_threads = min(threads, get_N(us)) - p = if mask isa NoMask - linear_partition(prod(size(out_fv)), n_max_threads) - else - masked_partition(mask, n_max_threads, us) - end - auto_launch!( - copyto_stencil_kernel!, - args; - threads_s = p.threads, - blocks_s = p.blocks, - ) + threads = threads_via_occupancy(copyto_stencil_kernel!, args) + n_max_threads = min(threads, get_N(us)) + p = if mask isa NoMask + linear_partition(prod(size(out_fv)), n_max_threads) + else + masked_partition(mask, n_max_threads, us) end + auto_launch!( + copyto_stencil_kernel!, + args; + threads_s = p.threads, + blocks_s = p.blocks, + ) call_post_op_callback() && post_op_callback(out, out, bc) return out end @@ -172,43 +128,3 @@ function copyto_stencil_kernel!( end return nothing end - -function copyto_stencil_kernel_shmem!( - out, - bc′::Union{StencilBroadcasted, Broadcasted}, - space, - bds, - us, - mask, - shmem_params::ShmemParams, -) - @inbounds begin - out_fv = Fields.field_values(out) - us = DataLayouts.UniversalSize(out_fv) - I = fd_shmem_stencil_universal_index(space, us) - if fd_shmem_stencil_is_valid_index(I, us) # check that hidx is in bounds - (li, lw, rw, ri) = bds - (i, j, _, v, h) = I.I - hidx = (i, j, h) - idx = v - 1 + li - bc = Operators.reconstruct_placeholder_broadcasted(space, bc′) - bc_shmem = fd_allocate_shmem(shmem_params, bc) # allocates shmem - - fd_resolve_shmem!(bc_shmem, idx, hidx, bds) # recursively fills shmem - CUDA.sync_threads() - - nv = Spaces.nlevels(space) - isactive = if space isa Operators.AllFaceFiniteDifferenceSpace # check that idx is in bounds - idx + half <= nv - else - idx <= nv - end - if isactive - # Call getidx overloaded in operators_fd_shmem_common.jl - val = Operators.getidx(space, bc_shmem, idx, hidx) - setidx!(space, out, idx, hidx, val) - end - end - end - return nothing -end diff --git a/test/Operators/finitedifference/benchmark_fd_ops_shared_memory.jl b/test/Operators/finitedifference/benchmark_fd_ops_shared_memory.jl deleted file mode 100644 index f1b8849f51..0000000000 --- a/test/Operators/finitedifference/benchmark_fd_ops_shared_memory.jl +++ /dev/null @@ -1,114 +0,0 @@ -#= -julia --project=.buildkite -using Revise; include("test/Operators/finitedifference/benchmark_fd_ops_shared_memory.jl") -=# -include("utils_fd_ops_shared_memory.jl") -using BenchmarkTools - -#! format: off -function bench_kernels!(fields) - (; f, ρ, ϕ) = fields - (; ᶜout1, ᶜout2, ᶜout3, ᶜout4, ᶜout5, ᶜout6, ᶜout7, ᶜout8) = fields - device = ClimaComms.device(f) - FT = Spaces.undertype(axes(ϕ)) - div_bcs = Operators.DivergenceF2C(; - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - top = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - div = Operators.DivergenceF2C() - ᶠwinterp = Operators.WeightedInterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - println("ᶜout1: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout1 = $div(Geometry.WVector($f)) - end) - @. ᶜout2 = 0 - println("ᶜout2: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout2 += $div(Geometry.WVector($f) * 2) - end) - println("ᶜout3: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout3 = $div(Geometry.WVector($ᶠwinterp($ϕ, $ρ))) - end) - - println("ᶜout4: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout4 = $div_bcs(Geometry.WVector($f)) - end) - @. ᶜout5 = 0 - println("ᶜout5: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout5 += $div_bcs(Geometry.WVector($f) * 2) - end) - println("ᶜout6: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout6 = $div_bcs(Geometry.WVector($ᶠwinterp($ϕ, $ρ))) - end) - - # from the wild - Ic2f = Operators.InterpolateC2F(; top = Operators.Extrapolate()) - divf2c = Operators.DivergenceF2C(; bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10)))) - # only upward component of divergence - println("ᶜout7: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout7 = $divf2c(Geometry.WVector($Ic2f($ϕ))) # works - end) - println("ᶜout8: ", @benchmark ClimaComms.@cuda_sync $device begin - @. $ᶜout8 = $divf2c($Ic2f(Geometry.WVector($ϕ))) - end) - return nothing -end; - -#! format: on - -@info "GPU benchmark results (Float64) z_elem = 10, h_elem = 30:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float64; z_elem = 10, h_elem = 30); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -@info "GPU benchmark results (Float64) z_elem = 63, h_elem = 30:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float64; z_elem = 63, h_elem = 30); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -@info "GPU benchmark results (Float64) z_elem = 10, h_elem = 100:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float64; z_elem = 10, h_elem = 100); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -@info "GPU benchmark results (Float32) z_elem = 10, h_elem = 30:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float32; z_elem = 10, h_elem = 30); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -@info "GPU benchmark results (Float32) z_elem = 63, h_elem = 30:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float32; z_elem = 63, h_elem = 30); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -@info "GPU benchmark results (Float32) z_elem = 10, h_elem = 100:" -ᶜspace = - get_space_extruded(ClimaComms.device(), Float32; z_elem = 10, h_elem = 100); -ᶠspace = Spaces.face_space(ᶜspace); -fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...); -bench_kernels!(fields) -Utilities.Cache.clean_cache!(); -GC.gc(true); - -nothing diff --git a/test/Operators/finitedifference/unit_fd_ops_shared_memory.jl b/test/Operators/finitedifference/unit_fd_ops_shared_memory.jl deleted file mode 100644 index 79eba774fb..0000000000 --- a/test/Operators/finitedifference/unit_fd_ops_shared_memory.jl +++ /dev/null @@ -1,222 +0,0 @@ -#= -julia --project=.buildkite -julia --check-bounds=yes -g2 --project=.buildkite -using Revise; include("test/Operators/finitedifference/unit_fd_ops_shared_memory.jl") -=# -include("utils_fd_ops_shared_memory.jl") - -Operators.use_fd_shmem() = true -@testset "FD shared memory: dispatch" begin # this ensures that we exercise the correct code-path - FT = Float64 - device = ClimaComms.device() - @test device isa ClimaComms.CUDADevice - ᶜspace = get_space_column(device, FT) - ᶠspace = Spaces.face_space(ᶜspace) - f = Fields.Field(FT, ᶠspace) - c = Fields.Field(FT, ᶜspace) - grad = Operators.GradientF2C() - bc = @. lazy(grad(f)) - @test !Operators.any_fd_shmem_supported(bc) - div = Operators.DivergenceF2C() - bc = @. lazy(div(Geometry.WVector(f))) - @test Operators.any_fd_shmem_supported(bc) - bc = @. lazy(c + div(Geometry.WVector(f))) - @test Operators.any_fd_shmem_supported(bc) - - ᶠinterp = Operators.InterpolateC2F() - div = Operators.DivergenceF2C(; - top = Operators.SetValue(10), - bottom = Operators.SetValue(10), - ) - bc = @. lazy(div(ᶠinterp(Geometry.WVector(c)))) - @test Operators.any_fd_shmem_supported(bc) -end - -@testset "Disable for high resolution" begin - FT = Float64 - device = ClimaComms.device() - ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt) - ᶜspace = get_space_column(device, FT; z_elem = 1000) - ᶠspace = Spaces.face_space(ᶜspace) - f = Fields.Field(FT, ᶠspace) - c = Fields.Field(FT, ᶜspace) - div = Operators.DivergenceF2C() - bc = @. lazy(div(Geometry.WVector(f))) - @test Operators.any_fd_shmem_supported(bc) - @test !ext.any_fd_shmem_style(ext.disable_shmem_style(bc)) - @. c = div(Geometry.WVector(f)) - ᶠgrad = Operators.GradientC2F(; - bottom = Operators.SetValue(FT(0)), - top = Operators.SetValue(FT(0)), - ) - bc = @. lazy(ᶠgrad(c)) - @test Operators.any_fd_shmem_supported(bc) - @test Operators.fd_shmem_is_supported(bc) -end - -@testset "Utility functions" begin - FT = Float64 - device = ClimaComms.device() - ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt) - ᶜspace = get_space_column(device, FT; z_elem = 10) - ᶠspace = Spaces.face_space(ᶜspace) - f = Fields.Field(FT, ᶠspace) - c = Fields.Field(FT, ᶜspace) - fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...) - (; ϕ, ρ) = fields - - div = Operators.DivergenceF2C() - bc = @. lazy(div(Geometry.WVector(f))) - test_face_windows(bc) - - div_bcs = Operators.DivergenceF2C(; - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - top = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - bc = @. lazy(div_bcs(Geometry.WVector(f) * 2)) - test_face_windows(bc) - - ᶠgrad = Operators.GradientC2F(; - bottom = Operators.SetValue(FT(0)), - top = Operators.SetValue(FT(0)), - ) - bc = @. lazy(ᶠgrad(c)) - test_center_windows(bc) - - div = Operators.DivergenceF2C() - ᶠwinterp = Operators.WeightedInterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - bc = @. lazy(div(Geometry.WVector(ᶠwinterp(ϕ, ρ)))) - test_face_windows(bc) - # highly nested cases - ᶜinterp = Operators.InterpolateF2C() - ᶠinterp = Operators.InterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - bc = @. lazy(ᶠgrad(ᶜinterp(ᶠinterp(ᶜinterp(f))))) # exercises very nested operation - test_center_windows(bc) - test_face_windows(bc.args[1]) - test_center_windows(bc.args[1].args[1]) -end - -#! format: off -@testset "Correctness column" begin - ᶜspace_cpu = get_space_column(ClimaComms.CPUSingleThreaded(), Float64) - ᶠspace_cpu = Spaces.face_space(ᶜspace_cpu) - fields_cpu = (; get_fields(ᶜspace_cpu)..., get_fields(ᶠspace_cpu)...) - kernels!(fields_cpu) - @info "Compiled CPU kernels" - - ᶜspace = get_space_column(ClimaComms.device(), Float64) - ClimaComms.device(ᶜspace) isa ClimaComms.CPUSingleThreaded && - @warn "Running on the CPU" - ᶠspace = Spaces.face_space(ᶜspace) - fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...) - kernels!(fields) - @info "Compiled GPU kernels" - - @test compare_cpu_gpu(fields_cpu.ᶜout1, fields.ᶜout1); @test !is_trivial(fields_cpu.ᶜout1) - @test compare_cpu_gpu(fields_cpu.ᶜout2, fields.ᶜout2); @test !is_trivial(fields_cpu.ᶜout2) - @test compare_cpu_gpu(fields_cpu.ᶜout3, fields.ᶜout3); @test !is_trivial(fields_cpu.ᶜout3) - @test compare_cpu_gpu(fields_cpu.ᶜout4, fields.ᶜout4); @test !is_trivial(fields_cpu.ᶜout4) - @test compare_cpu_gpu(fields_cpu.ᶜout5, fields.ᶜout5); @test !is_trivial(fields_cpu.ᶜout5) - @test compare_cpu_gpu(fields_cpu.ᶜout6, fields.ᶜout6); @test !is_trivial(fields_cpu.ᶜout6) - @test compare_cpu_gpu(fields_cpu.ᶜout7, fields.ᶜout7); @test !is_trivial(fields_cpu.ᶜout7) - @test compare_cpu_gpu(fields_cpu.ᶜout8, fields.ᶜout8); @test !is_trivial(fields_cpu.ᶜout8) - @test compare_cpu_gpu(fields_cpu.ᶠout1, fields.ᶠout1); @test !is_trivial(fields_cpu.ᶠout1) - @test compare_cpu_gpu(fields_cpu.ᶠout2, fields.ᶠout2); @test !is_trivial(fields_cpu.ᶠout2) - @test compare_cpu_gpu(fields_cpu.ᶠout1_contra, fields.ᶠout1_contra); @test !is_trivial(fields_cpu.ᶠout1_contra) - @test compare_cpu_gpu(fields_cpu.ᶠout2_contra, fields.ᶠout2_contra); @test !is_trivial(fields_cpu.ᶠout2_contra) - @test compare_cpu_gpu(fields_cpu.ᶜout9, fields.ᶜout9); @test !is_trivial(fields_cpu.ᶜout9) - @test compare_cpu_gpu(fields_cpu.ᶜout10, fields.ᶜout10); @test !is_trivial(fields_cpu.ᶜout10) - @test compare_cpu_gpu(fields_cpu.ᶜout11, fields.ᶜout11); @test !is_trivial(fields_cpu.ᶜout11) - @test compare_cpu_gpu(fields_cpu.ᶜout12, fields.ᶜout12); @test !is_trivial(fields_cpu.ᶜout12) - @test compare_cpu_gpu(fields_cpu.ᶜout13, fields.ᶜout13); @test !is_trivial(fields_cpu.ᶜout13) - @test compare_cpu_gpu(fields_cpu.ᶜout14, fields.ᶜout14); @test !is_trivial(fields_cpu.ᶜout14) - @test compare_cpu_gpu(fields_cpu.ᶜout_uₕ, fields.ᶜout_uₕ); @test !is_trivial(fields_cpu.ᶜout_uₕ) - @test compare_cpu_gpu(fields_cpu.ᶠout3_cov, fields.ᶠout3_cov); @test !is_trivial(fields_cpu.ᶠout3_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout4_cov, fields.ᶠout4_cov); @test !is_trivial(fields_cpu.ᶠout4_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout5_cov, fields.ᶠout5_cov); @test !is_trivial(fields_cpu.ᶠout5_cov) -end - -@testset "Correctness plane" begin - ᶜspace_cpu = get_space_plane(ClimaComms.CPUSingleThreaded(), Float64) - ᶠspace_cpu = Spaces.face_space(ᶜspace_cpu) - fields_cpu = (; get_fields(ᶜspace_cpu)..., get_fields(ᶠspace_cpu)...) - kernels!(fields_cpu) - @info "Compiled CPU kernels" - - ᶜspace = get_space_plane(ClimaComms.device(), Float64) - ClimaComms.device(ᶜspace) isa ClimaComms.CPUSingleThreaded && - @warn "Running on the CPU" - ᶠspace = Spaces.face_space(ᶜspace) - fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...) - kernels!(fields) - @info "Compiled GPU kernels" - - @test compare_cpu_gpu(fields_cpu.ᶜout1, fields.ᶜout1); @test !is_trivial(fields_cpu.ᶜout1) - @test compare_cpu_gpu(fields_cpu.ᶜout2, fields.ᶜout2); @test !is_trivial(fields_cpu.ᶜout2) - @test compare_cpu_gpu(fields_cpu.ᶜout3, fields.ᶜout3); @test !is_trivial(fields_cpu.ᶜout3) - @test compare_cpu_gpu(fields_cpu.ᶜout4, fields.ᶜout4); @test !is_trivial(fields_cpu.ᶜout4) - @test compare_cpu_gpu(fields_cpu.ᶜout5, fields.ᶜout5); @test !is_trivial(fields_cpu.ᶜout5) - @test compare_cpu_gpu(fields_cpu.ᶜout6, fields.ᶜout6); @test !is_trivial(fields_cpu.ᶜout6) - @test compare_cpu_gpu(fields_cpu.ᶜout7, fields.ᶜout7); @test !is_trivial(fields_cpu.ᶜout7) - @test compare_cpu_gpu(fields_cpu.ᶜout8, fields.ᶜout8); @test !is_trivial(fields_cpu.ᶜout8) - @test compare_cpu_gpu(fields_cpu.ᶠout1, fields.ᶠout1); @test !is_trivial(fields_cpu.ᶠout1) - @test compare_cpu_gpu(fields_cpu.ᶠout2, fields.ᶠout2); @test !is_trivial(fields_cpu.ᶠout2) - @test compare_cpu_gpu(fields_cpu.ᶠout1_contra, fields.ᶠout1_contra); @test !is_trivial(fields_cpu.ᶠout1_contra) - @test compare_cpu_gpu(fields_cpu.ᶠout2_contra, fields.ᶠout2_contra); @test !is_trivial(fields_cpu.ᶠout2_contra) - @test compare_cpu_gpu(fields_cpu.ᶜout9, fields.ᶜout9); @test !is_trivial(fields_cpu.ᶜout9) - @test compare_cpu_gpu(fields_cpu.ᶜout10, fields.ᶜout10); @test !is_trivial(fields_cpu.ᶜout10) - @test compare_cpu_gpu(fields_cpu.ᶜout11, fields.ᶜout11); @test !is_trivial(fields_cpu.ᶜout11) - @test compare_cpu_gpu(fields_cpu.ᶜout12, fields.ᶜout12); @test !is_trivial(fields_cpu.ᶜout12) - @test compare_cpu_gpu(fields_cpu.ᶜout13, fields.ᶜout13); @test !is_trivial(fields_cpu.ᶜout13) - @test compare_cpu_gpu(fields_cpu.ᶜout14, fields.ᶜout14); @test !is_trivial(fields_cpu.ᶜout14) - @test compare_cpu_gpu(fields_cpu.ᶜout_uₕ, fields.ᶜout_uₕ); @test !is_trivial(fields_cpu.ᶜout_uₕ) - @test compare_cpu_gpu(fields_cpu.ᶠout3_cov, fields.ᶠout3_cov); @test !is_trivial(fields_cpu.ᶠout3_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout4_cov, fields.ᶠout4_cov); @test !is_trivial(fields_cpu.ᶠout4_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout5_cov, fields.ᶠout5_cov); @test !is_trivial(fields_cpu.ᶠout5_cov) -end - -@testset "Correctness extruded cubed sphere" begin - ᶜspace_cpu = get_space_extruded(ClimaComms.CPUSingleThreaded(), Float64) - ᶠspace_cpu = Spaces.face_space(ᶜspace_cpu) - fields_cpu = (; get_fields(ᶜspace_cpu)..., get_fields(ᶠspace_cpu)...) - kernels!(fields_cpu) - @info "Compiled CPU kernels" - - ᶜspace = get_space_extruded(ClimaComms.device(), Float64) - ᶠspace = Spaces.face_space(ᶜspace) - fields = (; get_fields(ᶜspace)..., get_fields(ᶠspace)...) - kernels!(fields) - @info "Compiled GPU kernels" - - @test compare_cpu_gpu(fields_cpu.ᶜout1, fields.ᶜout1); @test !is_trivial(fields_cpu.ᶜout1) - @test compare_cpu_gpu(fields_cpu.ᶜout2, fields.ᶜout2); @test !is_trivial(fields_cpu.ᶜout2) - @test compare_cpu_gpu(fields_cpu.ᶜout3, fields.ᶜout3); @test !is_trivial(fields_cpu.ᶜout3) - @test compare_cpu_gpu(fields_cpu.ᶜout4, fields.ᶜout4); @test !is_trivial(fields_cpu.ᶜout4) - @test compare_cpu_gpu(fields_cpu.ᶜout5, fields.ᶜout5); @test !is_trivial(fields_cpu.ᶜout5) - @test compare_cpu_gpu(fields_cpu.ᶜout6, fields.ᶜout6); @test !is_trivial(fields_cpu.ᶜout6) - @test compare_cpu_gpu(fields_cpu.ᶜout7, fields.ᶜout7); @test !is_trivial(fields_cpu.ᶜout7) - @test compare_cpu_gpu(fields_cpu.ᶜout8, fields.ᶜout8); @test !is_trivial(fields_cpu.ᶜout8) - @test compare_cpu_gpu(fields_cpu.ᶠout1, fields.ᶠout1); @test !is_trivial(fields_cpu.ᶠout1) - @test compare_cpu_gpu(fields_cpu.ᶠout2, fields.ᶠout2); @test !is_trivial(fields_cpu.ᶠout2) - @test compare_cpu_gpu(fields_cpu.ᶠout1_contra, fields.ᶠout1_contra); @test !is_trivial(fields_cpu.ᶠout1_contra) - @test compare_cpu_gpu(fields_cpu.ᶠout2_contra, fields.ᶠout2_contra); @test !is_trivial(fields_cpu.ᶠout2_contra) - @test compare_cpu_gpu(fields_cpu.ᶜout9, fields.ᶜout9); @test !is_trivial(fields_cpu.ᶜout9) - @test compare_cpu_gpu(fields_cpu.ᶜout10, fields.ᶜout10); @test !is_trivial(fields_cpu.ᶜout10) - @test compare_cpu_gpu(fields_cpu.ᶜout11, fields.ᶜout11); @test !is_trivial(fields_cpu.ᶜout11) - @test compare_cpu_gpu(fields_cpu.ᶜout12, fields.ᶜout12); @test !is_trivial(fields_cpu.ᶜout12) - @test compare_cpu_gpu(fields_cpu.ᶜout13, fields.ᶜout13); @test !is_trivial(fields_cpu.ᶜout13) - @test compare_cpu_gpu(fields_cpu.ᶜout14, fields.ᶜout14); @test !is_trivial(fields_cpu.ᶜout14) - @test compare_cpu_gpu(fields_cpu.ᶜout_uₕ, fields.ᶜout_uₕ); @test !is_trivial(fields_cpu.ᶜout_uₕ) - @test compare_cpu_gpu(fields_cpu.ᶠout3_cov, fields.ᶠout3_cov); @test !is_trivial(fields_cpu.ᶠout3_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout4_cov, fields.ᶠout4_cov); @test !is_trivial(fields_cpu.ᶠout4_cov) - @test compare_cpu_gpu(fields_cpu.ᶠout5_cov, fields.ᶠout5_cov); @test !is_trivial(fields_cpu.ᶠout5_cov) -end - -#! format: on -nothing diff --git a/test/Operators/finitedifference/utils_fd_ops_shared_memory.jl b/test/Operators/finitedifference/utils_fd_ops_shared_memory.jl deleted file mode 100644 index e1c7976fe8..0000000000 --- a/test/Operators/finitedifference/utils_fd_ops_shared_memory.jl +++ /dev/null @@ -1,393 +0,0 @@ -ENV["CLIMACOMMS_DEVICE"] = "CUDA"; # requires cuda -using LazyBroadcast: lazy -using ClimaCore.Utilities: half -using Test, ClimaComms -ClimaComms.@import_required_backends; -using ClimaCore: Geometry, Spaces, Fields, Operators, ClimaCore; -using ClimaCore: Utilities, DataLayouts -using ClimaCore.CommonSpaces; -import ClimaCore.Geometry: ⊗ - -get_space_extruded(dev, FT; z_elem = 63, h_elem = 30) = - ExtrudedCubedSphereSpace( - FT; - device = dev, - z_elem, - z_min = 0, - z_max = 1, - radius = 10, - h_elem, - n_quad_points = 4, - staggering = CellCenter(), - ); - -get_space_plane(dev, FT; z_elem = 100) = SliceXZSpace( - FT; - device = dev, - z_elem, - x_min = 0, - x_max = 100000, - z_min = 0, - z_max = 21000, - periodic_x = false, - n_quad_points = 4, - x_elem = 100, - staggering = CellCenter(), -); - -get_space_column(dev, FT; z_elem = 10) = ColumnSpace( - FT; - device = dev, - z_elem = z_elem, - z_min = 0, - z_max = 1, - staggering = CellCenter(), -); - -function test_face_windows(Base.@nospecialize(bc)) - ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt) - @test bc isa Operators.StencilBroadcasted - space = axes(bc) - ᶜspace = Spaces.center_space(space) - arg_space = ext.get_arg_space(bc, bc.args) - bc_bds = Operators.window_bounds(space, bc) - for lev in 1:(Spaces.nlevels(ᶜspace) + 1) - idx = (lev - half)::Utilities.PlusHalf - args = (idx, arg_space, bc_bds, bc.op, bc.args...) - B = ( - ext.in_left_boundary_window(args...), - ext.in_interior(args...), - ext.in_right_boundary_window(args...), - ) - count(B) == 1 || @show idx, B - @test count(B) == 1 - end -end - -@inline function test_center_windows(Base.@nospecialize(bc)) - ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt) - @test bc isa Operators.StencilBroadcasted - space = axes(bc) - ᶜspace = Spaces.center_space(space) - arg_space = ext.get_arg_space(bc, bc.args) - bc_bds = Operators.window_bounds(space, bc) - for idx in 1:Spaces.nlevels(ᶜspace) - args = (idx, arg_space, bc_bds, bc.op, bc.args...) - B = ( - ext.in_left_boundary_window(args...), - ext.in_interior(args...), - ext.in_right_boundary_window(args...), - idx == Spaces.nlevels(ᶜspace) + 1, - ) - count(B) == 1 || @show idx, B - @test count(B) == 1 - end -end - -function kernels!(fields) - (; f, ρ, ϕ) = fields - (; ᶜout1, ᶜout2, ᶜout3, ᶜout4, ᶜout5, ᶜout6, ᶜout7, ᶜout8, ᶜout9) = fields - (; ᶜout10, ᶜout11, ᶜout12, ᶜout13, ᶜout14) = fields - (; ᶠout1, ᶠout2) = fields - (; ᶠout1_contra, ᶠout2_contra) = fields - (; ᶠout3_cov, ᶠout4_cov, ᶠout5_cov) = fields - (; w_cov) = fields - (; ᶜout_uₕ, ᶜuₕ) = fields - FT = Spaces.undertype(axes(ϕ)) - div_bcs = Operators.DivergenceF2C(; - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - top = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - div = Operators.DivergenceF2C() - ᶠwinterp = Operators.WeightedInterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - @. ᶜout1 = div(Geometry.WVector(f)) - @. ᶜout2 = 0 - @. ᶜout2 += div(Geometry.WVector(f) * 2) - @. ᶜout3 = div(Geometry.WVector(ᶠwinterp(ϕ, ρ))) - - @. ᶜout4 = div_bcs(Geometry.WVector(f)) - @. ᶜout5 = 0 - @. ᶜout5 += div_bcs(Geometry.WVector(f) * 2) - @. ᶜout6 = div_bcs(Geometry.WVector(ᶠwinterp(ϕ, ρ))) - - # from the wild - Ic2f = Operators.InterpolateC2F(; top = Operators.Extrapolate()) - divf2c = Operators.DivergenceF2C(; - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - # only upward component of divergence - @. ᶜout7 = divf2c(Geometry.WVector(Ic2f(ϕ))) - @. ᶜout8 = divf2c(Ic2f(Geometry.WVector(ϕ))) - - upwind = Operators.UpwindBiasedProductC2F(; - bottom = Operators.SetValue(FT(0)), - top = Operators.SetValue(FT(0)), - ) - @. ᶠout1_contra = upwind(w_cov, ϕ) - - upwind = Operators.UpwindBiasedProductC2F(; - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - @. ᶠout2_contra = upwind(w_cov, ϕ) - - upwind = Operators.Upwind3rdOrderBiasedProductC2F(; - # bottom = Operators.ThirdOrderOneSided(), - # top = Operators.ThirdOrderOneSided() - bottom = Operators.FirstOrderOneSided(), - top = Operators.FirstOrderOneSided(), - ) - outer = (; - bottom = Operators.SetValue(Geometry.Contravariant3Vector(FT(0.0))), - top = Operators.SetValue(Geometry.Contravariant3Vector(FT(0.0))), - ) - divf2c = Operators.DivergenceF2C(outer) - @. ᶜout9 = divf2c(upwind(w_cov, ϕ)) - - divf2c_vl = Operators.DivergenceF2C( - bottom = Operators.SetValue(Geometry.WVector(FT(10))), - top = Operators.SetValue(Geometry.WVector(FT(10))), - ) - limiter_method = Operators.AlgebraicMean() - VanLeerMethod = Operators.LinVanLeerC2F( - bottom = Operators.FirstOrderOneSided(), - top = Operators.FirstOrderOneSided(), - constraint = limiter_method, - ) - Δt = FT(1) - @. ᶜout10 = -divf2c_vl(VanLeerMethod(w_cov, ϕ, Δt)) - - inner = (;) - outer = set_value_divgrad_uₕ_maybe_field_bcs(axes(ϕ)) - - grad = Operators.GradientC2F(inner) - div_uh = Operators.DivergenceF2C(outer) - @. ᶜout_uₕ = div_uh(f * grad(ᶜuₕ)) - - ᶠgrad = Operators.GradientC2F(; - bottom = Operators.SetValue(FT(10)), - top = Operators.SetValue(FT(10)), - ) - @. ᶠout3_cov = ᶠgrad(ϕ) - - ᶠgrad = Operators.GradientC2F(; - bottom = Operators.SetValue(FT(10)), - top = Operators.SetValue(FT(10)), - ) - ᶜinterp = Operators.InterpolateF2C() - ᶠinterp = Operators.InterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - @. ᶠout4_cov = ᶠgrad(ᶜinterp(f)) - @. ᶠout2 = ᶠinterp(ᶜinterp(f)) # exercises very nested operation - @. ᶠout5_cov = ᶠgrad(ᶜinterp(ᶠinterp(ᶜinterp(f)))) # exercises very nested operation - @. ᶜout14 = ᶜinterp(ᶠinterp(ᶜinterp(f))) # exercises very nested operation - - @. ᶠout4_cov = ᶠgrad(ᶜinterp(f)) - ᶜdiv = Operators.DivergenceF2C( - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - top = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - @. ᶜout13 = ᶜdiv(Geometry.WVector(ᶠinterp(ᶜinterp(f)))) # exercises very nested operation - - ᶠgrad_top_bcs = Operators.GradientC2F(; top = Operators.SetValue(FT(10))) - divf2c = Operators.DivergenceF2C( - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - @. ᶜout11 = 0 - @. ᶜout11 += divf2c(ᶠgrad_top_bcs(ϕ)) - - ᶠgrad_no_bc = Operators.GradientC2F() - div_bcs = Operators.DivergenceF2C(; - bottom = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - top = Operators.SetValue(Geometry.Covariant3Vector(FT(10))), - ) - @. ᶜout12 = 0 - @. ᶜout12 += div_bcs(ᶠgrad_no_bc(ϕ)) - - ᶠinterp = Operators.InterpolateC2F( - bottom = Operators.Extrapolate(), - top = Operators.Extrapolate(), - ) - @. ᶠout1 = ᶠinterp(ϕ) - - return nothing -end; - -function get_fields(space::Operators.AllFaceFiniteDifferenceSpace) - FT = Spaces.undertype(space) - (; z) = Fields.coordinate_field(space) - K = (:f,) - V = ntuple(i -> Fields.zeros(space), length(K)) - K_contra = (ntuple(i -> Symbol("ᶠout$(i)_contra"), 8)...,) - V_contra = ntuple( - i -> Fields.Field(Geometry.Contravariant3Vector{FT}, space), - length(K_contra), - ) - K_scalar = (ntuple(i -> Symbol("ᶠout$(i)"), 3)...,) - V_scalar = ntuple(i -> Fields.Field(FT, space), length(K_scalar)) - K_cov_out = (ntuple(i -> Symbol("ᶠout$(i)_cov"), 8)...,) - V_cov_out = ntuple( - i -> Fields.zeros(Geometry.Covariant3Vector{FT}, space), - length(K_cov_out), - ) - K_cov = (:w_cov,) - V_cov = ntuple( - i -> Fields.Field(Geometry.Covariant3Vector{FT}, space), - length(K_cov), - ) - nt = (; - zip(K, V)..., - zip(K_scalar, V_scalar)..., - zip(K_contra, V_contra)..., - zip(K_cov, V_cov)..., - zip(K_cov_out, V_cov_out)..., - ) - @. nt.f = sin(z) - @. nt.w_cov.components.data.:1 = sin(z) - return nt -end - -function get_fields(space::Operators.AllCenterFiniteDifferenceSpace) - FT = Spaces.undertype(space) - K = (ntuple(i -> Symbol("ᶜout$i"), 14)..., :ρ, :ϕ) - V = ntuple(i -> Fields.zeros(space), length(K)) - nt = (; - zip(K, V)..., - ᶜout_uₕ = Fields.Field(Geometry.Covariant12Vector{FT}, space), - ᶜuₕ = Fields.Field(Geometry.Covariant12Vector{FT}, space), - ) - (; z) = Fields.coordinate_field(space) - @. nt.ρ = sin(z) - @. nt.ϕ = sin(z) - @. nt.ᶜout_uₕ.components.data.:1 = 0 - @. nt.ᶜout_uₕ.components.data.:2 = 0 - @. nt.ᶜuₕ.components.data.:1 = sin(z) - @. nt.ᶜuₕ.components.data.:2 = sin(z) - return nt -end - -function set_value_divgrad_uₕ_bcs(space) # real-world example - FT = Spaces.undertype(space) - top_val = - Geometry.Contravariant3Vector(FT(0)) ⊗ - Geometry.Covariant12Vector(FT(0), FT(0)) - bottom_val = - Geometry.Contravariant3Vector(FT(0)) ⊗ - Geometry.Covariant12Vector(FT(0), FT(0)) - return (; - top = Operators.SetValue(top_val), - bottom = Operators.Extrapolate(), - ) -end - -function set_value_divgrad_uₕ_maybe_field_bcs(space) # real-world example - FT = Spaces.undertype(space) - top_val = - Geometry.Contravariant3Vector(FT(0)) ⊗ - Geometry.Covariant12Vector(FT(0), FT(0)) - if hasproperty(space, :horizontal_space) - z_bottom = Spaces.level(Fields.coordinate_field(space).z, 1) - bottom_val = - Geometry.Contravariant3Vector.(zeros(axes(z_bottom))) .⊗ - Geometry.Covariant12Vector.( - zeros(axes(z_bottom)), - zeros(axes(z_bottom)), - ) - return (; - top = Operators.SetValue(top_val), - bottom = Operators.SetValue(.-bottom_val), - ) - else - return (; - top = Operators.SetValue(top_val), - bottom = Operators.SetValue(top_val), - ) - end -end - - -function compare_cpu_gpu(cpu, gpu; print_diff = true, C_best = 10) - # there are some odd errors that build up when run without debug / bounds checks: - space = axes(cpu) - are_boundschecks_forced = Base.JLOptions().check_bounds == 1 - absΔ = abs.(parent(cpu) .- Array(parent(gpu))) - max_allowed_err = if space isa Spaces.FiniteDifferenceSpace - are_boundschecks_forced ? 1000 * eps() : 10000000 * eps() - else - 1e-9 - end - max_err = maximum(absΔ) - gpu_matches_cpu = max_err <= max_allowed_err - gpu_matches_cpu || @show max_err - z = Fields.field_values(Fields.coordinate_field(space)) - C = count(x -> x <= max_allowed_err, absΔ) - if !gpu_matches_cpu && print_diff - N = 3 - if z isa DataLayouts.VF - @show parent(cpu)[1:N] - @show parent(gpu)[1:N] - @show parent(cpu)[(end - N):end] - @show parent(gpu)[(end - N):end] - elseif z isa DataLayouts.VIJFH - @show parent(cpu)[1:N, 1, 1, 1, end] - @show parent(gpu)[1:N, 1, 1, 1, end] - @show parent(cpu)[(end - N):end, 1, 1, 1, end] - @show parent(gpu)[(end - N):end, 1, 1, 1, end] - elseif z isa DataLayouts.VIFH - @show parent(cpu)[1:N, 1, 1, end] - @show parent(gpu)[1:N, 1, 1, end] - @show parent(cpu)[(end - N):end, 1, 1, end] - @show parent(gpu)[(end - N):end, 1, 1, end] - else - error("Unsupported type: $(typeof(z))") - end - end - return gpu_matches_cpu -end - -# This function is useful for debugging new cases. -function compare_cpu_gpu_incremental(cpu, gpu; print_diff = true, C_best = 10) - # there are some odd errors that build up when run without debug / bounds checks: - space = axes(cpu) - z = Fields.field_values(Fields.coordinate_field(space)) - are_boundschecks_forced = Base.JLOptions().check_bounds == 1 - absΔ = abs.(parent(cpu) .- Array(parent(gpu))) - max_err = are_boundschecks_forced ? 10000 * eps() : 10000000 * eps() - B = maximum(absΔ) <= max_err - C = count(x -> x <= max_err, absΔ) - @test C ≥ C_best - if !(C_best == 10) - C > C_best && @show C_best - @test_broken C > C_best - end - if !B && print_diff - if z isa DataLayouts.VF - @show parent(cpu)[1:3] - @show parent(gpu)[1:3] - @show parent(cpu)[(end - 3):end] - @show parent(gpu)[(end - 3):end] - elseif z isa DataLayouts.VIJFH - @show parent(cpu)[1:3, 1, 1, 1, end] - @show parent(gpu)[1:3, 1, 1, 1, end] - @show parent(cpu)[(end - 3):end, 1, 1, 1, end] - @show parent(gpu)[(end - 3):end, 1, 1, 1, end] - elseif z isa DataLayouts.VIJFH - @show parent(cpu)[1:3, 1, 1, end] - @show parent(gpu)[1:3, 1, 1, end] - @show parent(cpu)[(end - 3):end, 1, 1, end] - @show parent(gpu)[(end - 3):end, 1, 1, end] - else - error("Unsupported type: $(typeof(z))") - end - end - return true -end - -is_trivial(x) = length(parent(x)) == count(iszero, parent(x)) # Make sure we don't have a trivial solution - -nothing