Generalize shmem size

Author: charleskawczynski

ext/cuda/operators_fd_shmem.jl

@@ -6,15 +6,15 @@ import ClimaCore.RecursiveApply: ⊟, ⊞
 
 Base.@propagate_inbounds function fd_operator_shmem(
     space,
-    ::Val{Nvt},
+    ::Val{params},
     op::Operators.DivergenceF2C,
     args...,
-) where {Nvt}
+) where {params}
     # allocate temp output
     RT = return_eltype(op, args...)
-    Ju³ = CUDA.CuStaticSharedArray(RT, (Nvt,))
-    lJu³ = CUDA.CuStaticSharedArray(RT, (1,))
-    rJu³ = CUDA.CuStaticSharedArray(RT, (1,))
+    Ju³ = CUDA.CuStaticSharedArray(RT, shmem_size(Val(params)))
+    lJu³ = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params)))
+    rJu³ = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params)))
     return (Ju³, lJu³, rJu³)
 end
 
@@ -109,15 +109,15 @@ end
 
 Base.@propagate_inbounds function fd_operator_shmem(
     space,
-    ::Val{Nvt},
+    ::Val{params},
     op::Operators.GradientC2F,
     args...,
-) where {Nvt}
+) where {params}
     # 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
+    u = CUDA.CuStaticSharedArray(RT, shmem_size(Val(params))) # cell centers
+    lb = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params))) # left boundary
+    rb = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params))) # right boundary
     return (u, lb, rb)
 end
 
@@ -202,15 +202,15 @@ end
 
 Base.@propagate_inbounds function fd_operator_shmem(
     space,
-    ::Val{Nvt},
+    ::Val{params},
     op::Operators.InterpolateC2F,
     args...,
-) where {Nvt}
+) where {params}
     # 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
+    u = CUDA.CuStaticSharedArray(RT, shmem_size(Val(params))) # cell centers
+    lb = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params))) # left boundary
+    rb = CUDA.CuStaticSharedArray(RT, boundary_shmem_size(Val(params))) # right boundary
     return (u, lb, rb)
 end
 

ext/cuda/operators_fd_shmem_common.jl

@@ -209,19 +209,23 @@ Base.@propagate_inbounds function getidx(
 end
 
 """
-    fd_allocate_shmem(Val(Nvt), b)
+    fd_allocate_shmem(Val(params), b)
 
 Create a new broadcasted object with necessary share memory allocated,
-using `Nvt` nodal points per block.
+using `params` nodal points per block.
 """
-@inline function fd_allocate_shmem(::Val{Nvt}, obj) where {Nvt}
+@inline function fd_allocate_shmem(::Val{params}, obj) where {params}
     obj
 end
 @inline function fd_allocate_shmem(
-    ::Val{Nvt},
+    ::Val{params},
     bc::Broadcasted{Style},
-) where {Nvt, Style}
-    Broadcasted{Style}(bc.f, _fd_allocate_shmem(Val(Nvt), bc.args...), bc.axes)
+) where {params, Style}
+    Broadcasted{Style}(
+        bc.f,
+        _fd_allocate_shmem(Val(params), bc.args...),
+        bc.axes,
+    )
 end
 
 ######### MatrixFields
@@ -236,24 +240,27 @@ end
 #########
 
 @inline function fd_allocate_shmem(
-    ::Val{Nvt},
+    ::Val{params},
     sbc::StencilBroadcasted{Style},
-) where {Nvt, Style}
-    args = _fd_allocate_shmem(Val(Nvt), sbc.args...)
+) where {params, Style}
+    args = _fd_allocate_shmem(Val(params), sbc.args...)
     work = if Operators.fd_shmem_is_supported(sbc)
-        fd_operator_shmem(sbc.axes, Val(Nvt), sbc.op, args...)
+        fd_operator_shmem(sbc.axes, Val(params), sbc.op, args...)
     else
         nothing
     end
     StencilBroadcasted{Style}(sbc.op, args, sbc.axes, work)
 end
 
-@inline _fd_allocate_shmem(::Val{Nvt}) where {Nvt} = ()
-@inline _fd_allocate_shmem(::Val{Nvt}, arg, xargs...) where {Nvt} = (
-    fd_allocate_shmem(Val(Nvt), arg),
-    _fd_allocate_shmem(Val(Nvt), xargs...)...,
+@inline _fd_allocate_shmem(::Val{params}) where {params} = ()
+@inline _fd_allocate_shmem(::Val{params}, arg, xargs...) where {params} = (
+    fd_allocate_shmem(Val(params), arg),
+    _fd_allocate_shmem(Val(params), xargs...)...,
 )
 
+shmem_size(::Val{params}) where {params} = (params.Nvt,)
+boundary_shmem_size(::Val{params}) where {params} = (1,)
+
 """
     fd_shmem_needed_per_column(::Base.Broadcast.Broadcasted)
     fd_shmem_needed_per_column(::StencilBroadcasted)

ext/cuda/operators_finite_difference.jl

@@ -56,6 +56,7 @@ function Base.copyto!(
        mask isa NoMask &&
        enough_shmem &&
        Operators.use_fd_shmem()
+        shmem_params = (; Nv = n_face_levels)
         p = fd_shmem_stencil_partition(us, n_face_levels)
         args = (
             strip_space(out, space),
@@ -64,7 +65,7 @@ function Base.copyto!(
             bounds,
             us,
             mask,
-            Val(p.Nvthreads),
+            Val(shmem_params),
         )
         auto_launch!(
             copyto_stencil_kernel_shmem!,
@@ -153,8 +154,8 @@ function copyto_stencil_kernel_shmem!(
     bds,
     us,
     mask,
-    ::Val{Nvt},
-) where {Nvt}
+    ::Val{shmem_params},
+) where {shmem_params}
     @inbounds begin
         out_fv = Fields.field_values(out)
         us = DataLayouts.UniversalSize(out_fv)
@@ -165,7 +166,7 @@ function copyto_stencil_kernel_shmem!(
             hidx = (i, j, h)
             idx = v - 1 + li
             bc = Operators.reconstruct_placeholder_broadcasted(space, bc′)
-            bc_shmem = fd_allocate_shmem(Val(Nvt), bc) # allocates shmem
+            bc_shmem = fd_allocate_shmem(Val(shmem_params), bc) # allocates shmem
 
             fd_resolve_shmem!(bc_shmem, idx, hidx, bds) # recursively fills shmem
             CUDA.sync_threads()