Disable GradientC2F shmem, add tests (#2278)

Co-authored-by: Charlie Kawczynski <charliek@clima.gps.caltech.edu>

Author: charleskawczynski

ext/cuda/operators_fd_shmem_is_supported.jl

@@ -175,7 +175,7 @@ end
 @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
+) = false
+# all(values(bcs)) do bc
+#     all(supported_bc -> bc isa supported_bc, (Operators.SetValue,))
+# end

test/Operators/finitedifference/unit_fd_ops_shared_memory.jl

@@ -40,9 +40,9 @@ end
         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)
+    # bc = @. lazy(ᶠgrad(c))
+    # @test Operators.any_fd_shmem_supported(bc)
+    # @test Operators.fd_shmem_is_supported(bc)
 end
 
 #! format: off
@@ -72,6 +72,8 @@ end
     @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.ᶜ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
@@ -102,6 +104,8 @@ end
     @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.ᶜout_uₕ, fields.ᶜout_uₕ); @test !is_trivial(fields_cpu.ᶜout_uₕ)
 end
 
@@ -130,6 +134,8 @@ end
     @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.ᶜout_uₕ, fields.ᶜout_uₕ); @test !is_trivial(fields_cpu.ᶜout_uₕ)
 
 end

test/Operators/finitedifference/utils_fd_ops_shared_memory.jl

@@ -47,7 +47,8 @@ get_space_column(dev, FT; z_elem = 10) = ColumnSpace(
 function kernels!(fields)
     (; f, ρ, ϕ) = fields
     (; ᶜout1, ᶜout2, ᶜout3, ᶜout4, ᶜout5, ᶜout6, ᶜout7, ᶜout8, ᶜout9) = fields
-    (; ᶜout10) = fields
+    (; ᶜout10, ᶜout11, ᶜout12) = fields
+    (; ᶠout1) = fields
     (; ᶠout1_contra, ᶠout2_contra) = fields
     (; ᶠout3_cov) = fields
     (; w_cov) = fields
@@ -132,6 +133,27 @@ function kernels!(fields)
     )
     @. ᶠout3_cov = ᶠgrad(ϕ)
 
+    ᶠ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;
 
@@ -145,6 +167,8 @@ function get_fields(space::Operators.AllFaceFiniteDifferenceSpace)
         i -> Fields.Field(Geometry.Contravariant3Vector{FT}, space),
         length(K_contra),
     )
+    K_scalar = (ntuple(i -> Symbol("ᶠout$(i)"), 1)...,)
+    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),
@@ -157,6 +181,7 @@ function get_fields(space::Operators.AllFaceFiniteDifferenceSpace)
     )
     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)...,
@@ -168,7 +193,7 @@ end
 
 function get_fields(space::Operators.AllCenterFiniteDifferenceSpace)
     FT = Spaces.undertype(space)
-    K = (ntuple(i -> Symbol("ᶜout$i"), 10)..., :ρ, :ϕ)
+    K = (ntuple(i -> Symbol("ᶜout$i"), 12)..., :ρ, :ϕ)
     V = ntuple(i -> Fields.zeros(space), length(K))
     nt = (;
         zip(K, V)...,