Add shmem for GradientC2F (#2268)

charleskawczynski · web-flow · commit e506eb4bea8c · 2025-04-02T14:28:49.000-04:00
diff --git a/ext/cuda/operators_fd_shmem.jl b/ext/cuda/operators_fd_shmem.jl
@@ -1,6 +1,7 @@
 import ClimaCore: DataLayouts, Spaces, Geometry, RecursiveApply, DataLayouts
 import CUDA
 import ClimaCore.Operators: return_eltype, get_local_geometry
+import ClimaCore.Geometry: ⊗
 
 Base.@propagate_inbounds function fd_operator_shmem(
     space,
@@ -92,3 +93,103 @@ Base.@propagate_inbounds function fd_operator_evaluate(
         return (Ju³₊ ⊟ Ju³₋) ⊠ local_geometry.invJ
     end
 end
+
+Base.@propagate_inbounds function fd_operator_shmem(
+    space,
+    ::Val{Nvt},
+    op::Operators.GradientC2F,
+    args...,
+) where {Nvt}
+    # allocate temp output
+    RT = return_eltype(op, args...)
+    u = CUDA.CuStaticSharedArray(RT, (Nvt,)) # cell centers
+    lb = CUDA.CuStaticSharedArray(RT, (1,)) # left boundary
+    rb = CUDA.CuStaticSharedArray(RT, (1,)) # right boundary
+    return (u, lb, rb)
+end
+
+Base.@propagate_inbounds function fd_operator_fill_shmem_interior!(
+    op::Operators.GradientC2F,
+    (u, lb, rb),
+    loc, # can be any location
+    space,
+    idx::Integer,
+    hidx,
+    arg,
+)
+    @inbounds begin
+        vt = threadIdx().x
+        cov3 = Geometry.Covariant3Vector(1)
+        u[vt] = cov3 ⊗ Operators.getidx(space, arg, loc, idx, hidx)
+    end
+    return nothing
+end
+
+Base.@propagate_inbounds function fd_operator_fill_shmem_left_boundary!(
+    op::Operators.GradientC2F,
+    bc::Operators.SetValue,
+    (u, lb, rb),
+    loc,
+    space,
+    idx::Integer,
+    hidx,
+    arg,
+)
+    idx == Operators.left_center_boundary_idx(space) ||
+        error("Incorrect left idx")
+    @inbounds begin
+        vt = threadIdx().x
+        cov3 = Geometry.Covariant3Vector(1)
+        u[vt] = cov3 ⊗ Operators.getidx(space, arg, loc, idx, hidx)
+        lb[1] = cov3 ⊗ Operators.getidx(space, bc.val, loc, nothing, hidx)
+    end
+    return nothing
+end
+
+Base.@propagate_inbounds function fd_operator_fill_shmem_right_boundary!(
+    op::Operators.GradientC2F,
+    bc::Operators.SetValue,
+    (u, lb, rb),
+    loc,
+    space,
+    idx::Integer,
+    hidx,
+    arg,
+)
+    # The right boundary is called at `idx + 1`, so we need to subtract 1 from idx (shmem is loaded at vt+1)
+    idx == Operators.right_center_boundary_idx(space) ||
+        error("Incorrect right idx")
+    @inbounds begin
+        vt = threadIdx().x
+        cov3 = Geometry.Covariant3Vector(1)
+        u[vt] = cov3 ⊗ Operators.getidx(space, arg, loc, idx, hidx)
+        rb[1] = cov3 ⊗ Operators.getidx(space, bc.val, loc, nothing, hidx)
+    end
+    return nothing
+end
+
+Base.@propagate_inbounds function fd_operator_evaluate(
+    op::Operators.GradientC2F,
+    (u, lb, rb),
+    loc,
+    space,
+    idx::PlusHalf,
+    hidx,
+    args...,
+)
+    @inbounds begin
+        vt = threadIdx().x
+        # @assert idx.i == vt-1 # assertion passes, but commented to remove potential thrown exception in llvm output
+        if idx == Operators.right_face_boundary_idx(space)
+            u₋ = 2 * u[vt - 1]   # corresponds to idx - half
+            u₊ = 2 * rb[1] # corresponds to idx + half
+        elseif idx == Operators.left_face_boundary_idx(space)
+            u₋ = 2 * lb[1]   # corresponds to idx - half
+            u₊ = 2 * u[vt] # corresponds to idx + half
+        else
+            u₋ = u[vt - 1]   # corresponds to idx - half
+            u₊ = u[vt] # corresponds to idx + half
+        end
+        return u₊ ⊟ u₋
+    end
+end
diff --git a/ext/cuda/operators_fd_shmem_common.jl b/ext/cuda/operators_fd_shmem_common.jl
@@ -192,10 +192,11 @@ get_arg_space(
     args::Tuple,
 ) = axes(args[1])
 
-get_cent_idx(idx::Integer) = idx
-get_face_idx(idx::PlusHalf) = idx
-get_cent_idx(idx::PlusHalf) = idx + half
-get_face_idx(idx::Integer) = idx - half
+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!(
@@ -301,21 +302,69 @@ Base.@propagate_inbounds function fd_resolve_shmem!(
         else # this else should never run
         end
     else  # populate shmem on centers
-#! format: off
-        # TODO: this needs exercised
-        case1 = IP || get_cent_idx(bc_lw) ≤ ᶜidx ≤ get_cent_idx(bc_rw) + 1 # interior
-        case2 = get_cent_idx(bc_li) < ᶜidx < get_cent_idx(bc_lw) && Operators.has_boundary(op, lloc) # left
-        case3 = get_cent_idx(bc_rw) < ᶜidx < get_cent_idx(bc_ri) && Operators.has_boundary(op, rloc) # right
-        case4 = get_cent_idx(bc_li) < ᶜidx < get_cent_idx(bc_lw) && !Operators.has_boundary(op, lloc) # left
-        case5 = get_cent_idx(bc_rw) < ᶜidx < get_cent_idx(bc_ri) && !Operators.has_boundary(op, rloc) # right
-        if case1; fd_operator_fill_shmem_interior!(sbc.op,sbc.work,iloc,space,ᶜidx,hidx,sbc.args...)
-        elseif case2; fd_operator_fill_shmem_left_boundary!(sbc.op,Operators.get_boundary(op, lloc),sbc.work,lloc,space,ᶜidx,hidx,sbc.args...)
-        elseif case3; fd_operator_fill_shmem_right_boundary!(sbc.op,Operators.get_boundary(op, rloc),sbc.work,rloc,space,ᶜidx,hidx,sbc.args...)
-        elseif case4; fd_operator_fill_shmem_interior!(sbc.op,sbc.work,lloc,space,ᶜidx,hidx,sbc.args...)
-        elseif case5; fd_operator_fill_shmem_interior!(sbc.op,sbc.work,rloc,space,ᶜidx,hidx,sbc.args...)
-        else # this else should never run
+        if IP || get_cent_idx(bc_lw) ≤ ᶜidx < get_cent_idx(bc_rw) # interior
+            fd_operator_fill_shmem_interior!(
+                sbc.op,
+                sbc.work,
+                iloc,
+                space,
+                ᶜidx,
+                hidx,
+                sbc.args...,
+            )
+        elseif get_cent_idx(bc_li) ≤ ᶜidx < get_cent_idx(bc_lw) &&
+               Operators.has_boundary(op, lloc) # left
+            fd_operator_fill_shmem_left_boundary!(
+                sbc.op,
+                Operators.get_boundary(op, lloc),
+                sbc.work,
+                lloc,
+                space,
+                ᶜidx,
+                hidx,
+                sbc.args...,
+            )
+        elseif get_cent_idx(bc_rw) ≤ ᶜidx < get_cent_idx(bc_ri) &&
+               Operators.has_boundary(op, rloc) # right
+            fd_operator_fill_shmem_right_boundary!(
+                sbc.op,
+                Operators.get_boundary(op, rloc),
+                sbc.work,
+                rloc,
+                space,
+                ᶜidx,
+                hidx,
+                sbc.args...,
+            )
+        elseif get_cent_idx(bc_li) < ᶜidx < get_cent_idx(bc_lw) &&
+               !Operators.has_boundary(op, lloc) # left
+            fd_operator_fill_shmem_interior!(
+                sbc.op,
+                sbc.work,
+                lloc,
+                space,
+                ᶜidx,
+                hidx,
+                sbc.args...,
+            )
+        elseif get_cent_idx(bc_rw) < ᶜidx < get_cent_idx(bc_ri) &&
+               !Operators.has_boundary(op, rloc) # right
+            fd_operator_fill_shmem_interior!(
+                sbc.op,
+                sbc.work,
+                rloc,
+                space,
+                ᶜidx,
+                hidx,
+                sbc.args...,
+            )
+        else # this should only ever be exercised at Spaces.nlevels(ᶜspace)+1
+            # We don't have or need to fill shmem at `Spaces.nlevels
+            # (ᶜspace)+1`, but threads may have this ᶜidx because they may be
+            # filling shmem for an operator whose shmem exists on the face
+            # space, which extends beyond the center space.
+            # @assert ᶜidx == Spaces.nlevels(ᶜspace) + 1 # assertion passes, but commented to remove potential thrown exception in llvm output
         end
-#! format: on
     end
     return nothing
 end
diff --git a/ext/cuda/operators_fd_shmem_is_supported.jl b/ext/cuda/operators_fd_shmem_is_supported.jl
@@ -149,6 +149,8 @@ end
 ) = 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(
@@ -162,3 +164,18 @@ end
     all(values(bcs)) do bc
         all(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
+        all(supported_bc -> bc isa supported_bc, (Operators.SetValue,))
+    end
diff --git a/test/Operators/finitedifference/unit_fd_ops_shared_memory.jl b/test/Operators/finitedifference/unit_fd_ops_shared_memory.jl
@@ -36,6 +36,13 @@ end
     @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
 
 #! format: off
@@ -66,6 +73,7 @@ end
     @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.ᶜ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)
 end
 
 @testset "Correctness plane" begin
diff --git a/test/Operators/finitedifference/utils_fd_ops_shared_memory.jl b/test/Operators/finitedifference/utils_fd_ops_shared_memory.jl
@@ -49,6 +49,7 @@ function kernels!(fields)
     (; ᶜout1, ᶜout2, ᶜout3, ᶜout4, ᶜout5, ᶜout6, ᶜout7, ᶜout8, ᶜout9) = fields
     (; ᶜout10) = fields
     (; ᶠout1_contra, ᶠout2_contra) = fields
+    (; ᶠout3_cov) = fields
     (; w_cov) = fields
     (; ᶜout_uₕ, ᶜuₕ) = fields
     FT = Spaces.undertype(axes(ϕ))
@@ -125,6 +126,12 @@ function kernels!(fields)
     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(ϕ)
+
     return nothing
 end;
 
@@ -138,12 +145,22 @@ function get_fields(space::Operators.AllFaceFiniteDifferenceSpace)
         i -> Fields.Field(Geometry.Contravariant3Vector{FT}, space),
         length(K_contra),
     )
+    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_contra, V_contra)..., zip(K_cov, V_cov)...)
+    nt = (;
+        zip(K, V)...,
+        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
