Compute stencil location on the fly (#2284)

Make non-breaking with promote_bcs

Author: charleskawczynski

examples/column/hydrostatic_ekman.jl

@@ -133,7 +133,7 @@ function tendency!(dY, Y, _, t)
     # TODO!: Undesirable casting to vector required
     @. dρθ = -∂c(w * If(ρθ)) + ρ * ∂c(ν * ∂f(ρθ / ρ))
 
-    uv_1 = Operators.getidx(axes(uv), uv, Operators.Interior(), 1)
+    uv_1 = Operators.getidx(axes(uv), uv, 1)
     u_wind = LinearAlgebra.norm(uv_1)
 
     A = Operators.AdvectionC2C(

ext/cuda/operators_fd_shmem.jl

@@ -20,7 +20,6 @@ end
 Base.@propagate_inbounds function fd_operator_fill_shmem!(
     op::Operators.DivergenceF2C,
     (Ju³, lJu³, rJu³),
-    loc,
     bc_bds,
     arg_space,
     space,
@@ -31,12 +30,24 @@ Base.@propagate_inbounds function fd_operator_fill_shmem!(
     @inbounds begin
         vt = threadIdx().x
         lg = Geometry.LocalGeometry(space, idx, hidx)
-        if !on_boundary(space, op, loc, idx)
-            u³ = Operators.getidx(space, arg, loc, idx, hidx)
+        if !on_boundary(idx, space, op)
+            u³ = Operators.getidx(space, arg, idx, hidx)
             Ju³[vt] = Geometry.Jcontravariant3(u³, lg)
-        else
-            bc = Operators.get_boundary(op, loc)
-            ub = Operators.getidx(space, bc.val, loc, nothing, hidx)
+        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)
@@ -52,7 +63,6 @@ end
 Base.@propagate_inbounds function fd_operator_evaluate(
     op::Operators.DivergenceF2C,
     (Ju³, lJu³, rJu³),
-    loc,
     space,
     idx::Integer,
     hidx,
@@ -61,12 +71,16 @@ Base.@propagate_inbounds function fd_operator_evaluate(
     @inbounds begin
         vt = threadIdx().x
         lg = Geometry.LocalGeometry(space, idx, hidx)
-        if !on_boundary(space, op, loc, idx)
+        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
-            bc = Operators.get_boundary(op, loc)
+            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
@@ -109,7 +123,6 @@ end
 Base.@propagate_inbounds function fd_operator_fill_shmem!(
     op::Operators.GradientC2F,
     (u, lb, rb),
-    loc, # can be any location
     bc_bds,
     arg_space,
     space,
@@ -121,23 +134,20 @@ Base.@propagate_inbounds function fd_operator_fill_shmem!(
         vt = threadIdx().x
         cov3 = Geometry.Covariant3Vector(1)
         if in_domain(idx, arg_space)
-            u[vt] = cov3 ⊗ Operators.getidx(space, arg, loc, idx, hidx)
+            u[vt] = cov3 ⊗ Operators.getidx(space, arg, idx, hidx)
         else # idx can be Spaces.nlevels(ᶜspace)+1 because threads must extend to faces
             ᶜspace = Spaces.center_space(arg_space)
             @assert idx == Spaces.nlevels(ᶜspace) + 1
         end
         if on_any_boundary(idx, space, op)
-            lloc =
-                Operators.LeftBoundaryWindow{Spaces.left_boundary_name(space)}()
-            rloc = Operators.RightBoundaryWindow{
-                Spaces.right_boundary_name(space),
-            }()
+            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, bloc, nothing, hidx)
+            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
@@ -154,7 +164,6 @@ end
 Base.@propagate_inbounds function fd_operator_evaluate(
     op::Operators.GradientC2F,
     (u, lb, rb),
-    loc,
     space,
     idx::PlusHalf,
     hidx,
@@ -163,12 +172,16 @@ Base.@propagate_inbounds function fd_operator_evaluate(
     @inbounds begin
         vt = threadIdx().x
         lg = Geometry.LocalGeometry(space, idx, hidx)
-        if !on_boundary(space, op, loc, idx)
+        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
-            bc = Operators.get_boundary(op, loc)
+            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

ext/cuda/operators_fd_shmem_common.jl

@@ -9,17 +9,13 @@ import ClimaCore.Utilities
 ##### Boundary helpers
 #####
 
-@inline function has_left_boundary(space, op)
-    lloc = Operators.LeftBoundaryWindow{Spaces.left_boundary_name(space)}()
-    return Operators.has_boundary(op, lloc)
-end
-@inline function has_right_boundary(space, op)
-    rloc = Operators.RightBoundaryWindow{Spaces.right_boundary_name(space)}()
-    return Operators.has_boundary(op, rloc)
-end
+@inline has_left_boundary(space, op) =
+    Operators.has_boundary(op, Operators.left_boundary_window(space))
+@inline has_right_boundary(space, op) =
+    Operators.has_boundary(op, Operators.right_boundary_window(space))
 
-@inline on_boundary(space, op, loc, idx) =
-    Operators.has_boundary(op, loc) && on_boundary(idx, space)
+@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) =
     has_left_boundary(space, op) && on_left_boundary(idx, space)
@@ -92,8 +88,7 @@ end
     op,
     args...,
 )
-    lloc = Operators.LeftBoundaryWindow{Spaces.left_boundary_name(space)}()
-    Operators.should_call_left_boundary(idx, space, lloc, op, args...) ||
+    Operators.should_call_left_boundary(idx, space, op, args...) ||
         in_left_boundary_window_range(idx, bc_bds)
 end
 
@@ -104,8 +99,7 @@ end
     op,
     args...,
 )
-    rloc = Operators.RightBoundaryWindow{Spaces.right_boundary_name(space)}()
-    Operators.should_call_right_boundary(idx, space, rloc, op, args...) ||
+    Operators.should_call_right_boundary(idx, space, op, args...) ||
         in_right_boundary_window_range(idx, bc_bds)
 end
 
@@ -146,8 +140,7 @@ end
     op,
     args...,
 )
-    lloc = Operators.LeftBoundaryWindow{Spaces.left_boundary_name(space)}()
-    Operators.should_call_left_boundary(idx, space, lloc, op, args...) ||
+    Operators.should_call_left_boundary(idx, space, op, args...) ||
         in_left_boundary_window_range(idx, bc_bds)
 end
 
@@ -158,10 +151,9 @@ end
     op,
     args...,
 )
-    rloc = Operators.RightBoundaryWindow{Spaces.right_boundary_name(space)}()
     ᶜspace = Spaces.center_space(space)
     idx > Spaces.nlevels(ᶜspace) && return false # short-circuit if
-    Operators.should_call_right_boundary(idx, space, rloc, op, args...) ||
+    Operators.should_call_right_boundary(idx, space, op, args...) ||
         in_right_boundary_window_range(idx, bc_bds)
 end
 
@@ -172,30 +164,6 @@ end
 Base.@propagate_inbounds function getidx(
     parent_space,
     bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle},
-    loc::Interior,
-    idx,
-    hidx,
-)
-    space = axes(bc)
-    if Operators.fd_shmem_is_supported(bc)
-        return fd_operator_evaluate(
-            bc.op,
-            bc.work,
-            loc,
-            space,
-            idx,
-            hidx,
-            bc.args...,
-        )
-    end
-    Operators.stencil_interior(bc.op, loc, space, idx, hidx, bc.args...)
-end
-
-
-Base.@propagate_inbounds function getidx(
-    parent_space,
-    bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle},
-    loc::Operators.LeftBoundaryWindow,
     idx,
     hidx,
 )
@@ -204,63 +172,34 @@ Base.@propagate_inbounds function getidx(
         return fd_operator_evaluate(
             bc.op,
             bc.work,
-            loc,
             space,
             idx,
             hidx,
             bc.args...,
         )
     end
     op = bc.op
-    if Operators.should_call_left_boundary(idx, space, loc, bc.op, bc.args...)
+    if Operators.should_call_left_boundary(idx, space, bc.op, bc.args...)
         Operators.stencil_left_boundary(
             op,
-            Operators.get_boundary(op, loc),
-            loc,
-            space,
-            idx,
-            hidx,
-            bc.args...,
-        )
-    else
-        # fallback to interior stencil
-        Operators.stencil_interior(op, loc, space, idx, hidx, bc.args...)
-    end
-end
-
-Base.@propagate_inbounds function getidx(
-    parent_space,
-    bc::StencilBroadcasted{CUDAWithShmemColumnStencilStyle},
-    loc::Operators.RightBoundaryWindow,
-    idx,
-    hidx,
-)
-    space = axes(bc)
-    if Operators.fd_shmem_is_supported(bc)
-        return fd_operator_evaluate(
-            bc.op,
-            bc.work,
-            loc,
+            Operators.get_boundary(op, Operators.left_boundary_window(space)),
             space,
             idx,
             hidx,
             bc.args...,
         )
-    end
-    op = bc.op
-    if Operators.should_call_right_boundary(idx, space, loc, bc.op, bc.args...)
+    elseif Operators.should_call_right_boundary(idx, space, bc.op, bc.args...)
         Operators.stencil_right_boundary(
             op,
-            Operators.get_boundary(op, loc),
-            loc,
+            Operators.get_boundary(op, Operators.right_boundary_window(space)),
             space,
             idx,
             hidx,
             bc.args...,
         )
     else
         # fallback to interior stencil
-        Operators.stencil_interior(op, loc, space, idx, hidx, bc.args...)
+        Operators.stencil_interior(op, space, idx, hidx, bc.args...)
     end
 end
 
@@ -375,9 +314,6 @@ Base.@propagate_inbounds function fd_resolve_shmem!(
 )
     (li, lw, rw, ri) = bds
     space = axes(sbc)
-
-    ᶜspace = Spaces.center_space(space)
-    ᶠspace = Spaces.face_space(space)
     arg_space = get_arg_space(sbc, sbc.args)
     ᶜidx = get_cent_idx(idx)
     ᶠidx = get_face_idx(idx)
@@ -387,13 +323,6 @@ Base.@propagate_inbounds function fd_resolve_shmem!(
     # After recursion, check if shmem is supported for this operator
     Operators.fd_shmem_is_supported(sbc) || return nothing
 
-    (; op) = sbc
-    lloc = Operators.LeftBoundaryWindow{Spaces.left_boundary_name(space)}()
-    rloc = Operators.RightBoundaryWindow{Spaces.right_boundary_name(space)}()
-    iloc = Operators.Interior()
-
-    IP = Topologies.isperiodic(Spaces.vertical_topology(space))
-
     # 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)
@@ -403,52 +332,18 @@ Base.@propagate_inbounds function fd_resolve_shmem!(
     # (the space of all broadcasted arguments must all match, so using the first is valid).
 
     bc_bds = Operators.window_bounds(space, sbc)
-    (bc_li, bc_lw, bc_rw, bc_ri) = bc_bds
     ᵃidx = arg_space isa Operators.AllFaceFiniteDifferenceSpace ? ᶠidx : ᶜidx
 
-    if in_interior(ᵃidx, arg_space, bc_bds, sbc.op, sbc.args...)
-        fd_operator_fill_shmem!(
-            sbc.op,
-            sbc.work,
-            iloc,
-            bc_bds,
-            arg_space,
-            space,
-            ᵃidx,
-            hidx,
-            sbc.args...,
-        )
-    elseif in_left_boundary_window(ᵃidx, arg_space, bc_bds, sbc.op, sbc.args...)
-        fd_operator_fill_shmem!(
-            sbc.op,
-            sbc.work,
-            lloc,
-            bc_bds,
-            arg_space,
-            space,
-            ᵃidx,
-            hidx,
-            sbc.args...,
-        )
-    elseif in_right_boundary_window(
-        ᵃidx,
-        arg_space,
-        bc_bds,
+    fd_operator_fill_shmem!(
         sbc.op,
+        sbc.work,
+        bc_bds,
+        arg_space,
+        space,
+        ᵃidx,
+        hidx,
         sbc.args...,
     )
-        fd_operator_fill_shmem!(
-            sbc.op,
-            sbc.work,
-            rloc,
-            bc_bds,
-            arg_space,
-            space,
-            ᵃidx,
-            hidx,
-            sbc.args...,
-        )
-    end
     CUDA.sync_threads()
     return nothing
 end