[imreddyTeja]

slice-loop statements that need any shared memory were
noinlined so that equally sized buffers of different statements share
memory. A call that is noinlined copies its arguments onto the stack of
every GPU thread. This PR first makes statements that use a small buffer get inlined. It also uses a placeholder grid to make the stack smaller for the cases where it does spill to the stack

script I used to benchmark (not added, unless reviewers think it should be):

function setup(zelem)
    context = ClimaComms.context(ClimaComms.device())
    ᶜspace =
        TU.CenterExtrudedFiniteDifferenceSpace(FT; zelem, helem = HELEM, Nq = 4, context)
    ᶠspace = Spaces.face_space(ᶜspace)
    ᶜz = Fields.coordinate_field(ᶜspace).z
    ᶠz = Fields.coordinate_field(ᶠspace).z
    c = (;
        ρ = (@. 1 + 0.1 * sin(ᶜz)),
        ρe = (@. 2 + 0.1 * cos(ᶜz)),
        ρq = (@. 0.01 + 0.001 * sin(ᶜz)),
        h = (@. 3 + 0.2 * sin(2ᶜz)),
        p = (@. 1e5 * exp(0 - ᶜz)),
        T = (@. 280 - 10 * ᶜz),
    )
    ᶠu³ = @. Geometry.WVector(0.1 * sin(3ᶠz))
    ᶠgΦ = @. Geometry.Covariant3Vector(9.8 + 0 * ᶠz)
    Yₜc = (; ρ = zero(c.ρ), ρe = zero(c.ρ), ρq = zero(c.ρ))
    Yₜf = (; u₃ = zero(ᶠgΦ))
    tmp = (; ᶠρ = zero(ᶠz), ᶜq = zero(ᶜz), ᶜθ = zero(ᶜz), ᶜΠ = zero(ᶜz))
    return (; c, ᶠu³, ᶠgΦ, Yₜc, Yₜf, tmp)
end

const ᶠinterp = Operators.InterpolateC2F(;
    bottom = Operators.Extrapolate(), top = Operators.Extrapolate())
const ᶜdiv = Operators.DivergenceF2C(;
    bottom = Operators.SetValue(Geometry.WVector(0.0)),
    top = Operators.SetValue(Geometry.WVector(0.0)))
const ᶠgrad = Operators.GradientC2F(;
    bottom = Operators.SetGradient(Geometry.WVector(0.0)),
    top = Operators.SetGradient(Geometry.WVector(0.0)))

# ---- body A: 3 advective divergences sharing ᶠρ (implicit_vertical_advection)
@inline function adv_body!(ρₜ, ρeₜ, ρqₜ, ρ, ρq, h, ᶠu³)
    ᶠρ = @. ᶠinterp(ρ)
    @. ρₜ -= ᶜdiv(ᶠρ * ᶠu³)
    @. ρeₜ += -(ᶜdiv(ᶠρ * ᶠu³ * ᶠinterp(h)))
    ᶜq = @. ρq / ρ
    @. ρqₜ += -(ᶜdiv(ᶠρ * ᶠu³ * ᶠinterp(ᶜq)))
    return nothing
end
adv_fused!(s) = DataLayouts.foreach_column(
    s.Yₜc.ρ, s.Yₜc.ρe, s.Yₜc.ρq, s.c.ρ, s.c.ρq, s.c.h, s.ᶠu³,
) do args...
    adv_body!(args...)
end
function adv_unfused!(s)
    (; c, ᶠu³, Yₜc, tmp) = s
    @. tmp.ᶠρ = ᶠinterp(c.ρ)
    @. Yₜc.ρ -= ᶜdiv(tmp.ᶠρ * ᶠu³)
    @. Yₜc.ρe += -(ᶜdiv(tmp.ᶠρ * ᶠu³ * ᶠinterp(c.h)))
    @. tmp.ᶜq = c.ρq / c.ρ
    @. Yₜc.ρq += -(ᶜdiv(tmp.ᶠρ * ᶠu³ * ᶠinterp(tmp.ᶜq)))
    return nothing
end

# ---- body B: u₃ tendency with pointwise temporaries (add_implicit_u₃_tendency!)
lazy(x) = x  # ClimaAtmos' lazy is a no-op marker here
Base.Broadcast.broadcasted(::typeof(lazy), x) = x
@inline θv(T, p) = T * (1e5 / p)^0.2857
@inline exner(p) = (p / 1e5)^0.2857
@inline function u3_body!(u₃ₜ, ᶠgΦ, p, T)
    ᶜθ = @. θv(T, p)
    ᶜΠ = @. lazy(exner(p))
    @. u₃ₜ -= ᶠgΦ + 1004 * ᶠinterp(ᶜθ - 300) * ᶠgrad(ᶜΠ)
    return nothing
end
u3_fused!(s) = DataLayouts.foreach_column(
    s.Yₜf.u₃, s.ᶠgΦ, s.c.p, s.c.T,
) do args...
    u3_body!(args...)
end
function u3_unfused!(s)
    (; c, ᶠgΦ, Yₜf, tmp) = s
    @. tmp.ᶜθ = θv(c.T, c.p)
    @. tmp.ᶜΠ = exner(c.p)
    @. Yₜf.u₃ -= ᶠgΦ + 1004 * ᶠinterp(tmp.ᶜθ - 300) * ᶠgrad(tmp.ᶜΠ)
    return nothing
end

# ---- body C: single FD statement (no fusion at all) for reference
single!(s) = (@. s.Yₜc.ρe = ᶜdiv(ᶠinterp(s.c.ρ) * s.ᶠu³ * ᶠinterp(s.c.h)); nothing)

claude generated results table:

Timings of perf/fused_column_fd.jl

GPU kernel time per call (µs, CUPTI) on one A100, HELEM=16, Nq=4, Float64, for main (96a1129) and this branch (9a00beb). adv is the three advective divergences of ρ, ρe_tot and ρq_tot that share ᶠρ; u₃ is the implicit u₃ tendency, both simplified from ClimaAtmos zs/column-fusion. Nv is the number of vertical elements (Nv + 1 faces). Unfused times are the same on both branches (within 1%), so they are listed once. Fused results match unfused ones exactly at every Nv on both branches.

Nv	adv unfused	adv fused (main)	adv fused (branch)	u₃ unfused	u₃ fused (main)	u₃ fused (branch)
25	196	103	103	65	58	58
31	231	112	112	74	62	62
32	259	12,087	468	95	6,334	72
63	444	18,216	204	140	6,379	112
64	470	17,593	409	170	12,595	152
79	576	35,565	298	180	12,404	192
95	663	35,608	322	210	12,618	191

With columns of more than 32 faces, the fused bodies on main are 37-74x slower than unfused. On this branch they are faster than unfused, except adv at Nv = 32 (1.8x slower) and u₃ at Nv = 79 (7% slower).

• Code follows the style guidelines OR N/A.
• Unit tests are included OR N/A.
• Code is exercised in an integration test OR N/A.
• Documentation has been added/updated OR N/A.