JuliaGPU / JuliaGPU/KernelAbstractions.jl

Performance discrepancy against CUDA

Open
#598 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
523
Forks
88
Avg merge
1d 11h
Merged PRs (30d)
25

Description

Reproducer:

CUDA.jl:

using CUDA
using Adapt
CUDA.allowscalar(false)
# using KernelAbstractions
is_valid_index(meta, ui) =
    1 ≤ ui[1] ≤ params(meta)[4] &&
    1 ≤ ui[2] ≤ params(meta)[1] &&
    1 ≤ ui[3] ≤ params(meta)[2] &&
    1 ≤ ui[4] ≤ params(meta)[5]
@inline function universal_index_cuda(meta, src)
    (v,) = CUDA.threadIdx()
    (h, ij) = CUDA.blockIdx()
    (Ni, Nj, Nv, _, Nh) = params(meta)
    Ni * Nj < ij && return CartesianIndex((-1, -1, 1, -1, -1))
    @inbounds (i, j) = CartesianIndices((Ni, Nj))[ij].I
    return CartesianIndex((v, i, j, 1, h))
end
function my_memcpy_cuda!(dest, src, meta)
    (Ni, Nj, _, Nv, Nh) = params(meta)
    kernel = CUDA.@cuda(
        always_inline = true,
        launch = false,
        my_memcpy_kernel!(dest, src, meta)
    )
    threads = (Nv, )
    blocks = (Nh, Ni * Nj)
    kernel(dest, src, meta; threads, blocks)
    return nothing
end
function my_memcpy_kernel!(dest, src, meta)
    ui = universal_index_cuda(meta, src)
    (Ni, Nj, _, Nv, Nh) = params(meta)
    src_arr = CUDA.CuStaticSharedArray(eltype(src), (Nv,))
    is_valid_index(meta, ui) && (src_arr[ui[1]] = src[ui])
    CUDA.sync_threads()
    if is_valid_index(meta, ui)
        src_shmem = reconstruct_src(src, src_arr)
        dest[ui] = src_shmem[ui[1]]
    end
    return nothing
end
reconstruct_src(src, src_arr) = src_arr
params(p::Tuple) = map(unval, p)
unval(::Val{n}) where {n} = n
unval(n) = n
Ni = 4
Nj = 4
Nv = 64
Nh = 5400
meta = (Val(Ni), Val(Nj), 1, Val(Nv), Val(Nh))
_src1 = CUDA.CuArray(rand(Float32, Nv, Ni, Nj, 1, Nh));
_src2 = similar(_src1); @. _src2 = _src1
_dest1 = similar(_src1);
_dest2 = similar(_src1);
my_memcpy_cuda!(_dest1, _src1, meta)
@. _dest2 = _src2
using Test
@test all(Array(_dest2) .== Array(_dest1))

CUDA.@profile begin
    my_memcpy_cuda!(_dest1, _src1, meta)
    my_memcpy_cuda!(_dest1, _src1, meta)
    my_memcpy_cuda!(_dest1, _src1, meta)
end

KA.jl

using CUDA
using KernelAbstractions
using Adapt
CUDA.allowscalar(false)
using CUDA.CUDAKernels
# using KernelAbstractions
@inline function universal_index_KA(
    meta,
    ilc, # @index(Local, Cartesian)
    igc, # @index(Group, Cartesian)
    gs, # @groupsize()
)
    @inbounds begin
        v = ilc[1]
        (_, h, ij) = igc.I
        (Ni, Nj, _, Nv, Nh) = params(meta)
        if Ni * Nj < ij
            return CartesianIndex((-1, -1, 1, -1, -1))
        else
            @inbounds (i, j) = CartesianIndices((Ni, Nj))[ij].I
            return CartesianIndex((v, i, j, 1, h))
        end
    end
end

is_valid_index(meta, ui) =
    1 ≤ ui[1] ≤ params(meta)[4] &&
    1 ≤ ui[2] ≤ params(meta)[1] &&
    1 ≤ ui[3] ≤ params(meta)[2] &&
    1 ≤ ui[4] ≤ params(meta)[5]
function my_memcpy_KA!(dest, src, meta)
    (Ni, Nj, _, Nv, Nh) = params(meta)
    backend = CUDABackend()
    group = (Nv, )
    ndrange = (Nv, Nh, Ni * Nj)
    kernel = my_memcpy_kernel!(backend, group)
    kernel(dest, src, meta, ndrange = ndrange)
    return nothing
end
@kernel function my_memcpy_kernel!(dest, src, meta)
    ilc = @index(Local, Cartesian); igc = @index(Group, Cartesian); gs = @groupsize()
    ui = universal_index_KA(meta, ilc, igc, gs)
    (Ni, Nj, _, Nv, Nh) = @uniform params(meta)
    src_arr = @localmem eltype(src) (Nv,)
    is_valid_index(meta, ui) && (src_arr[ui[1]] = src[ui])
    @synchronize
    ilc = @index(Local, Cartesian); igc = @index(Group, Cartesian); gs = @groupsize()
    ui = universal_index_KA(meta, ilc, igc, gs)
    if is_valid_index(meta, ui)
        src_shmem = reconstruct_src(src, src_arr)
        dest[ui] = src_shmem[ui[1]]
    end
end
reconstruct_src(src, src_arr) = src_arr
params(p::Tuple) = map(unval, p)
unval(::Val{n}) where {n} = n
unval(n) = n
Ni = 4
Nj = 4
Nv = 64
Nh = 5400
meta = (Val(Ni), Val(Nj), 1, Val(Nv), Val(Nh))
_src1 = CUDA.CuArray(rand(Float32, Nv, Ni, Nj, 1, Nh));
_src2 = similar(_src1); @. _src2 = _src1
_dest1 = similar(_src1);
_dest2 = similar(_src1);
my_memcpy_KA!(_dest1, _src1, meta)
@. _dest2 = _src2
using Test
@test all(Array(_dest2) .== Array(_dest1))

CUDA.@profile begin
    my_memcpy_KA!(_dest1, _src1, meta)
    my_memcpy_KA!(_dest1, _src1, meta)
    my_memcpy_KA!(_dest1, _src1, meta)
end

Output:

julia> using Revise; include("../ka_reproducer.jl")
Profiler ran for 429.39 µs, capturing 68 events.

Host-side activity: calling CUDA APIs took 77.96 µs (18.16% of the trace)
┌──────────┬────────────┬───────┬────────────────────────────────────┬────────────────┐
│ Time (%) │ Total time │ Calls │ Time distribution                  │ Name           │
├──────────┼────────────┼───────┼────────────────────────────────────┼────────────────┤
│   17.38% │   74.63 µs │     3 │  24.88 µs ± 33.83  (  3.81 ‥ 63.9) │ cuLaunchKernel │
└──────────┴────────────┴───────┴────────────────────────────────────┴────────────────┘

Device-side activity: GPU was busy for 304.94 µs (71.02% of the trace)
┌──────────┬────────────┬───────┬──────────────────────────────────────┬─────────────────────────────────────────────────────────────────────
│ Time (%) │ Total time │ Calls │ Time distribution                    │ Name                                                               ⋯
├──────────┼────────────┼───────┼──────────────────────────────────────┼─────────────────────────────────────────────────────────────────────
│   71.02% │  304.94 µs │     3 │ 101.65 µs ± 0.36   (101.33 ‥ 102.04) │ gpu_my_memcpy_kernel_(CompilerMetadata<DynamicSize, DynamicCheck,  ⋯
└──────────┴────────────┴───────┴──────────────────────────────────────┴─────────────────────────────────────────────────────────────────────
                                                                                                                             1 column omitted


julia> using Revise; include("../cuda_reproducer.jl")
Profiler ran for 341.42 µs, capturing 68 events.

Host-side activity: calling CUDA APIs took 70.1 µs (20.53% of the trace)
┌──────────┬────────────┬───────┬─────────────────────────────────────┬────────────────┐
│ Time (%) │ Total time │ Calls │ Time distribution                   │ Name           │
├──────────┼────────────┼───────┼─────────────────────────────────────┼────────────────┤
│   19.55% │   66.76 µs │     3 │  22.25 µs ± 28.88  (  4.05 ‥ 55.55) │ cuLaunchKernel │
└──────────┴────────────┴───────┴─────────────────────────────────────┴────────────────┘

Device-side activity: GPU was busy for 250.34 µs (73.32% of the trace)
┌──────────┬────────────┬───────┬─────────────────────────────────────┬──────────────────────────────────────────────────────────────────────
│ Time (%) │ Total time │ Calls │ Time distribution                   │ Name                                                                ⋯
├──────────┼────────────┼───────┼─────────────────────────────────────┼──────────────────────────────────────────────────────────────────────
│   73.32% │  250.34 µs │     3 │  83.45 µs ± 0.24   ( 83.21 ‥ 83.68) │ my_memcpy_kernel_(CuDeviceArray<Float32, 5, 1>, CuDeviceArray<Float ⋯
└──────────┴────────────┴───────┴─────────────────────────────────────┴──────────────────────────────────────────────────────────────────────

My real-world example is much more complicated, and incurs a 25x slowdown. I'm not sure why. Here is the real-world case:

(Toggle implicit_tendency! at the top of the script)

Click to see script 
ENV["CLIMACOMMS_DEVICE"] = "CUDA";
high_res = true;
# implicit_tendency!(Yₜ, Y, p, t) = implicit_tendency_KA!(Yₜ, Y, p, t)
implicit_tendency!(Yₜ, Y, p, t) = implicit_tendency_cuda!(Yₜ, Y, p, t)

using CUDA
import ClimaComms
ClimaComms.@import_required_backends
using ClimaCore.CommonSpaces
import ClimaAtmos as CA
import ClimaCore.Fields.StaticArrays: MArray
using LazyBroadcast: lazy
using LinearAlgebra: ×, dot, norm
import ClimaAtmos.Parameters as CAP
import Thermodynamics as TD
import SciMLBase
import ClimaCore.Grids
import ClimaCore
using KernelAbstractions
import KernelAbstractions as KA
import ClimaTimeSteppers as CTS
import ClimaCore.Geometry
import ClimaCore.MatrixFields: @name, ⋅
import ClimaCore.MatrixFields: DiagonalMatrixRow, BidiagonalMatrixRow
import LinearAlgebra: Adjoint
import LinearAlgebra: adjoint
import LinearAlgebra as LA
import ClimaCore: Operators, Topologies, DataLayouts
import ClimaCore.MatrixFields
import ClimaCore.Spaces
import ClimaCore.Fields
# import KernelAbstractions as KA
# using KernelAbstractions

# Unless the kernel here fills the shmem, we cannot use the shmem getidx.
# So, let's disable for now. Maybe we can write a simple shmem transformation
# of the broadcasted object + the state.
const cccuda_ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt);
# LazyBroadcast calls instantiate on intermediate broadcast expressions, so we
# should be determining AbstractStencilStyle locally.
# ClimaCore bug: stencil style should be locally determined at the top level, not using the recursive Operators.any_fd_shmem_supported(bc)
Operators.AbstractStencilStyle(bc, ::ClimaComms.CUDADevice) =
    cccuda_ext.CUDAColumnStencilStyle
Operators.fd_shmem_is_supported(bc::Base.Broadcast.Broadcasted) = false
ClimaCore.Operators.use_fd_shmem() = false

@info "Arch: $(ClimaComms.device())"
if get(ENV, "CLIMACOMMS_DEVICE", "CPU") == "CUDA"
    using CUDA
    using CUDA.CUDAKernels
    CUDA.allowscalar(false)
else
end

# allow on-device use of lazy broadcast objects
DataLayouts.parent_array_type(::Type{<:CUDA.CuDeviceArray{T, N, A} where {N}}) where {T, A} =
    CUDA.CuDeviceArray{T, N, A} where {N}

# allow on-device use of lazy broadcast objects
DataLayouts.promote_parent_array_type(
    ::Type{CUDA.CuDeviceArray{T1, N, B} where {N}},
    ::Type{CUDA.CuDeviceArray{T2, N, B} where {N}},
) where {T1, T2, B} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N}

# Ditch sizes (they're never actually used!)
DataLayouts.promote_parent_array_type(
    ::Type{MArray{S1, T1}},
    ::Type{MArray{S2, T2}},
) where {S1, T1, S2, T2} = MArray{S, promote_type(T1, T2)} where {S}
DataLayouts.promote_parent_array_type(
    ::Type{MArray{S1, T1} where {S1}},
    ::Type{MArray{S2, T2}},
) where {T1, S2, T2} = MArray{S, promote_type(T1, T2)} where {S}
DataLayouts.promote_parent_array_type(
    ::Type{MArray{S1, T1}},
    ::Type{MArray{S2, T2} where {S2}},
) where {S1, T1, T2} = MArray{S, promote_type(T1, T2)} where {S}

# allow on-device use of lazy broadcast objects with different type params
DataLayouts.promote_parent_array_type(
    ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}},
    ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}},
) where {T1, T2, B1, B2} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}

# allow on-device use of lazy broadcast objects with different type params
DataLayouts.promote_parent_array_type(
    ::Type{CUDA.CuDeviceArray{T1}},
    ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}},
) where {T1, T2, B2} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}

DataLayouts.promote_parent_array_type(
    ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}},
    ::Type{CUDA.CuDeviceArray{T2} where {N}},
) where {T1, T2, B1} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}

# Specialize to allow on-device call of `device` for `DeviceExtrudedFiniteDifferenceGrid`
ClimaComms.device(grid::Grids.DeviceExtrudedFiniteDifferenceGrid) =
    ClimaComms.device(Grids.vertical_topology(grid))

# The existing implementation limits our ability to apply the same expressions from within kernels
ClimaComms.device(topology::Topologies.DeviceIntervalTopology) = ClimaComms.CUDADevice()

Fields.error_mismatched_spaces(::Type, ::Type) = nothing # causes unsupported dynamic function invocation

import ClimaAtmos: C1, C2, C12, C3, C123, CT1, CT2, CT12, CT3, CT123, UVW
import ClimaAtmos:
    divₕ, wdivₕ, gradₕ, wgradₕ, curlₕ, wcurlₕ, ᶜinterp, ᶜdivᵥ, ᶜgradᵥ
import ClimaAtmos: ᶠinterp, ᶠgradᵥ, ᶠcurlᵥ, ᶜinterp_matrix, ᶠgradᵥ_matrix
import ClimaAtmos: ᶜadvdivᵥ, ᶜadvdivᵥ_matrix, ᶠwinterp, ᶠinterp_matrix

Fields.local_geometry_field(bc::Base.Broadcast.Broadcasted) =
    Fields.local_geometry_field(axes(bc))

ᶜtendencies(ρ, uₕ, ρe_tot) = (; ρ, uₕ, ρe_tot)
ᶠtendencies(u₃) = (; u₃)

@inline is_valid_index(us, ui) = 1 ≤ ui[4] ≤ DataLayouts.get_Nv(us)

function implicit_tendency_bc!(Yₜ, Y, p, t)
    Yₜ .= zero(eltype(Yₜ))
    set_precomputed_quantities!(Y, p, t)
    (; rayleigh_sponge, params, dt) = p
    (; ᶜh_tot, ᶠu³, ᶜp) = p.precomputed
    ᶜJ = Fields.local_geometry_field(Y.c).J
    ᶜz = Fields.coordinate_field(Y.c).z
    ᶠz = Fields.coordinate_field(Y.f).z
    grav = FT(CAP.grav(params))
    zmax = CA.z_max(axes(Y.f))

    @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠwinterp(ᶜJ, Y.c.ρ) * ᶠu³)
    # Central advection of active tracers (e_tot and q_tot)
    Yₜ.c.ρe_tot .+= CA.vertical_transport(Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none))
    @. Yₜ.f.u₃ -= ᶠgradᵥ(ᶜp) / ᶠinterp(Y.c.ρ) + ᶠgradᵥ(Φ(grav, ᶜz))

    @. Yₜ.f.u₃ -= CA.β_rayleigh_w(rayleigh_sponge, ᶠz, zmax) * Y.f.u₃
    return nothing
end

function implicit_tendency_cuda!(Yₜ, Y, p, t)
    ᶜspace = axes(Y.c)
    ᶠspace = Spaces.face_space(ᶜspace)
    ᶠNv = Spaces.nlevels(ᶠspace)
    ᶜcf = Fields.coordinate_field(ᶜspace)
    us = DataLayouts.UniversalSize(Fields.field_values(ᶜcf))
    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
    nitems = Ni * Nj * 1 * ᶠNv * Nh
    ᶜYₜ = Yₜ.c
    ᶠYₜ = Yₜ.f
    ᶜY = Y.c
    ᶠY = Y.f
    (; rayleigh_sponge, params, dt) = p
    p_kernel = (; rayleigh_sponge, params, dt)
    zmax = Spaces.z_max(axes(ᶠY)) # DeviceIntervalTopology does not have mesh, and therefore cannot compute zmax
    
    kernel = CUDA.@cuda(
        always_inline = true,
        launch = false,
        implicit_tendency_kernel_cuda!(ᶜYₜ, ᶠYₜ, ᶜY, ᶠY, p_kernel, t, zmax)
    )
    threads = (ᶠNv, )
    blocks = (Nh, Ni * Nj)
    kernel(ᶜYₜ, ᶠYₜ, ᶜY, ᶠY, p_kernel, t, zmax; threads, blocks)
end

function implicit_tendency_kernel_cuda!(ᶜYₜ, ᶠYₜ, _ᶜY, _ᶠY, p, t, zmax)
    ᶜY_fv = Fields.field_values(_ᶜY)
    ᶠY_fv = Fields.field_values(_ᶠY)
    FT = Spaces.undertype(axes(_ᶜY))
    ᶜNv = Spaces.nlevels(axes(_ᶜY))
    ᶠNv = Spaces.nlevels(axes(_ᶠY))
    ᶜus = DataLayouts.UniversalSize(ᶜY_fv)
    ᶠus = DataLayouts.UniversalSize(ᶠY_fv)
    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(ᶠus)
    ᶜTS = DataLayouts.typesize(FT, eltype(ᶜY_fv))
    ᶠTS = DataLayouts.typesize(FT, eltype(ᶠY_fv))
    ᶜlg = Spaces.local_geometry_data(axes(_ᶜY))
    ᶠlg = Spaces.local_geometry_data(axes(_ᶠY))
    ᶜTS_lg = DataLayouts.typesize(FT, eltype(ᶜlg))

    ᶜui = universal_index_cuda(ᶜus)
    ᶠui = universal_index_cuda(ᶠus)
    # ilc = @index(Local, Cartesian)
    # igc = @index(Group, Cartesian)
    # gs = @groupsize()
    # ᶜui = universal_index_KA(ᶠus, ilc, igc, gs)
    # ᶠui = universal_index_KA(ᶠus, ilc, igc, gs)

    ᶜY_arr = CUDA.CuStaticSharedArray(FT, (ᶜNv, ᶜTS)) # ᶜY_arr = @localmem FT (ᶜNv, ᶜTS)
    ᶠY_arr = CUDA.CuStaticSharedArray(FT, (ᶠNv, ᶠTS)) # ᶠY_arr = @localmem FT (ᶠNv, ᶠTS)
    ᶜdata_col = rebuild_column(ᶜY_fv, ᶜY_arr)
    ᶠdata_col = rebuild_column(ᶠY_fv, ᶠY_arr)
    
    ᶜlg_arr = CUDA.CuStaticSharedArray(FT, (ᶜNv, ᶜTS_lg)) # ᶜlg_arr = @localmem FT (ᶜNv, ᶜTS_lg)
    ᶠlg_arr = CUDA.CuStaticSharedArray(FT, (ᶠNv, ᶜTS_lg)) # ᶠlg_arr = @localmem FT (ᶠNv, ᶜTS_lg)

    (ᶜspace_col, ᶠspace_col) = column_spaces(_ᶜY, _ᶠY, ᶠui, ᶜlg_arr, ᶠlg_arr)

    is_valid_index(ᶜus, ᶜui) && (ᶜdata_col[ᶜui] = ᶜY_fv[ᶜui])
    is_valid_index(ᶠus, ᶠui) && (ᶠdata_col[ᶠui] = ᶠY_fv[ᶠui])

    ᶜlg_col = Spaces.local_geometry_data(ᶜspace_col)
    ᶠlg_col = Spaces.local_geometry_data(ᶠspace_col)
    # is_valid_index(ᶜus, ᶜui) && (ᶜlg_col[ᶜui] = ᶜlg[ᶜui])
    # is_valid_index(ᶠus, ᶠui) && (ᶠlg_col[ᶠui] = ᶠlg[ᶠui])

    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.coordinates.z[ᶜui] = ᶜlg.coordinates.z[ᶜui]) # needed
    is_valid_index(ᶠus, ᶠui) && (ᶠlg_col.coordinates.z[ᶠui] = ᶠlg.coordinates.z[ᶠui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.J[ᶜui] = ᶜlg.J[ᶜui]) # needed
    is_valid_index(ᶠus, ᶠui) && (ᶠlg_col.J[ᶠui] = ᶠlg.J[ᶠui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.invJ[ᶜui] = ᶜlg.invJ[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:1[ᶜui] = ᶜlg.gⁱʲ.components.data.:1[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:2[ᶜui] = ᶜlg.gⁱʲ.components.data.:2[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:3[ᶜui] = ᶜlg.gⁱʲ.components.data.:3[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:4[ᶜui] = ᶜlg.gⁱʲ.components.data.:4[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:5[ᶜui] = ᶜlg.gⁱʲ.components.data.:5[ᶜui]) # needed
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col.gⁱʲ.components.data.:6[ᶜui] = ᶜlg.gⁱʲ.components.data.:6[ᶜui]) # needed
    is_valid_index(ᶠus, ᶠui) && (ᶠlg_col.gⁱʲ.components.data.:9[ᶠui] = ᶠlg.gⁱʲ.components.data.:9[ᶠui]) # needed

    CUDA.sync_threads()

    # ilc = @index(Local, Cartesian)
    # igc = @index(Group, Cartesian)
    # gs = @groupsize()
    # ᶜui = universal_index_KA(ᶠus, ilc, igc, gs)
    # ᶠui = universal_index_KA(ᶠus, ilc, igc, gs)

    ᶜdata_col = rebuild_column(ᶜY_fv, ᶜY_arr)
    ᶠdata_col = rebuild_column(ᶠY_fv, ᶠY_arr)

    # (ᶜspace_col, ᶠspace_col) = column_spaces(_ᶜY, _ᶠY, ᶠui, ᶜlg_arr, ᶠlg_arr)

    if is_valid_index(ᶜus, ᶜui)
        (ᶜY, ᶠY) = column_states(_ᶜY, _ᶠY, ᶜdata_col, ᶠdata_col, ᶠui, ᶜspace_col, ᶠspace_col)
        ᶜbc = ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
        (ᶜidx, ᶜhidx) = operator_inds(axes(ᶜY), ᶜui)
        Fields.field_values(ᶜYₜ)[ᶜui] = Operators.getidx(axes(ᶜY), ᶜbc, ᶜidx, ᶜhidx)
        # ᶜYₜ[ᶜui] = ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t)[ᶜui] # might be possible?
    end
    if is_valid_index(ᶠus, ᶠui)
        (ᶜY, ᶠY) = column_states(_ᶜY, _ᶠY, ᶜdata_col, ᶠdata_col, ᶠui, ᶜspace_col, ᶠspace_col)
        ᶠbc = ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
        (ᶠidx, ᶠhidx) = operator_inds(axes(ᶠY), ᶠui)
        Fields.field_values(ᶠYₜ)[ᶠui] = Operators.getidx(axes(ᶠY), ᶠbc, ᶠidx, ᶠhidx)
        # ᶠYₜ[ᶠui] = ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t)[ᶠui] # might be possible?
    end
    return nothing
end

@inline function universal_index_cuda(us)
    (v,) = CUDA.threadIdx()
    (h, ij) = CUDA.blockIdx()
    (Ni, Nj, _, _, _) = DataLayouts.universal_size(us)
    Ni * Nj < ij && return CartesianIndex((-1, -1, 1, -1, -1))
    @inbounds (i, j) = CartesianIndices((Ni, Nj))[ij].I
    return CartesianIndex((i, j, 1, v, h))
end

@inline function universal_index_KA(
        us,
        ilc, # @index(Local, Cartesian)
        igc, # @index(Group, Cartesian)
        gs, # @groupsize()
    )
    @inbounds begin
        v = ilc[1]
        (_, h, ij) = igc.I
        # @print("igc = $(igc.I), ilc = $(ilc.I)\n")
        (Ni, Nj, _, _, _) = DataLayouts.universal_size(us)
        if Ni * Nj < ij
            ui = CartesianIndex((-1, -1, 1, -1, -1))
        else
            @inbounds (i, j) = CartesianIndices((Ni, Nj))[ij].I
            ui = CartesianIndex((i, j, 1, v, h))
        end
        return ui
    end
end

function implicit_tendency_KA!(@nospecialize(Yₜ), @nospecialize(Y), @nospecialize(p), @nospecialize(t))
    ᶜspace = axes(Y.c)
    ᶠspace = Spaces.face_space(ᶜspace)
    ᶠNv = Spaces.nlevels(ᶠspace)
    ᶜcf = Fields.coordinate_field(ᶜspace)
    us = DataLayouts.UniversalSize(Fields.field_values(ᶜcf))
    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
    nitems = Ni * Nj * 1 * ᶠNv * Nh
    ᶜYₜ = Yₜ.c
    ᶠYₜ = Yₜ.f
    ᶜY = Y.c
    ᶠY = Y.f
    (; rayleigh_sponge, params, dt) = p
    p_kernel = (; rayleigh_sponge, params, dt)
    zmax = Spaces.z_max(axes(ᶠY)) # DeviceIntervalTopology does not have mesh, and therefore cannot compute zmax

    backend = if ClimaComms.device(ᶜspace) isa ClimaComms.CUDADevice
        CUDABackend()
    else
        KA.CPU()
    end
    group = (ᶠNv, )
    kernel = implicit_tendency_kernel_KA!(backend, group)
    # ndrange = (Nh, 1, Ni * Nj)
    ndrange = (ᶠNv,  Nh, Ni * Nj)
    # @show ᶠNv, Ni, Nj, Nh, prod(ndrange) # (ᶠNv, Ni, Nj, Nh, prod(ndrange)) = (9, 2, 2, 24, 864)
    kernel(ᶜYₜ, ᶠYₜ, ᶜY, ᶠY, p_kernel, t, zmax, ndrange = ndrange)
end
@kernel function implicit_tendency_kernel_KA!(ᶜYₜ, ᶠYₜ, @Const(_ᶜY), @Const(_ᶠY), @Const(p), @Const(t), @Const(zmax))
    ᶜY_fv = @uniform Fields.field_values(_ᶜY)
    ᶠY_fv = @uniform Fields.field_values(_ᶠY)
    FT = @uniform Spaces.undertype(axes(_ᶜY))
    ᶜNv = @uniform Spaces.nlevels(axes(_ᶜY))
    ᶠNv = @uniform Spaces.nlevels(axes(_ᶠY))
    ᶜus = @uniform DataLayouts.UniversalSize(ᶜY_fv)
    ᶠus = @uniform DataLayouts.UniversalSize(ᶠY_fv)
    (Ni, Nj, _, _, Nh) = @uniform DataLayouts.universal_size(ᶠus)
    ᶜTS = @uniform DataLayouts.typesize(FT, eltype(ᶜY_fv))
    ᶠTS = @uniform DataLayouts.typesize(FT, eltype(ᶠY_fv))
    ᶜlg = @uniform Spaces.local_geometry_data(axes(_ᶜY))
    ᶠlg = @uniform Spaces.local_geometry_data(axes(_ᶠY))
    ᶜTS_lg = @uniform DataLayouts.typesize(FT, eltype(ᶜlg))

    ilc = @index(Local, Cartesian)
    igc = @index(Group, Cartesian)
    gs = @groupsize()
    ᶜui = universal_index_KA(ᶠus, ilc, igc, gs)
    ᶠui = universal_index_KA(ᶠus, ilc, igc, gs)

    # @print("ᶜui = $(ᶜui.I)\n")
    ᶜY_arr = @localmem FT (ᶜNv, ᶜTS)
    ᶠY_arr = @localmem FT (ᶠNv, ᶠTS)
    ᶜdata_col = rebuild_column(ᶜY_fv, ᶜY_arr)
    ᶠdata_col = rebuild_column(ᶠY_fv, ᶠY_arr)
    
    ᶜlg_arr = @localmem FT (ᶜNv, ᶜTS_lg)
    ᶠlg_arr = @localmem FT (ᶠNv, ᶜTS_lg)

    (ᶜspace_col, ᶠspace_col) = column_spaces(_ᶜY, _ᶠY, ᶠui, ᶜlg_arr, ᶠlg_arr)

    is_valid_index(ᶜus, ᶜui) && (ᶜdata_col[ᶜui] = ᶜY_fv[ᶜui])
    is_valid_index(ᶠus, ᶠui) && (ᶠdata_col[ᶠui] = ᶠY_fv[ᶠui])

    ᶜlg_col = Spaces.local_geometry_data(ᶜspace_col)
    ᶠlg_col = Spaces.local_geometry_data(ᶠspace_col)
    is_valid_index(ᶜus, ᶜui) && (ᶜlg_col[ᶜui] = ᶜlg[ᶜui])
    is_valid_index(ᶠus, ᶠui) && (ᶠlg_col[ᶠui] = ᶠlg[ᶠui])

    @synchronize

    ilc = @index(Local, Cartesian)
    igc = @index(Group, Cartesian)
    gs = @groupsize()
    ᶜui = universal_index_KA(ᶠus, ilc, igc, gs)
    ᶠui = universal_index_KA(ᶠus, ilc, igc, gs)

    ᶜdata_col = rebuild_column(ᶜY_fv, ᶜY_arr)
    ᶠdata_col = rebuild_column(ᶠY_fv, ᶠY_arr)

    (ᶜspace_col, ᶠspace_col) = column_spaces(_ᶜY, _ᶠY, ᶠui, ᶜlg_arr, ᶠlg_arr)

    if is_valid_index(ᶜus, ᶜui)
        (ᶜY, ᶠY) = column_states(_ᶜY, _ᶠY, ᶜdata_col, ᶠdata_col, ᶠui, ᶜspace_col, ᶠspace_col)
        ᶜbc = ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
        (ᶜidx, ᶜhidx) = operator_inds(axes(ᶜY), ᶜui)
        Fields.field_values(ᶜYₜ)[ᶜui] = Operators.getidx(axes(ᶜY), ᶜbc, ᶜidx, ᶜhidx)
        # ᶜYₜ[ᶜui] = ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t)[ᶜui] # might be possible?
    end
    if is_valid_index(ᶠus, ᶠui)
        (ᶜY, ᶠY) = column_states(_ᶜY, _ᶠY, ᶜdata_col, ᶠdata_col, ᶠui, ᶜspace_col, ᶠspace_col)
        ᶠbc = ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
        (ᶠidx, ᶠhidx) = operator_inds(axes(ᶠY), ᶠui)
        Fields.field_values(ᶠYₜ)[ᶠui] = Operators.getidx(axes(ᶠY), ᶠbc, ᶠidx, ᶠhidx)
        # ᶠYₜ[ᶠui] = ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t)[ᶠui] # might be possible?
    end
end

@inline function operator_inds(space, I)
    li = Operators.left_idx(space)
    (i, j, _, v, h) = I.I
    hidx = (i, j, h)
    idx = v - 1 + li
    return (idx, hidx)
end

@inline cartesian_indices(field::Fields.Field) =
    cartesian_indices(Fields.field_values(field))
@inline cartesian_indices(data::DataLayouts.AbstractData) =
    cartesian_indices(DataLayouts.UniversalSize(data))
@inline cartesian_indices(us::DataLayouts.UniversalSize) =
    CartesianIndices(map(Base.OneTo, DataLayouts.universal_size(us)))
@inline universal_index(x) = cartesian_indices(x)


function thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
    return @. lazy(TD.PhaseDry_ρe(
            thermo_params,
            ᶜρ,
            ᶜρe_tot / ᶜρ - ᶜK - Φ(grav, ᶜz),
        ))
end

# Drop everything except Nv and S:
@inline column_type_params(data::DataLayouts.AbstractData) = column_type_params(typeof(data))
@inline column_type_params(::Type{DataLayouts.IJFH{S, Nij, A}}) where {S, Nij, A} = (S, )
@inline column_type_params(::Type{DataLayouts.IJHF{S, Nij, A}}) where {S, Nij, A} = (S, )
@inline column_type_params(::Type{DataLayouts.IFH{S, Ni, A}}) where {S, Ni, A} = (S, )
@inline column_type_params(::Type{DataLayouts.IHF{S, Ni, A}}) where {S, Ni, A} = (S, )
@inline column_type_params(::Type{DataLayouts.DataF{S, A}}) where {S, A} = (S,)
@inline column_type_params(::Type{DataLayouts.IJF{S, Nij, A}}) where {S, Nij, A} = (S, )
@inline column_type_params(::Type{DataLayouts.IF{S, Ni, A}}) where {S, Ni, A} = (S, )
@inline column_type_params(::Type{DataLayouts.VF{S, Nv, A}}) where {S, Nv, A} = (S, Nv)
@inline column_type_params(::Type{DataLayouts.VIJFH{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = (S, Nv)
@inline column_type_params(::Type{DataLayouts.VIJHF{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = (S, Nv)
@inline column_type_params(::Type{DataLayouts.VIFH{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = (S, Nv)
@inline column_type_params(::Type{DataLayouts.VIHF{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = (S, Nv)

# Drop everything except V and F:
@inline column_singleton(::DataLayouts.IJFH) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.IJHF) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.IFH) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.IHF) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.DataF) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.IJF) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.IF) = DataLayouts.DataFSingleton()
@inline column_singleton(::DataLayouts.VF) = DataLayouts.VFSingleton()
@inline column_singleton(::DataLayouts.VIJFH) = DataLayouts.VFSingleton()
@inline column_singleton(::DataLayouts.VIJHF) = DataLayouts.VFSingleton()
@inline column_singleton(::DataLayouts.VIFH) = DataLayouts.VFSingleton()
@inline column_singleton(::DataLayouts.VIHF) = DataLayouts.VFSingleton()

function rebuild_column(data, array::AbstractArray)
    s_column = column_singleton(data)
    return DataLayouts.union_all(s_column){column_type_params(data)...}(array)
end

function column_lg_shmem(f, ui)
    (i, j, _, _, h) = ui.I
    colidx = Grids.ColumnIndex((i, j), h)
    lg = Spaces.local_geometry_data(axes(f))
    lg_col = Spaces.column(lg, colidx)
    FT = Spaces.undertype(axes(f))
    Nv = Spaces.nlevels(axes(f))
    TS = DataLayouts.typesize(FT, eltype(lg_col))
    lg_arr = CUDA.CuStaticSharedArray(FT, (Nv, TS))
    return rebuild_column(lg_col, lg_arr)
end

function column_lg_shmem_KA(f, ui, lg_arr)
    (i, j, _, _, h) = ui.I
    colidx = Grids.ColumnIndex((i, j), h)
    lg = Spaces.local_geometry_data(axes(f))
    lg_col = Spaces.column(lg, colidx)
    return rebuild_column(lg_col, lg_arr)
end

function column_spaces(ᶜY, ᶠY, ui, ᶜlg_arr, ᶠlg_arr)
    (i, j, _, _, h) = ui.I
    colidx = Grids.ColumnIndex((i, j), h)
    ᶜlg_col = column_lg_shmem_KA(ᶜY, ui, ᶜlg_arr)
    ᶠlg_col = column_lg_shmem_KA(ᶠY, ui, ᶠlg_arr)
    col_space = Spaces.column(axes(ᶜY), colidx)
    col_grid = Spaces.grid(col_space)
    if col_grid isa Grids.ColumnGrid && col_grid.full_grid isa Grids.DeviceExtrudedFiniteDifferenceGrid
        (; full_grid) = col_grid
        (; vertical_topology, global_geometry) = full_grid
        col_grid_shmem = Grids.DeviceFiniteDifferenceGrid(vertical_topology, global_geometry, ᶜlg_col, ᶠlg_col)
        ᶜspace_col = Spaces.space(col_grid_shmem, Grids.CellCenter())
        ᶠspace_col = Spaces.space(col_grid_shmem, Grids.CellFace())
    elseif col_grid isa Grids.ColumnGrid && col_grid.full_grid isa Grids.ExtrudedFiniteDifferenceGrid
        (; full_grid) = col_grid
        (; vertical_grid, global_geometry) = full_grid
        col_grid_shmem = Grids.FiniteDifferenceGrid(vertical_grid.topology, global_geometry, ᶜlg_col, ᶠlg_col)
        ᶜspace_col = Spaces.space(col_grid_shmem, Grids.CellCenter())
        ᶠspace_col = Spaces.space(col_grid_shmem, Grids.CellFace())
    else
        error("Uncaught case")
    end
    return (ᶜspace_col, ᶠspace_col)
end

function column_states(ᶜY, ᶠY, ᶜdata_col, ᶠdata_col, ui, ᶜspace_col, ᶠspace_col)
    ᶜY_col = Fields.Field(ᶜdata_col, ᶜspace_col)
    ᶠY_col = Fields.Field(ᶠdata_col, ᶠspace_col)
    return (ᶜY_col, ᶠY_col)
end

function vindex()
    (tv,) = CUDA.threadIdx()
    (h, bv, ij) = CUDA.blockIdx()
    v = tv + (bv - 1) * CUDA.blockDim().x
    return v
end

function ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
    (; rayleigh_sponge, params, dt) = p
    ᶜz = Fields.coordinate_field(ᶜY).z
    ᶜJ = Fields.local_geometry_field(ᶜY).J
    ᶠz = Fields.coordinate_field(ᶠY).z
    FT = Spaces.undertype(axes(ᶜY))
    grav = FT(CAP.grav(params))
    thermo_params = CAP.thermodynamics_params(params)
    ᶜρ = ᶜY.ρ
    ᶜρe_tot = ᶜY.ρe_tot
    ᶜuₕ = ᶜY.uₕ
    ᶠu₃ = ᶠY.u₃

    ᶜK = CA.compute_kinetic(ᶜuₕ, ᶠu₃)
    ᶜts = thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts))
    ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜρe_tot / ᶜρ))
    # Central advection of active tracers (e_tot and q_tot)
    ᶠuₕ³ = @. lazy(ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ)))
    ᶠu³ = @. lazy(ᶠuₕ³ + CT3(ᶠu₃))
    tend_ρ_1 = @. lazy( - ᶜdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠuₕ³))
    tend_ρe_tot_1 = CA.vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, Val(:none))
    ᶜuₕ₀ = (zero(eltype(ᶜuₕ)),)

    return @. lazy(ᶜtendencies(
        tend_ρ_1,
        - ᶜuₕ₀,
        tend_ρe_tot_1,
    ))
end

function ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t, zmax)
    (; rayleigh_sponge, params) = p
    ᶜz = Fields.coordinate_field(ᶜY).z
    ᶠz = Fields.coordinate_field(ᶠY).z
    FT = Spaces.undertype(axes(ᶜY))
    grav = FT(CAP.grav(params))
    thermo_params = CAP.thermodynamics_params(params)
    ᶜρ = ᶜY.ρ
    ᶜρe_tot = ᶜY.ρe_tot
    ᶜuₕ = ᶜY.uₕ
    ᶠu₃ = ᶠY.u₃
    ᶜK = CA.compute_kinetic(ᶜuₕ, ᶠu₃)
    ᶜts = thermo_state(thermo_params, ᶜρ, ᶜρe_tot, ᶜK, grav, ᶜz)
    ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts))
    bc1 = @. lazy(- (ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) + ᶠgradᵥ(Φ(grav, ᶜz))))
    bc2 = @. lazy(- CA.β_rayleigh_w(rayleigh_sponge, ᶠz, zmax) * ᶠu₃)
    return @. lazy(ᶠtendencies(bc1 + bc2))
end

function ImplicitEquationJacobian(
    Y::Fields.FieldVector;
    approximate_solve_iters = 1,
    transform_flag = false,
)
    FT = Spaces.undertype(axes(Y.c))
    CTh = CA.CTh_vector_type(axes(Y.c))

    BidiagonalRow_C3 = MatrixFields.BidiagonalMatrixRow{CA.C3{FT}}
    BidiagonalRow_ACT3 =
        MatrixFields.BidiagonalMatrixRow{LA.Adjoint{FT, CA.CT3{FT}}}
    BidiagonalRow_C3xACTh = MatrixFields.BidiagonalMatrixRow{
        typeof(zero(CA.C3{FT}) * zero(CTh{FT})'),
    }
    TridiagonalRow_C3xACT3 = MatrixFields.TridiagonalMatrixRow{
        typeof(zero(CA.C3{FT}) * zero(CA.CT3{FT})'),
    }

    is_in_Y(name) = MatrixFields.has_field(Y, name)

    sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : ()


    # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0,
    # which means that multiplying inv(-1) by a Float32 will yield a Float64.
    identity_blocks = MatrixFields.unrolled_map(
        name -> (name, name) => FT(-1) * LA.I,
        (@name(c.ρ), sfc_if_available...),
    )

    active_scalar_names = (@name(c.ρ), @name(c.ρe_tot))
    advection_blocks = (
        MatrixFields.unrolled_map(
            name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
            active_scalar_names,
        )...,
        MatrixFields.unrolled_map(
            name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),
            active_scalar_names,
        )...,
        (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh),
        (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
    )

    diffused_scalar_names = (@name(c.ρe_tot),)
    diffusion_blocks = MatrixFields.unrolled_map(
        name -> (name, name) => FT(-1) * LA.I,
        (diffused_scalar_names..., @name(c.uₕ)),
    )

    matrix = MatrixFields.FieldMatrix(
        identity_blocks...,
        advection_blocks...,
        diffusion_blocks...,
    )

    names₁_group₁ = (@name(c.ρ), sfc_if_available...)
    names₁_group₃ = (@name(c.ρe_tot),)
    names₁ = (names₁_group₁..., names₁_group₃...)

    alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ))
    alg = MatrixFields.BlockArrowheadSolve(names₁...; alg₂)

    return CA.ImplicitEquationJacobian(
        matrix,
        MatrixFields.FieldMatrixSolver(alg, matrix, Y),
        CA.IgnoreDerivative(), # diffusion_flag
        CA.IgnoreDerivative(), # topography_flag
        CA.IgnoreDerivative(), # sgs_advection_flag
        CA.IgnoreDerivative(), # sgs_entr_detr_flag
        CA.IgnoreDerivative(), # sgs_nh_pressure_flag
        CA.IgnoreDerivative(), # sgs_mass_flux_flag
        similar(Y),
        similar(Y),
        transform_flag,
        Ref{FT}(),
    )
end

function Wfact!(A, Y, p, dtγ, t)
    FT = Spaces.undertype(axes(Y.c))
    dtγ′ = FT(float(dtγ))
    A.dtγ_ref[] = dtγ′
    update_implicit_equation_jacobian!(A, Y, p, dtγ′)
end

Φ(grav, z) = grav * z

function update_implicit_equation_jacobian!(A, Y, p, dtγ)
    (; matrix) = A
    (; ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed
    (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p
    (; params) = p

    FT = Spaces.undertype(axes(Y.c))
    CTh = CA.CTh_vector_type(axes(Y.c))
    one_C3xACT3 = C3(FT(1)) * CT3(FT(1))'
    rs = p.rayleigh_sponge
    ᶠz = Fields.coordinate_field(Y.f).z
    zmax = CA.z_max(axes(Y.f))

    T_0 = FT(CAP.T_0(params))
    cp_d = FT(CAP.cp_d(params))
    thermo_params = CAP.thermodynamics_params(params)
    ᶜz = Fields.coordinate_field(Y.c).z
    grav = FT(CAP.grav(params))

    ᶜρ = Y.c.ρ
    ᶜuₕ = Y.c.uₕ
    ᶠu₃ = Y.f.u₃
    ᶜJ = Fields.local_geometry_field(Y.c).J
    ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ

    ᶜkappa_m = p.ᶜtemp_scalar
    @. ᶜkappa_m =
        TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts)

    @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ)))
    @. ∂ᶜK_∂ᶠu₃ =
        ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) +
        DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix()

    @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix()

    @. ᶜadvection_matrix =
        -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ))

    ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)]
    @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(CA.g³³(ᶠgⁱʲ))

    ∂ᶜρχ_err_∂ᶠu₃ = matrix[@name(c.ρe_tot), @name(f.u₃)]
    @. ∂ᶜρχ_err_∂ᶠu₃ =
        dtγ * ᶜadvection_matrix ⋅
        DiagonalMatrixRow(ᶠinterp(ᶜh_tot) * CA.g³³(ᶠgⁱʲ))

    ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)]
    ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)]

    @. ∂ᶠu₃_err_∂ᶜρ =
        dtγ * (
            ᶠp_grad_matrix ⋅
            DiagonalMatrixRow(ᶜkappa_m * (T_0 * cp_d - ᶜK - Φ(grav, ᶜz))) +
            DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / abs2(ᶠinterp(ᶜρ))) ⋅
            ᶠinterp_matrix()
        )
    @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m)

    ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)]
    ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)]
    I_u₃ = DiagonalMatrixRow(one_C3xACT3)
    @. ∂ᶠu₃_err_∂ᶜuₕ =
        dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ

    @. ∂ᶠu₃_err_∂ᶠu₃ =
        dtγ * (
            ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶠu₃ +
            DiagonalMatrixRow(-CA.β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,))
        ) - (I_u₃,)

end

function set_precomputed_quantities!(Y, p, t)
    thermo_params = CAP.thermodynamics_params(p.params)
    (; ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp) = p.precomputed

    ᶜρ = Y.c.ρ
    ᶜuₕ = Y.c.uₕ
    ᶜz = Fields.coordinate_field(Y.c).z
    grav = FT(CAP.grav(params))
    ᶠu₃ = Y.f.u₃
    @. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
    ᶠu³ .= CA.compute_ᶠuₕ³(ᶜuₕ, ᶜρ) .+ CT3.(ᶠu₃)
    ᶜK .= CA.compute_kinetic(ᶜuₕ, ᶠu₃)

    @. ᶜts = TD.PhaseDry_ρe(
        thermo_params,
        Y.c.ρ,
        Y.c.ρe_tot / Y.c.ρ - ᶜK - Φ(grav, ᶜz),
    )
    @. ᶜp = TD.air_pressure(thermo_params, ᶜts)

    (; ᶜh_tot) = p.precomputed
    @. ᶜh_tot =
        TD.total_specific_enthalpy(thermo_params, ᶜts, Y.c.ρe_tot / Y.c.ρ)
    return nothing
end

function dss!(Y, p, t)
    Spaces.weighted_dss!(Y.c => p.ghost_buffer.c, Y.f => p.ghost_buffer.f)
    return nothing
end

function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
    # Yₜ_lim .= zero(eltype(Yₜ_lim))
    Yₜ .= zero(eltype(Yₜ))
    (; dt, params, rayleigh_sponge) = p
    (; ᶜh_tot) = p.precomputed
    (; ᶠu³, ᶜu, ᶜK, ᶜp) = p.precomputed
    (; ᶜf³, ᶠf¹²) = p.precomputed
    ᶜz = Fields.coordinate_field(Y.c).z
    ᶜJ = Fields.local_geometry_field(Y.c).J
    grav = FT(CAP.grav(params))
    ᶜuₕ = Y.c.uₕ
    ᶠu₃ = Y.f.u₃
    ᶜρ = Y.c.ρ

    @. Yₜ.c.ρ -= wdivₕ(ᶜρ * ᶜu)
    @. Yₜ.c.ρe_tot -= wdivₕ(ᶜρ * ᶜh_tot * ᶜu)
    @. Yₜ.c.uₕ -= C12(gradₕ(ᶜp) / ᶜρ + gradₕ(ᶜK + Φ(grav, ᶜz)))

    ᶜω³ = p.scratch.ᶜtemp_CT3
    ᶠω¹² = p.scratch.ᶠtemp_CT12

    point_type = eltype(Fields.coordinate_field(Y.c))
    if point_type <: Geometry.Abstract3DPoint
        @. ᶜω³ = curlₕ(ᶜuₕ)
    elseif point_type <: Geometry.Abstract2DPoint
        @. ᶜω³ = zero(ᶜω³)
    end

    @. ᶠω¹² = ᶠcurlᵥ(ᶜuₕ)
    @. ᶠω¹² += CT12(curlₕ(ᶠu₃))
    # Without the CT12(), the right-hand side would be a CT1 or CT2 in 2D space.

    ᶠω¹²′ = if isnothing(ᶠf¹²)
        ᶠω¹² # shallow atmosphere
    else
        @. lazy(ᶠf¹² + ᶠω¹²) # deep atmosphere
    end

    @. Yₜ.c.uₕ -=
        ᶜinterp(ᶠω¹²′ × (ᶠinterp(ᶜρ * ᶜJ) * ᶠu³)) / (ᶜρ * ᶜJ) +
        (ᶜf³ + ᶜω³) × CT12(ᶜu)
    @. Yₜ.f.u₃ -= ᶠω¹²′ × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)

    Yₜ.c.uₕ .+= CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)

    return Yₜ
end

# This block:
# @time if !@isdefined(integrator)
    FT = Float64;
    if high_res
        ᶜspace = ExtrudedCubedSphereSpace(
            FT;
            z_elem = 63,
            z_min = 0,
            z_max = 30000.0,
            radius = 6.371e6,
            h_elem = 30,
            n_quad_points = 4,
            staggering = CellCenter(),
        );
    else
        ᶜspace = ExtrudedCubedSphereSpace(
            FT;
            z_elem = 8,
            z_min = 0,
            z_max = 30000.0,
            radius = 6.371e6,
            h_elem = 2,
            n_quad_points = 2,
            staggering = CellCenter(),
        );
    end
    ᶠspace = Spaces.face_space(ᶜspace);
    cnt = (; ρ = zero(FT), uₕ = zero(CA.C12{FT}), ρe_tot = zero(FT));
    Yc = Fields.fill(cnt, ᶜspace);
    fill!(parent(Yc.ρ), 1)
    fill!(parent(Yc.uₕ), 0.01)
    fill!(parent(Yc.ρe_tot), 1000.0)
    Yf = Fields.fill((; u₃ = zero(CA.C3{FT})), ᶠspace);
    Y = Fields.FieldVector(; c = Yc, f = Yf);

    A = ImplicitEquationJacobian(
        Y;
        approximate_solve_iters = 2,
        transform_flag = false, # assumes use_transform returns false
    )

    implicit_func = SciMLBase.ODEFunction(
        implicit_tendency!;
        jac_prototype = A,
        Wfact = Wfact!, # assumes use_transform returns false
        tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0),
    )

    func = CTS.ClimaODEFunction(;
        T_exp_T_lim! = remaining_tendency!,
        T_imp! = implicit_func,
        # Can we just pass implicit_tendency! and jac_prototype etc.?
        lim! = (Y, p, t, ref_Y) -> nothing, # limiters_func!
        dss!,
        cache! = set_precomputed_quantities!,
        cache_imp! = set_precomputed_quantities!,
    )

    newtons_method = CTS.NewtonsMethod(; max_iters = 2)
    params = CA.ClimaAtmosParameters(FT)
    ᶠcoord = Fields.coordinate_field(ᶠspace);
    ᶜcoord = Fields.coordinate_field(ᶜspace);
    (; ᶜf³, ᶠf¹²) = CA.compute_coriolis(ᶜcoord, ᶠcoord, params);
    scratch = (;
        ᶜtemp_CT3 = Fields.Field(CT3{FT}, ᶜspace),
        ᶠtemp_CT12 = Fields.Field(CT12{FT}, ᶠspace),
    )
    precomputed = (;
        ᶜh_tot = Fields.Field(FT, ᶜspace),
        ᶠu³ = Fields.Field(CA.CT3{FT}, ᶠspace),
        ᶜf³,
        ᶠf¹²,
        ᶜp = Fields.Field(FT, ᶜspace),
        ᶜK = Fields.Field(FT, ᶜspace),
        ᶜts = Fields.Field(TD.PhaseDry{FT}, ᶜspace),
        ᶠu = Fields.Field(C123{FT}, ᶠspace),
        ᶜu = Fields.Field(C123{FT}, ᶜspace),
    )
    dt = FT(0.1)

    ghost_buffer =
        !CA.do_dss(axes(Y.c)) ? (;) :
        (; c = Spaces.create_dss_buffer(Y.c), f = Spaces.create_dss_buffer(Y.f))

    CTh = CA.CTh_vector_type(axes(Y.c))
    p = (;
        rayleigh_sponge = CA.RayleighSponge{FT}(;
            zd = params.zd_rayleigh,
            α_uₕ = params.alpha_rayleigh_uh,
            α_w = params.alpha_rayleigh_w,
        ),
        params,
        ∂ᶜK_∂ᶜuₕ = Fields.Field(DiagonalMatrixRow{Adjoint{FT, CTh{FT}}}, ᶜspace),
        ∂ᶜK_∂ᶠu₃ = Fields.Field(BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}, ᶜspace),
        ᶜadvection_matrix = Fields.Field(
            BidiagonalMatrixRow{Adjoint{FT, C3{FT}}},
            ᶜspace,
        ),
        ᶜtemp_scalar = Fields.Field(FT, ᶜspace),
        ᶠp_grad_matrix = Fields.Field(BidiagonalMatrixRow{C3{FT}}, ᶠspace),
        scratch,
        ghost_buffer,
        dt,
        precomputed,
    )
    ode_algo = CTS.IMEXAlgorithm(CTS.ARS343(), newtons_method)
    problem = SciMLBase.ODEProblem(func, Y, (FT(0), FT(1)), p)
    integrator = SciMLBase.init(problem, ode_algo; dt)
    Yₜ = similar(integrator.u);
# end

function main!(integrator, Yₜ, n)
    for _ in 1:n
        # @time SciMLBase.step!(integrator)
        @time implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
    end
    return nothing
end
using Test
if ClimaComms.device() isa ClimaComms.CUDADevice
    Yₜ_bc = similar(Yₜ);
    @. Yₜ_bc = 0
    @. Yₜ = 0
    Yc = integrator.u.c;
    Yf = integrator.u.f;
    fill!(parent(Yc.ρ), 1);
    zc = Fields.coordinate_field(Yc).z;
    zf = Fields.coordinate_field(Yf).z;
    @. Yc.ρ += 0.1*sin(zc);
    parent(Yf.u₃) .+= 0.001 .* sin.(parent(zf));
    fill!(parent(Yc.uₕ), 0.01);
    fill!(parent(Yc.ρe_tot), 100000.0);

    implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
    implicit_tendency_bc!(Yₜ_bc, integrator.u, integrator.p, integrator.t)
    abs_err_c = maximum(Array(abs.(parent(Yₜ.c) .- parent(Yₜ_bc.c))))
    abs_err_f = maximum(Array(abs.(parent(Yₜ.f) .- parent(Yₜ_bc.f))))
    results_match = abs_err_c < 6e-9 && abs_err_c < 6e-9
    if !results_match
        @show norm(Array(parent(Yₜ_bc.c))), norm(Array(parent(Yₜ.c)))
        @show norm(Array(parent(Yₜ_bc.f))), norm(Array(parent(Yₜ.f)))
        @show abs_err_c
        @show abs_err_f
    end
    @test results_match
    println(CUDA.@profile trace=true begin
        # SciMLBase.step!(integrator)
        implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
        implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
        implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
        implicit_tendency!(Yₜ, integrator.u, integrator.p, integrator.t)
    end)
    println(CUDA.@profile begin
        # SciMLBase.step!(integrator)
        @. Yₜ += 1
        @. Yₜ += 1
        @. Yₜ += 1
        @. Yₜ += 1
    end)
else
    @info "Compiling main loop"
    @time main!(integrator, Yₜ, 1)
    @info "Running main loop"
    @time main!(integrator, Yₜ, 3)
end


nothing

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the minimal reproducers in ka_reproducer.jl and cuda_reproducer.jl, then compare their generated kernels and profiler timings. Use the real-world script's implicit_tendency! toggle to isolate the slowdown. Done means identifying the source of the CUDA versus KernelAbstractions performance discrepancy and documenting or correcting it without breaking the existing correctness check.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.