text:
using .NumericalFluxes:
    NumericalFluxGradient,
    numerical_boundary_flux_gradient!,
    numerical_flux_gradient!,
    NumericalFluxNonDiffusive,
    NumericalFluxDiffusive,
    numerical_flux_nondiffusive!,
    numerical_boundary_flux_nondiffusive!,
    numerical_flux_diffusive!,
    numerical_boundary_flux_diffusive!,
    divergence_penalty!,
    divergence_boundary_penalty!,
    divergence_penalty!,
    numerical_boundary_flux_hyperdiffusive!,
    numerical_flux_hyperdiffusive!

using ..Mesh.Geometry

using Requires
@init @require CUDAnative = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17" begin
    using .CUDAnative
end

# {{{ FIXME: remove this after we've figure out how to pass through to kernel
const _ξ1x1, _ξ2x1, _ξ3x1 = Grids._ξ1x1, Grids._ξ2x1, Grids._ξ3x1
const _ξ1x2, _ξ2x2, _ξ3x2 = Grids._ξ1x2, Grids._ξ2x2, Grids._ξ3x2
const _ξ1x3, _ξ2x3, _ξ3x3 = Grids._ξ1x3, Grids._ξ2x3, Grids._ξ3x3
const _M, _MI = Grids._M, Grids._MI
const _x1, _x2, _x3 = Grids._x1, Grids._x2, Grids._x3
const _JcV = Grids._JcV

const _n1, _n2, _n3 = Grids._n1, Grids._n2, Grids._n3
const _sM, _vMI = Grids._sM, Grids._vMI
# }}}

@doc """
    volumerhs!(bl::BalanceLaw, Val(polyorder), rhs, Q, Qvisc, auxstate,
               vgeo, t, D, elems)

Computational kernel: Evaluate the volume integrals on right-hand side of a
`DGBalanceLaw` semi-discretization.

See [`odefun!`](@ref) for usage.
""" volumerhs!
@kernel function volumerhs!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    rhs,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    t,
    ω,
    D,
    elems,
    increment,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)

        ngradlapstate = num_gradient_laplacian(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        l_S = MArray{Tuple{nstate}, FT}(undef)
        l_Q = MArray{Tuple{nstate}, FT}(undef)
        l_Qvisc = MArray{Tuple{nviscstate}, FT}(undef)
        l_Qhypervisc = MArray{Tuple{nhyperviscstate}, FT}(undef)
        l_aux = MArray{Tuple{nauxstate}, FT}(undef)
        l_F = MArray{Tuple{3, nstate}, FT}(undef)
    end

    s_F = @localmem FT (3, Nq, Nq, Nqk, nstate)
    s_ω = @localmem FT (Nq,)
    s_D = @localmem FT (Nq, Nq)
    l_rhs = @private FT (nstate,)

    e = @index(Group, Linear)
    ijk = @index(Local, Linear)
    i, j, k = @index(Local, NTuple)

    @inbounds begin
        s_ω[j] = ω[j]
        s_D[i, j] = D[i, j]

        M = vgeo[ijk, _M, e]
        ξ1x1 = vgeo[ijk, _ξ1x1, e]
        ξ1x2 = vgeo[ijk, _ξ1x2, e]
        ξ1x3 = vgeo[ijk, _ξ1x3, e]
        if dim == 3 || (dim == 2 && direction isa EveryDirection)
            ξ2x1 = vgeo[ijk, _ξ2x1, e]
            ξ2x2 = vgeo[ijk, _ξ2x2, e]
            ξ2x3 = vgeo[ijk, _ξ2x3, e]
        end
        if dim == 3 && direction isa EveryDirection
            ξ3x1 = vgeo[ijk, _ξ3x1, e]
            ξ3x2 = vgeo[ijk, _ξ3x2, e]
            ξ3x3 = vgeo[ijk, _ξ3x3, e]
        end

        @unroll for s in 1:nstate
            l_rhs[s] = increment ? rhs[ijk, s, e] : zero(FT)
        end

        @unroll for s in 1:nstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        @unroll for s in 1:nviscstate
            l_Qvisc[s] = Qvisc[ijk, s, e]
        end

        @unroll for s in 1:nhyperviscstate
            l_Qhypervisc[s] = Qhypervisc_grad[ijk, s, e]
        end

        fill!(l_F, -zero(eltype(l_F)))
        flux_nondiffusive!(
            bl,
            Grad{vars_state(bl, FT)}(l_F),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
        )

        @unroll for s in 1:nstate
            s_F[1, i, j, k, s] = l_F[1, s]
            s_F[2, i, j, k, s] = l_F[2, s]
            s_F[3, i, j, k, s] = l_F[3, s]
        end

        fill!(l_F, -zero(eltype(l_F)))
        flux_diffusive!(
            bl,
            Grad{vars_state(bl, FT)}(l_F),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_diffusive(bl, FT)}(l_Qvisc),
            Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
        )

        @unroll for s in 1:nstate
            s_F[1, i, j, k, s] += l_F[1, s]
            s_F[2, i, j, k, s] += l_F[2, s]
            s_F[3, i, j, k, s] += l_F[3, s]
        end

        # Build "inside metrics" flux
        @unroll for s in 1:nstate
            F1, F2, F3 =
                s_F[1, i, j, k, s], s_F[2, i, j, k, s], s_F[3, i, j, k, s]

            s_F[1, i, j, k, s] = M * (ξ1x1 * F1 + ξ1x2 * F2 + ξ1x3 * F3)
            if dim == 3 || (dim == 2 && direction isa EveryDirection)
                s_F[2, i, j, k, s] = M * (ξ2x1 * F1 + ξ2x2 * F2 + ξ2x3 * F3)
            end
            if dim == 3 && direction isa EveryDirection
                s_F[3, i, j, k, s] = M * (ξ3x1 * F1 + ξ3x2 * F2 + ξ3x3 * F3)
            end
        end

        fill!(l_S, -zero(eltype(l_S)))
        source!(
            bl,
            Vars{vars_state(bl, FT)}(l_S),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_diffusive(bl, FT)}(l_Qvisc),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
            direction,
        )

        @unroll for s in 1:nstate
            l_rhs[s] += l_S[s]
        end
        @synchronize

        # Weak "inside metrics" derivative
        MI = vgeo[ijk, _MI, e]
        @unroll for s in 1:nstate
            @unroll for n in 1:Nq
                # ξ1-grid lines
                l_rhs[s] += MI * s_D[n, i] * s_F[1, n, j, k, s]

                # ξ2-grid lines
                if dim == 3 || (dim == 2 && direction isa EveryDirection)
                    l_rhs[s] += MI * s_D[n, j] * s_F[2, i, n, k, s]
                end

                # ξ3-grid lines
                if dim == 3 && direction isa EveryDirection
                    l_rhs[s] += MI * s_D[n, k] * s_F[3, i, j, n, s]
                end
            end
        end
        ijk = i + Nq * ((j - 1) + Nq * (k - 1))
        @unroll for s in 1:nstate
            rhs[ijk, s, e] = l_rhs[s]
        end
    end
    @synchronize
end

@kernel function volumerhs!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction::VerticalDirection,
    rhs,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    t,
    ω,
    D,
    elems,
    increment,
) where {dim, polyorder}

    @uniform begin
        N = polyorder
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)

        ngradlapstate = num_gradient_laplacian(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        l_S = MArray{Tuple{nstate}, FT}(undef)
        l_Q = MArray{Tuple{nstate}, FT}(undef)
        l_Qvisc = MArray{Tuple{nviscstate}, FT}(undef)
        l_Qhypervisc = MArray{Tuple{nhyperviscstate}, FT}(undef)
        l_aux = MArray{Tuple{nauxstate}, FT}(undef)
        l_F = MArray{Tuple{3, nstate}, FT}(undef)

        _ζx1 = dim == 2 ? _ξ2x1 : _ξ3x1
        _ζx2 = dim == 2 ? _ξ2x2 : _ξ3x2
        _ζx3 = dim == 2 ? _ξ2x3 : _ξ3x3
    end

    s_F = @localmem FT (3, Nq, Nq, Nqk, nstate)
    s_ω = @localmem FT (Nq,)
    s_D = @localmem FT (Nq, Nq)
    l_rhs = @private FT (nstate,)

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)
    ijk = @index(Local, Linear)

    @inbounds begin
        s_ω[j] = ω[j]
        s_D[i, j] = D[i, j]

        M = vgeo[ijk, _M, e]
        ζx1 = vgeo[ijk, _ζx1, e]
        ζx2 = vgeo[ijk, _ζx2, e]
        ζx3 = vgeo[ijk, _ζx3, e]

        @unroll for s in 1:nstate
            l_rhs[s] = increment ? rhs[ijk, s, e] : zero(FT)
        end

        @unroll for s in 1:nstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        @unroll for s in 1:nviscstate
            l_Qvisc[s] = Qvisc[ijk, s, e]
        end

        @unroll for s in 1:nhyperviscstate
            l_Qhypervisc[s] = Qhypervisc_grad[ijk, s, e]
        end

        fill!(l_F, -zero(eltype(l_F)))
        flux_nondiffusive!(
            bl,
            Grad{vars_state(bl, FT)}(l_F),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
        )

        @unroll for s in 1:nstate
            s_F[1, i, j, k, s] = l_F[1, s]
            s_F[2, i, j, k, s] = l_F[2, s]
            s_F[3, i, j, k, s] = l_F[3, s]
        end

        fill!(l_F, -zero(eltype(l_F)))
        flux_diffusive!(
            bl,
            Grad{vars_state(bl, FT)}(l_F),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_diffusive(bl, FT)}(l_Qvisc),
            Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
        )

        @unroll for s in 1:nstate
            s_F[1, i, j, k, s] += l_F[1, s]
            s_F[2, i, j, k, s] += l_F[2, s]
            s_F[3, i, j, k, s] += l_F[3, s]
        end

        # Build "inside metrics" flux
        @unroll for s in 1:nstate
            F1, F2, F3 =
                s_F[1, i, j, k, s], s_F[2, i, j, k, s], s_F[3, i, j, k, s]
            s_F[3, i, j, k, s] = M * (ζx1 * F1 + ζx2 * F2 + ζx3 * F3)
        end

        fill!(l_S, -zero(eltype(l_S)))
        source!(
            bl,
            Vars{vars_state(bl, FT)}(l_S),
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_diffusive(bl, FT)}(l_Qvisc),
            Vars{vars_aux(bl, FT)}(l_aux),
            t,
            direction,
        )

        @unroll for s in 1:nstate
            l_rhs[s] += l_S[s]
        end
        @synchronize

        # Weak "inside metrics" derivative
        MI = vgeo[ijk, _MI, e]
        @unroll for s in 1:nstate
            @unroll for n in 1:Nq
                if dim == 2
                    Dnj = s_D[n, j]
                    l_rhs[s] += MI * Dnj * s_F[3, i, n, k, s]
                else
                    Dnk = s_D[n, k]
                    l_rhs[s] += MI * Dnk * s_F[3, i, j, n, s]
                end
            end
        end
        @unroll for s in 1:nstate
            rhs[ijk, s, e] = l_rhs[s]
        end
        @synchronize
    end
end

@doc """
    facerhs!(bl::BalanceLaw, Val(polyorder),
            numfluxnondiff::NumericalFluxNonDiffusive,
            numfluxdiff::NumericalFluxDiffusive,
            rhs, Q, Qvisc, auxstate,
            vgeo, sgeo, t, vmap⁻, vmap⁺, elemtobndy,
            elems)

Computational kernel: Evaluate the surface integrals on right-hand side of a
`BalanceLaw` semi-discretization.

See [`odefun!`](@ref) for usage.
""" facerhs!
@kernel function facerhs!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    numfluxnondiff::NumericalFluxNonDiffusive,
    numfluxdiff::NumericalFluxDiffusive,
    rhs,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    sgeo,
    t,
    vmap⁻,
    vmap⁺,
    elemtobndy,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)
        nauxstate = num_aux(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)

        if dim == 1
            Np = (N + 1)
            Nfp = 1
            nface = 2
        elseif dim == 2
            Np = (N + 1) * (N + 1)
            Nfp = (N + 1)
            nface = 4
        elseif dim == 3
            Np = (N + 1) * (N + 1) * (N + 1)
            Nfp = (N + 1) * (N + 1)
            nface = 6
        end

        faces = 1:nface
        if direction isa VerticalDirection
            faces = (nface - 1):nface
        elseif direction isa HorizontalDirection
            faces = 1:(nface - 2)
        end

        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq

        l_Q⁻ = MArray{Tuple{nstate}, FT}(undef)
        l_Qvisc⁻ = MArray{Tuple{nviscstate}, FT}(undef)
        l_Qhypervisc⁻ = MArray{Tuple{nhyperviscstate}, FT}(undef)
        l_aux⁻ = MArray{Tuple{nauxstate}, FT}(undef)

        # Need two copies since numerical_flux_nondiffusive! can modify Q⁺
        l_Q⁺nondiff = MArray{Tuple{nstate}, FT}(undef)
        l_Q⁺diff = MArray{Tuple{nstate}, FT}(undef)

        # Need two copies since numerical_flux_nondiffusive! can modify aux⁺
        l_aux⁺nondiff = MArray{Tuple{nauxstate}, FT}(undef)
        l_aux⁺diff = MArray{Tuple{nauxstate}, FT}(undef)

        l_Qvisc⁺ = MArray{Tuple{nviscstate}, FT}(undef)
        l_Qhypervisc⁺ = MArray{Tuple{nhyperviscstate}, FT}(undef)

        l_Q_bot1 = MArray{Tuple{nstate}, FT}(undef)
        l_Qvisc_bot1 = MArray{Tuple{nviscstate}, FT}(undef)
        l_aux_bot1 = MArray{Tuple{nauxstate}, FT}(undef)

        l_F = MArray{Tuple{nstate}, FT}(undef)
    end

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    e = @private Int (1,)
    @inbounds e[1] = elems[eI]

    @inbounds for f in faces
        e⁻ = e[1]
        n⁻ = SVector(
            sgeo[_n1, n, f, e⁻],
            sgeo[_n2, n, f, e⁻],
            sgeo[_n3, n, f, e⁻],
        )
        sM, vMI = sgeo[_sM, n, f, e⁻], sgeo[_vMI, n, f, e⁻]
        id⁻, id⁺ = vmap⁻[n, f, e⁻], vmap⁺[n, f, e⁻]
        e⁺ = ((id⁺ - 1) ÷ Np) + 1

        vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

        # Load minus side data
        @unroll for s in 1:nstate
            l_Q⁻[s] = Q[vid⁻, s, e⁻]
        end

        @unroll for s in 1:nviscstate
            l_Qvisc⁻[s] = Qvisc[vid⁻, s, e⁻]
        end

        @unroll for s in 1:nhyperviscstate
            l_Qhypervisc⁻[s] = Qhypervisc_grad[vid⁻, s, e⁻]
        end

        @unroll for s in 1:nauxstate
            l_aux⁻[s] = auxstate[vid⁻, s, e⁻]
        end

        # Load plus side data
        @unroll for s in 1:nstate
            l_Q⁺diff[s] = l_Q⁺nondiff[s] = Q[vid⁺, s, e⁺]
        end

        @unroll for s in 1:nviscstate
            l_Qvisc⁺[s] = Qvisc[vid⁺, s, e⁺]
        end

        @unroll for s in 1:nhyperviscstate
            l_Qhypervisc⁺[s] = Qhypervisc_grad[vid⁺, s, e⁺]
        end

        @unroll for s in 1:nauxstate
            l_aux⁺diff[s] = l_aux⁺nondiff[s] = auxstate[vid⁺, s, e⁺]
        end

        bctype = elemtobndy[f, e⁻]
        fill!(l_F, -zero(eltype(l_F)))
        if bctype == 0
            numerical_flux_nondiffusive!(
                numfluxnondiff,
                bl,
                Vars{vars_state(bl, FT)}(l_F),
                SVector(n⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_state(bl, FT)}(l_Q⁺nondiff),
                Vars{vars_aux(bl, FT)}(l_aux⁺nondiff),
                t,
            )
            numerical_flux_diffusive!(
                numfluxdiff,
                bl,
                Vars{vars_state(bl, FT)}(l_F),
                n⁻,
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_diffusive(bl, FT)}(l_Qvisc⁻),
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_state(bl, FT)}(l_Q⁺diff),
                Vars{vars_diffusive(bl, FT)}(l_Qvisc⁺),
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺diff),
                t,
            )
        else
            if (dim == 2 && f == 3) || (dim == 3 && f == 5)
                # Loop up the first element along all horizontal elements
                @unroll for s in 1:nstate
                    l_Q_bot1[s] = Q[n + Nqk^2, s, e⁻]
                end
                @unroll for s in 1:nviscstate
                    l_Qvisc_bot1[s] = Qvisc[n + Nqk^2, s, e⁻]
                end
                @unroll for s in 1:nauxstate
                    l_aux_bot1[s] = auxstate[n + Nqk^2, s, e⁻]
                end
            end
            numerical_boundary_flux_nondiffusive!(
                numfluxnondiff,
                bl,
                Vars{vars_state(bl, FT)}(l_F),
                SVector(n⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_state(bl, FT)}(l_Q⁺nondiff),
                Vars{vars_aux(bl, FT)}(l_aux⁺nondiff),
                bctype,
                t,
                Vars{vars_state(bl, FT)}(l_Q_bot1),
                Vars{vars_aux(bl, FT)}(l_aux_bot1),
            )
            numerical_boundary_flux_diffusive!(
                numfluxdiff,
                bl,
                Vars{vars_state(bl, FT)}(l_F),
                n⁻,
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_diffusive(bl, FT)}(l_Qvisc⁻),
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_state(bl, FT)}(l_Q⁺diff),
                Vars{vars_diffusive(bl, FT)}(l_Qvisc⁺),
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺diff),
                bctype,
                t,
                Vars{vars_state(bl, FT)}(l_Q_bot1),
                Vars{vars_diffusive(bl, FT)}(l_Qvisc_bot1),
                Vars{vars_aux(bl, FT)}(l_aux_bot1),
            )
        end

        #Update RHS
        @unroll for s in 1:nstate
            # FIXME: Should we pretch these?
            rhs[vid⁻, s, e⁻] -= vMI * sM * l_F[s]
        end
        # Need to wait after even faces to avoid race conditions
        @synchronize(f % 2 == 0)
    end
end

@kernel function volumeviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    t,
    D,
    hypervisc_indexmap,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder

        FT = eltype(Q)
        nstate = num_state(bl, FT)
        ngradstate = num_gradient(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        ngradtransformstate = nstate

        l_G = MArray{Tuple{ngradstate}, FT}(undef)
        l_Qvisc = MArray{Tuple{nviscstate}, FT}(undef)
        l_gradG = MArray{Tuple{3, ngradstate}, FT}(undef)
    end

    s_G = @localmem FT (Nq, Nq, Nqk, ngradstate)
    s_D = @localmem FT (Nq, Nq)

    l_Q = @private FT (ngradtransformstate,)
    l_aux = @private FT (nauxstate,)

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)
    ijk = @index(Local, Linear)

    @inbounds @views begin
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradtransformstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        fill!(l_G, -zero(eltype(l_G)))
        gradvariables!(
            bl,
            Vars{vars_gradient(bl, FT)}(l_G),
            Vars{vars_state(bl, FT)}(l_Q[:]),
            Vars{vars_aux(bl, FT)}(l_aux[:]),
            t,
        )
        @unroll for s in 1:ngradstate
            s_G[i, j, k, s] = l_G[s]
        end
        @synchronize

        # Compute gradient of each state
        ξ1x1, ξ1x2, ξ1x3 =
            vgeo[ijk, _ξ1x1, e], vgeo[ijk, _ξ1x2, e], vgeo[ijk, _ξ1x3, e]
        if dim == 3 || (dim == 2 && direction isa EveryDirection)
            ξ2x1, ξ2x2, ξ2x3 =
                vgeo[ijk, _ξ2x1, e], vgeo[ijk, _ξ2x2, e], vgeo[ijk, _ξ2x3, e]
        end
        if dim == 3 && direction isa EveryDirection
            ξ3x1, ξ3x2, ξ3x3 =
                vgeo[ijk, _ξ3x1, e], vgeo[ijk, _ξ3x2, e], vgeo[ijk, _ξ3x3, e]
        end

        @unroll for s in 1:ngradstate
            Gξ1 = Gξ2 = Gξ3 = zero(FT)
            @unroll for n in 1:Nq
                Gξ1 += s_D[i, n] * s_G[n, j, k, s]
                if dim == 3 || (dim == 2 && direction isa EveryDirection)
                    Gξ2 += s_D[j, n] * s_G[i, n, k, s]
                end
                if dim == 3 && direction isa EveryDirection
                    Gξ3 += s_D[k, n] * s_G[i, j, n, s]
                end
            end
            l_gradG[1, s] = ξ1x1 * Gξ1
            l_gradG[2, s] = ξ1x2 * Gξ1
            l_gradG[3, s] = ξ1x3 * Gξ1

            if dim == 3 || (dim == 2 && direction isa EveryDirection)
                l_gradG[1, s] += ξ2x1 * Gξ2
                l_gradG[2, s] += ξ2x2 * Gξ2
                l_gradG[3, s] += ξ2x3 * Gξ2
            end

            if dim == 3 && direction isa EveryDirection
                l_gradG[1, s] += ξ3x1 * Gξ3
                l_gradG[2, s] += ξ3x2 * Gξ3
                l_gradG[3, s] += ξ3x3 * Gξ3
            end
        end

        @unroll for s in 1:ngradlapstate
            Qhypervisc_grad[ijk, 3 * (s - 1) + 1, e] =
                l_gradG[1, hypervisc_indexmap[s]]
            Qhypervisc_grad[ijk, 3 * (s - 1) + 2, e] =
                l_gradG[2, hypervisc_indexmap[s]]
            Qhypervisc_grad[ijk, 3 * (s - 1) + 3, e] =
                l_gradG[3, hypervisc_indexmap[s]]
        end

        if nviscstate > 0
            fill!(l_Qvisc, -zero(eltype(l_Qvisc)))
            diffusive!(
                bl,
                Vars{vars_diffusive(bl, FT)}(l_Qvisc),
                Grad{vars_gradient(bl, FT)}(l_gradG),
                Vars{vars_state(bl, FT)}(l_Q[:]),
                Vars{vars_aux(bl, FT)}(l_aux[:]),
                t,
            )

            @unroll for s in 1:nviscstate
                Qvisc[ijk, s, e] = l_Qvisc[s]
            end
        end
    end
    @synchronize
end

@kernel function volumeviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    ::VerticalDirection,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    t,
    D,
    hypervisc_indexmap,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder

        FT = eltype(Q)
        nstate = num_state(bl, FT)
        ngradstate = num_gradient(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)
        ngradtransformstate = nstate

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq
        l_G = MArray{Tuple{ngradstate}, FT}(undef)
        l_Qvisc = MArray{Tuple{nviscstate}, FT}(undef)
        l_gradG = MArray{Tuple{3, ngradstate}, FT}(undef)

        _ζx1 = dim == 2 ? _ξ2x1 : _ξ3x1
        _ζx2 = dim == 2 ? _ξ2x2 : _ξ3x2
        _ζx3 = dim == 2 ? _ξ2x3 : _ξ3x3
    end

    s_G = @localmem FT (Nq, Nq, Nqk, ngradstate)
    s_D = @localmem FT (Nq, Nq)

    l_Q = @private FT (ngradtransformstate,)
    l_aux = @private FT (nauxstate,)

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)
    ijk = @index(Local, Linear)

    @inbounds @views begin
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradtransformstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        fill!(l_G, -zero(eltype(l_G)))
        gradvariables!(
            bl,
            Vars{vars_gradient(bl, FT)}(l_G),
            Vars{vars_state(bl, FT)}(l_Q[:]),
            Vars{vars_aux(bl, FT)}(l_aux[:]),
            t,
        )
        @unroll for s in 1:ngradstate
            s_G[i, j, k, s] = l_G[s]
        end
        @synchronize

        # Compute gradient of each state
        ζx1 = vgeo[ijk, _ζx1, e]
        ζx2 = vgeo[ijk, _ζx2, e]
        ζx3 = vgeo[ijk, _ζx3, e]

        @unroll for s in 1:ngradstate
            Gζ = zero(FT)
            @unroll for n in 1:Nq
                if dim == 2
                    Gζ += s_D[j, n] * s_G[i, n, k, s]
                elseif dim == 3
                    Gζ += s_D[k, n] * s_G[i, j, n, s]
                end
            end
            l_gradG[1, s] = ζx1 * Gζ
            l_gradG[2, s] = ζx2 * Gζ
            l_gradG[3, s] = ζx3 * Gζ
        end

        @unroll for s in 1:ngradlapstate
            Qhypervisc_grad[ijk, 3 * (s - 1) + 1, e] =
                l_gradG[1, hypervisc_indexmap[s]]
            Qhypervisc_grad[ijk, 3 * (s - 1) + 2, e] =
                l_gradG[2, hypervisc_indexmap[s]]
            Qhypervisc_grad[ijk, 3 * (s - 1) + 3, e] =
                l_gradG[3, hypervisc_indexmap[s]]
        end

        if nviscstate > 0
            fill!(l_Qvisc, -zero(eltype(l_Qvisc)))
            diffusive!(
                bl,
                Vars{vars_diffusive(bl, FT)}(l_Qvisc),
                Grad{vars_gradient(bl, FT)}(l_gradG),
                Vars{vars_state(bl, FT)}(l_Q[:]),
                Vars{vars_aux(bl, FT)}(l_aux[:]),
                t,
            )
        end

        @unroll for s in 1:nviscstate
            Qvisc[ijk, s, e] = l_Qvisc[s]
        end
        @synchronize
    end
end

@kernel function faceviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    gradnumflux::NumericalFluxGradient,
    Q,
    Qvisc,
    Qhypervisc_grad,
    auxstate,
    vgeo,
    sgeo,
    t,
    vmap⁻,
    vmap⁺,
    elemtobndy,
    hypervisc_indexmap,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        ngradstate = num_gradient(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)

        if dim == 1
            Np = (N + 1)
            Nfp = 1
            nface = 2
        elseif dim == 2
            Np = (N + 1) * (N + 1)
            Nfp = (N + 1)
            nface = 4
        elseif dim == 3
            Np = (N + 1) * (N + 1) * (N + 1)
            Nfp = (N + 1) * (N + 1)
            nface = 6
        end

        faces = 1:nface
        if direction isa VerticalDirection
            faces = (nface - 1):nface
        elseif direction isa HorizontalDirection
            faces = 1:(nface - 2)
        end

        Nqk = dim == 2 ? 1 : N + 1

        ngradtransformstate = nstate

        l_Q⁻ = MArray{Tuple{ngradtransformstate}, FT}(undef)
        l_aux⁻ = MArray{Tuple{nauxstate}, FT}(undef)
        l_G⁻ = MArray{Tuple{ngradstate}, FT}(undef)
        l_nG⁻ = MArray{Tuple{3, ngradstate}, FT}(undef)

        l_Q⁺ = MArray{Tuple{ngradtransformstate}, FT}(undef)
        l_aux⁺ = MArray{Tuple{nauxstate}, FT}(undef)
        l_G⁺ = MArray{Tuple{ngradstate}, FT}(undef)

        # FIXME Qvisc is sort of a terrible name...
        l_Qvisc = MArray{Tuple{nviscstate}, FT}(undef)
        l_gradG = MArray{Tuple{3, ngradstate}, FT}(undef)
        l_Q⁻visc = MArray{Tuple{nviscstate}, FT}(undef)

        l_Q_bot1 = MArray{Tuple{nstate}, FT}(undef)
        l_aux_bot1 = MArray{Tuple{nauxstate}, FT}(undef)
    end

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    e = @private Int (1,)
    @inbounds e[1] = elems[eI]

    @inbounds for f in faces
        e⁻ = e[1]
        n⁻ = SVector(
            sgeo[_n1, n, f, e⁻],
            sgeo[_n2, n, f, e⁻],
            sgeo[_n3, n, f, e⁻],
        )
        sM, vMI = sgeo[_sM, n, f, e⁻], sgeo[_vMI, n, f, e⁻]
        id⁻, id⁺ = vmap⁻[n, f, e⁻], vmap⁺[n, f, e⁻]
        e⁺ = ((id⁺ - 1) ÷ Np) + 1

        vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

        # Load minus side data
        @unroll for s in 1:ngradtransformstate
            l_Q⁻[s] = Q[vid⁻, s, e⁻]
        end

        @unroll for s in 1:nauxstate
            l_aux⁻[s] = auxstate[vid⁻, s, e⁻]
        end

        fill!(l_G⁻, -zero(eltype(l_G⁻)))
        gradvariables!(
            bl,
            Vars{vars_gradient(bl, FT)}(l_G⁻),
            Vars{vars_state(bl, FT)}(l_Q⁻),
            Vars{vars_aux(bl, FT)}(l_aux⁻),
            t,
        )

        # Load plus side data
        @unroll for s in 1:ngradtransformstate
            l_Q⁺[s] = Q[vid⁺, s, e⁺]
        end

        @unroll for s in 1:nauxstate
            l_aux⁺[s] = auxstate[vid⁺, s, e⁺]
        end

        fill!(l_G⁺, -zero(eltype(l_G⁺)))
        gradvariables!(
            bl,
            Vars{vars_gradient(bl, FT)}(l_G⁺),
            Vars{vars_state(bl, FT)}(l_Q⁺),
            Vars{vars_aux(bl, FT)}(l_aux⁺),
            t,
        )

        bctype = elemtobndy[f, e⁻]
        fill!(l_Qvisc, -zero(eltype(l_Qvisc)))
        if bctype == 0
            numerical_flux_gradient!(
                gradnumflux,
                bl,
                l_gradG,
                SVector(n⁻),
                Vars{vars_gradient(bl, FT)}(l_G⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_gradient(bl, FT)}(l_G⁺),
                Vars{vars_state(bl, FT)}(l_Q⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺),
                t,
            )
            if nviscstate > 0
                diffusive!(
                    bl,
                    Vars{vars_diffusive(bl, FT)}(l_Qvisc),
                    Grad{vars_gradient(bl, FT)}(l_gradG),
                    Vars{vars_state(bl, FT)}(l_Q⁻),
                    Vars{vars_aux(bl, FT)}(l_aux⁻),
                    t,
                )
            end
        else
            if (dim == 2 && f == 3) || (dim == 3 && f == 5)
                # Loop up the first element along all horizontal elements
                @unroll for s in 1:nstate
                    l_Q_bot1[s] = Q[n + Nqk^2, s, e⁻]
                end
                @unroll for s in 1:nauxstate
                    l_aux_bot1[s] = auxstate[n + Nqk^2, s, e⁻]
                end
            end
            numerical_boundary_flux_gradient!(
                gradnumflux,
                bl,
                l_gradG,
                SVector(n⁻),
                Vars{vars_gradient(bl, FT)}(l_G⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_gradient(bl, FT)}(l_G⁺),
                Vars{vars_state(bl, FT)}(l_Q⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺),
                bctype,
                t,
                Vars{vars_state(bl, FT)}(l_Q_bot1),
                Vars{vars_aux(bl, FT)}(l_aux_bot1),
            )
            if nviscstate > 0
                diffusive!(
                    bl,
                    Vars{vars_diffusive(bl, FT)}(l_Qvisc),
                    Grad{vars_gradient(bl, FT)}(l_gradG),
                    Vars{vars_state(bl, FT)}(l_Q⁻),
                    Vars{vars_aux(bl, FT)}(l_aux⁻),
                    t,
                )
            end
        end

        @unroll for j in 1:ngradstate
            @unroll for i in 1:3
                l_nG⁻[i, j] = n⁻[i] * l_G⁻[j]
            end
        end

        @unroll for s in 1:ngradlapstate
            j = hypervisc_indexmap[s]
            Qhypervisc_grad[vid⁻, 3 * (s - 1) + 1, e⁻] +=
                vMI * sM * (l_gradG[1, j] - l_nG⁻[1, j])
            Qhypervisc_grad[vid⁻, 3 * (s - 1) + 2, e⁻] +=
                vMI * sM * (l_gradG[2, j] - l_nG⁻[2, j])
            Qhypervisc_grad[vid⁻, 3 * (s - 1) + 3, e⁻] +=
                vMI * sM * (l_gradG[3, j] - l_nG⁻[3, j])
        end

        diffusive!(
            bl,
            Vars{vars_diffusive(bl, FT)}(l_Q⁻visc),
            Grad{vars_gradient(bl, FT)}(l_nG⁻),
            Vars{vars_state(bl, FT)}(l_Q⁻),
            Vars{vars_aux(bl, FT)}(l_aux⁻),
            t,
        )

        @unroll for s in 1:nviscstate
            Qvisc[vid⁻, s, e⁻] += vMI * sM * (l_Qvisc[s] - l_Q⁻visc[s])
        end
        # Need to wait after even faces to avoid race conditions
        @synchronize(f % 2 == 0)
    end
end

@kernel function initstate!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    state,
    auxstate,
    vgeo,
    elems,
    args...,
) where {dim, polyorder}
    N = polyorder
    FT = eltype(auxstate)
    nauxstate = num_aux(bl, FT)
    nstate = num_state(bl, FT)

    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq
    Np = Nq * Nq * Nqk

    l_state = MArray{Tuple{nstate}, FT}(undef)
    l_aux = MArray{Tuple{nauxstate}, FT}(undef)

    e = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        coords = SVector(vgeo[n, _x1, e], vgeo[n, _x2, e], vgeo[n, _x3, e])
        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[n, s, e]
        end
        @unroll for s in 1:nstate
            l_state[s] = state[n, s, e]
        end
        init_state!(
            bl,
            Vars{vars_state(bl, FT)}(l_state),
            Vars{vars_aux(bl, FT)}(l_aux),
            coords,
            args...,
        )
        @unroll for s in 1:nstate
            state[n, s, e] = l_state[s]
        end
    end
end

@doc """
    initauxstate!(bl::BalanceLaw, Val(polyorder), auxstate, vgeo, elems)

Computational kernel: Initialize the auxiliary state

See [`DGBalanceLaw`](@ref) for usage.
""" initauxstate!
@kernel function initauxstate!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    auxstate,
    vgeo,
    elems,
) where {dim, polyorder}
    N = polyorder
    FT = eltype(auxstate)
    nauxstate = num_aux(bl, FT)

    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq
    Np = Nq * Nq * Nqk

    l_aux = MArray{Tuple{nauxstate}, FT}(undef)

    e = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[n, s, e]
        end

        init_aux!(
            bl,
            Vars{vars_aux(bl, FT)}(l_aux),
            LocalGeometry(Val(polyorder), vgeo, n, e),
        )

        @unroll for s in 1:nauxstate
            auxstate[n, s, e] = l_aux[s]
        end
    end
end

@doc """
    knl_nodal_update_aux!(bl::BalanceLaw, ::Val{dim}, ::Val{N}, f!, Q, auxstate,
                          t, elems, activedofs) where {dim, N}

Update the auxiliary state array
""" knl_nodal_update_aux!
@kernel function knl_nodal_update_aux!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{N},
    f!,
    Q,
    auxstate,
    t,
    elems,
    activedofs,
) where {dim, N}
    FT = eltype(Q)
    nstate = num_state(bl, FT)
    nauxstate = num_aux(bl, FT)

    Nq = N + 1

    Nqk = dim == 2 ? 1 : Nq

    Np = Nq * Nq * Nqk

    l_Q = MArray{Tuple{nstate}, FT}(undef)
    l_aux = MArray{Tuple{nauxstate}, FT}(undef)

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        e = elems[eI]

        active = activedofs[n + (e - 1) * Np]

        if active
            @unroll for s in 1:nstate
                l_Q[s] = Q[n, s, e]
            end

            @unroll for s in 1:nauxstate
                l_aux[s] = auxstate[n, s, e]
            end

            f!(
                bl,
                Vars{vars_state(bl, FT)}(l_Q),
                Vars{vars_aux(bl, FT)}(l_aux),
                t,
            )

            @unroll for s in 1:nauxstate
                auxstate[n, s, e] = l_aux[s]
            end
        end
    end
end

@doc """
    knl_nodal_update_aux!(bl::BalanceLaw, ::Val{dim}, ::Val{N}, f!, Q, auxstate, diffstate,
                          t, elems, activedofs) where {dim, N}

Update the auxiliary state array
""" knl_nodal_update_aux!
@kernel function knl_nodal_update_aux!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{N},
    f!,
    Q,
    auxstate,
    diffstate,
    t,
    elems,
    activedofs,
) where {dim, N}
    FT = eltype(Q)
    nstate = num_state(bl, FT)
    nviscstate = num_diffusive(bl, FT)
    nauxstate = num_aux(bl, FT)

    Nq = N + 1

    Nqk = dim == 2 ? 1 : Nq

    Np = Nq * Nq * Nqk

    l_Q = MArray{Tuple{nstate}, FT}(undef)
    l_aux = MArray{Tuple{nauxstate}, FT}(undef)
    l_diff = MArray{Tuple{nviscstate}, FT}(undef)

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        e = elems[eI]

        active = activedofs[n + (e - 1) * Np]

        if active
            @unroll for s in 1:nstate
                l_Q[s] = Q[n, s, e]
            end

            @unroll for s in 1:nauxstate
                l_aux[s] = auxstate[n, s, e]
            end

            @unroll for s in 1:nviscstate
                l_diff[s] = diffstate[n, s, e]
            end

            f!(
                bl,
                Vars{vars_state(bl, FT)}(l_Q),
                Vars{vars_aux(bl, FT)}(l_aux),
                Vars{vars_diffusive(bl, FT)}(l_diff),
                t,
            )

            @unroll for s in 1:nauxstate
                auxstate[n, s, e] = l_aux[s]
            end
        end
    end
end

@doc """
    knl_indefinite_stack_integral!(bl::BalanceLaw, ::Val{dim}, ::Val{N},
                                  ::Val{nvertelem}, Q, auxstate, vgeo,
                                  Imat, elems) where {dim, N, nvertelem}
Computational kernel: compute indefinite integral along the vertical stack
See [`DGBalanceLaw`](@ref) for usage.
""" knl_indefinite_stack_integral!
@kernel function knl_indefinite_stack_integral!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{N},
    ::Val{nvertelem},
    Q,
    auxstate,
    vgeo,
    Imat,
    elems,
) where {dim, N, nvertelem}
    @uniform begin
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        nauxstate = num_aux(bl, FT)
        nout = num_integrals(bl, FT)

        Nq = N + 1
        Nqj = dim == 2 ? 1 : Nq

        l_Q = MArray{Tuple{nstate}, FT}(undef)
        l_aux = MArray{Tuple{nauxstate}, FT}(undef)
        l_knl = MArray{Tuple{nout, Nq}, FT}(undef)
    end

    l_int = @private FT (nout, Nq)
    s_I = @localmem FT (Nq, Nq)

    _eh = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    @inbounds begin
        @unroll for n in 1:Nq
            s_I[i, n] = Imat[i, n]
        end
        @synchronize

        # Initialize the constant state at zero
        @unroll for k in 1:Nq
            @unroll for s in 1:nout
                l_int[s, k] = 0
            end
        end

        eh = elems[_eh]

        # Loop up the stack of elements
        for ev in 1:nvertelem
            e = ev + (eh - 1) * nvertelem

            # Evaluate the integral kernel at each DOF in the slabk
            # loop up the pencil
            @unroll for k in 1:Nq
                ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                Jc = vgeo[ijk, _JcV, e]
                @unroll for s in 1:nstate
                    l_Q[s] = Q[ijk, s, e]
                end

                @unroll for s in 1:nauxstate
                    l_aux[s] = auxstate[ijk, s, e]
                end

                integral_load_aux!(
                    bl,
                    Vars{vars_integrals(bl, FT)}(view(l_knl, :, k)),
                    Vars{vars_state(bl, FT)}(l_Q),
                    Vars{vars_aux(bl, FT)}(l_aux),
                )

                # multiply in the curve jacobian
                @unroll for s in 1:nout
                    l_knl[s, k] *= Jc
                end
            end

            # Evaluate the integral up the element
            @unroll for s in 1:nout
                @unroll for k in 1:Nq
                    @unroll for n in 1:Nq
                        l_int[s, k] += s_I[k, n] * l_knl[s, n]
                    end
                end
            end

            # Store out to memory and reset the background value for next element
            @unroll for k in 1:Nq
                @unroll for s in 1:nout
                    l_knl[s, k] = l_int[s, k]
                end
                ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                integral_set_aux!(
                    bl,
                    Vars{vars_aux(bl, FT)}(view(auxstate, ijk, :, e)),
                    Vars{vars_integrals(bl, FT)}(view(l_knl, :, k)),
                )
                @unroll for ind_out in 1:nout
                    l_int[ind_out, k] = l_int[ind_out, Nq]
                end
            end
        end
    end
end

@kernel function knl_reverse_indefinite_stack_integral!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{N},
    ::Val{nvertelem},
    state,
    auxstate,
    elems,
) where {dim, N, nvertelem}
    @uniform begin
        FT = eltype(auxstate)

        Nq = N + 1
        Nqj = dim == 2 ? 1 : Nq
        nout = num_reverse_integrals(bl, FT)

        # note that k is the second not 4th index (since this is scratch memory and k
        # needs to be persistent across threads)
        l_T = MArray{Tuple{nout}, FT}(undef)
        l_V = MArray{Tuple{nout}, FT}(undef)
    end

    _eh = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    @inbounds begin
        eh = elems[_eh]

        # Initialize the constant state at zero
        ijk = i + Nq * ((j - 1) + Nqj * (Nq - 1))
        et = nvertelem + (eh - 1) * nvertelem
        reverse_integral_load_aux!(
            bl,
            Vars{vars_reverse_integrals(bl, FT)}(l_T),
            Vars{vars_state(bl, FT)}(view(state, ijk, :, et)),
            Vars{vars_aux(bl, FT)}(view(auxstate, ijk, :, et)),
        )

        # Loop up the stack of elements
        for ev in 1:nvertelem
            e = ev + (eh - 1) * nvertelem
            @unroll for k in 1:Nq
                ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                reverse_integral_load_aux!(
                    bl,
                    Vars{vars_reverse_integrals(bl, FT)}(l_V),
                    Vars{vars_state(bl, FT)}(view(state, ijk, :, et)),
                    Vars{vars_aux(bl, FT)}(view(auxstate, ijk, :, e)),
                )
                l_V .= l_T .- l_V
                reverse_integral_set_aux!(
                    bl,
                    Vars{vars_aux(bl, FT)}(view(auxstate, ijk, :, e)),
                    Vars{vars_reverse_integrals(bl, FT)}(l_V),
                )
            end
        end
    end
end

# TODO: Generalize to more than one field?
@kernel function knl_copy_stack_field_down!(
    ::Val{dim},
    ::Val{N},
    ::Val{nvertelem},
    auxstate,
    elems,
    ::Val{fldin},
    ::Val{fldout},
) where {dim, N, nvertelem, fldin, fldout}
    DFloat = eltype(auxstate)

    Nq = N + 1
    Nqj = dim == 2 ? 1 : Nq

    _eh = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    # note that k is the second not 4th index (since this is scratch memory and k
    # needs to be persistent across threads)
    @inbounds begin
        # Initialize the constant state at zero
        ijk = i + Nq * ((j - 1) + Nqj * (Nq - 1))
        eh = elems[_eh]
        et = nvertelem + (eh - 1) * nvertelem
        val = auxstate[ijk, fldin, et]

        # Loop up the stack of elements
        for ev in 1:nvertelem
            e = ev + (eh - 1) * nvertelem
            @unroll for k in 1:Nq
                ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                auxstate[ijk, fldout, e] = val
            end
        end
    end
end

@kernel function volumedivgrad!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    Qhypervisc_grad,
    Qhypervisc_div,
    vgeo,
    D,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Qhypervisc_grad)
        ngradlapstate = num_gradient_laplacian(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        l_div = MArray{Tuple{ngradlapstate}, FT}(undef)
    end

    s_grad = @localmem FT (Nq, Nq, Nqk, ngradlapstate, 3)
    s_D = @localmem FT (Nq, Nq)

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)
    ijk = @index(Local, Linear)

    @inbounds begin
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradlapstate
            s_grad[i, j, k, s, 1] = Qhypervisc_grad[ijk, 3 * (s - 1) + 1, e]
            s_grad[i, j, k, s, 2] = Qhypervisc_grad[ijk, 3 * (s - 1) + 2, e]
            s_grad[i, j, k, s, 3] = Qhypervisc_grad[ijk, 3 * (s - 1) + 3, e]
        end
        @synchronize

        ξ1x1, ξ1x2, ξ1x3 =
            vgeo[ijk, _ξ1x1, e], vgeo[ijk, _ξ1x2, e], vgeo[ijk, _ξ1x3, e]
        if dim == 3 || (dim == 2 && direction isa EveryDirection)
            ξ2x1, ξ2x2, ξ2x3 =
                vgeo[ijk, _ξ2x1, e], vgeo[ijk, _ξ2x2, e], vgeo[ijk, _ξ2x3, e]
        end
        if dim == 3 && direction isa EveryDirection
            ξ3x1, ξ3x2, ξ3x3 =
                vgeo[ijk, _ξ3x1, e], vgeo[ijk, _ξ3x2, e], vgeo[ijk, _ξ3x3, e]
        end

        @unroll for s in 1:ngradlapstate
            g1ξ1 = g1ξ2 = g1ξ3 = zero(FT)
            g2ξ1 = g2ξ2 = g2ξ3 = zero(FT)
            g3ξ1 = g3ξ2 = g3ξ3 = zero(FT)
            @unroll for n in 1:Nq
                Din = s_D[i, n]
                g1ξ1 += Din * s_grad[n, j, k, s, 1]
                g2ξ1 += Din * s_grad[n, j, k, s, 2]
                g3ξ1 += Din * s_grad[n, j, k, s, 3]
                if dim == 3 || (dim == 2 && direction isa EveryDirection)
                    Djn = s_D[j, n]
                    g1ξ2 += Djn * s_grad[i, n, k, s, 1]
                    g2ξ2 += Djn * s_grad[i, n, k, s, 2]
                    g3ξ2 += Djn * s_grad[i, n, k, s, 3]
                end
                if dim == 3 && direction isa EveryDirection
                    Dkn = s_D[k, n]
                    g1ξ3 += Dkn * s_grad[i, j, n, s, 1]
                    g2ξ3 += Dkn * s_grad[i, j, n, s, 2]
                    g3ξ3 += Dkn * s_grad[i, j, n, s, 3]
                end
            end
            l_div[s] = ξ1x1 * g1ξ1 + ξ1x2 * g2ξ1 + ξ1x3 * g3ξ1

            if dim == 3 || (dim == 2 && direction isa EveryDirection)
                l_div[s] += ξ2x1 * g1ξ2 + ξ2x2 * g2ξ2 + ξ2x3 * g3ξ2
            end

            if dim == 3 && direction isa EveryDirection
                l_div[s] += ξ3x1 * g1ξ3 + ξ3x2 * g2ξ3 + ξ3x3 * g3ξ3
            end
        end

        @unroll for s in 1:ngradlapstate
            Qhypervisc_div[ijk, s, e] = l_div[s]
        end
        @synchronize
    end
end

@kernel function volumedivgrad!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    ::VerticalDirection,
    Qhypervisc_grad,
    Qhypervisc_div,
    vgeo,
    D,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Qhypervisc_grad)
        ngradlapstate = num_gradient_laplacian(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        l_div = MArray{Tuple{ngradlapstate}, FT}(undef)
    end

    s_grad = @localmem FT (Nq, Nq, Nqk, ngradlapstate, 3)
    s_D = @localmem FT (Nq, Nq)

    e = @index(Group, Linear)
    ijk = @index(Local, Linear)
    i, j, k = @index(Local, NTuple)
    @inbounds begin
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradlapstate
            s_grad[i, j, k, s, 1] = Qhypervisc_grad[ijk, 3 * (s - 1) + 1, e]
            s_grad[i, j, k, s, 2] = Qhypervisc_grad[ijk, 3 * (s - 1) + 2, e]
            s_grad[i, j, k, s, 3] = Qhypervisc_grad[ijk, 3 * (s - 1) + 3, e]
        end
        @synchronize

        if dim == 2
            ξ2x1, ξ2x2, ξ2x3 =
                vgeo[ijk, _ξ2x1, e], vgeo[ijk, _ξ2x2, e], vgeo[ijk, _ξ2x3, e]
        else
            ξ3x1, ξ3x2, ξ3x3 =
                vgeo[ijk, _ξ3x1, e], vgeo[ijk, _ξ3x2, e], vgeo[ijk, _ξ3x3, e]
        end

        @unroll for s in 1:ngradlapstate
            g1ξv = g2ξv = g3ξv = zero(FT)
            @unroll for n in 1:Nq
                if dim == 2
                    Djn = s_D[j, n]
                    g1ξv += Djn * s_grad[i, n, k, s, 1]
                    g2ξv += Djn * s_grad[i, n, k, s, 2]
                    g3ξv += Djn * s_grad[i, n, k, s, 3]
                else
                    Dkn = s_D[k, n]
                    g1ξv += Dkn * s_grad[i, j, n, s, 1]
                    g2ξv += Dkn * s_grad[i, j, n, s, 2]
                    g3ξv += Dkn * s_grad[i, j, n, s, 3]
                end
            end

            if dim == 2
                l_div[s] = ξ2x1 * g1ξv + ξ2x2 * g2ξv + ξ2x3 * g3ξv
            else
                l_div[s] = ξ3x1 * g1ξv + ξ3x2 * g2ξv + ξ3x3 * g3ξv
            end
        end

        @unroll for s in 1:ngradlapstate
            Qhypervisc_div[ijk, s, e] = l_div[s]
        end
        @synchronize
    end
end

@kernel function facedivgrad!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    divgradnumpenalty,
    Qhypervisc_grad,
    Qhypervisc_div,
    vgeo,
    sgeo,
    vmap⁻,
    vmap⁺,
    elemtobndy,
    elems,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Qhypervisc_grad)
        ngradlapstate = num_gradient_laplacian(bl, FT)

        if dim == 1
            Np = (N + 1)
            Nfp = 1
            nface = 2
        elseif dim == 2
            Np = (N + 1) * (N + 1)
            Nfp = (N + 1)
            nface = 4
        elseif dim == 3
            Np = (N + 1) * (N + 1) * (N + 1)
            Nfp = (N + 1) * (N + 1)
            nface = 6
        end

        faces = 1:nface
        if direction isa VerticalDirection
            faces = (nface - 1):nface
        elseif direction isa HorizontalDirection
            faces = 1:(nface - 2)
        end

        Nqk = dim == 2 ? 1 : N + 1

        l_grad⁻ = MArray{Tuple{3, ngradlapstate}, FT}(undef)
        l_grad⁺ = MArray{Tuple{3, ngradlapstate}, FT}(undef)
        l_div = MArray{Tuple{ngradlapstate}, FT}(undef)
    end

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    e = @private Int (1,)
    @inbounds e[1] = elems[eI]

    @inbounds for f in faces
        e⁻ = e[1]
        n⁻ = SVector(
            sgeo[_n1, n, f, e⁻],
            sgeo[_n2, n, f, e⁻],
            sgeo[_n3, n, f, e⁻],
        )
        sM, vMI = sgeo[_sM, n, f, e⁻], sgeo[_vMI, n, f, e⁻]
        id⁻, id⁺ = vmap⁻[n, f, e⁻], vmap⁺[n, f, e⁻]
        e⁺ = ((id⁺ - 1) ÷ Np) + 1

        vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

        # Load minus side data
        @unroll for s in 1:ngradlapstate
            l_grad⁻[1, s] = Qhypervisc_grad[vid⁻, 3 * (s - 1) + 1, e⁻]
            l_grad⁻[2, s] = Qhypervisc_grad[vid⁻, 3 * (s - 1) + 2, e⁻]
            l_grad⁻[3, s] = Qhypervisc_grad[vid⁻, 3 * (s - 1) + 3, e⁻]
        end

        # Load plus side data
        @unroll for s in 1:ngradlapstate
            l_grad⁺[1, s] = Qhypervisc_grad[vid⁺, 3 * (s - 1) + 1, e⁺]
            l_grad⁺[2, s] = Qhypervisc_grad[vid⁺, 3 * (s - 1) + 2, e⁺]
            l_grad⁺[3, s] = Qhypervisc_grad[vid⁺, 3 * (s - 1) + 3, e⁺]
        end

        bctype = elemtobndy[f, e⁻]
        if bctype == 0
            divergence_penalty!(
                divgradnumpenalty,
                bl,
                Vars{vars_gradient_laplacian(bl, FT)}(l_div),
                n⁻,
                Grad{vars_gradient_laplacian(bl, FT)}(l_grad⁻),
                Grad{vars_gradient_laplacian(bl, FT)}(l_grad⁺),
            )
        else
            divergence_boundary_penalty!(
                divgradnumpenalty,
                bl,
                Vars{vars_gradient_laplacian(bl, FT)}(l_div),
                n⁻,
                Grad{vars_gradient_laplacian(bl, FT)}(l_grad⁻),
                Grad{vars_gradient_laplacian(bl, FT)}(l_grad⁺),
                bctype,
            )
        end

        @unroll for s in 1:ngradlapstate
            Qhypervisc_div[vid⁻, s, e⁻] += vMI * sM * l_div[s]
        end
        # Need to wait after even faces to avoid race conditions
        @synchronize(f % 2 == 0)
    end
end

@kernel function volumehyperviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    Qhypervisc_grad,
    Qhypervisc_div,
    Q,
    auxstate,
    vgeo,
    ω,
    D,
    elems,
    t,
) where {dim, polyorder}
    @uniform begin
        N = polyorder

        FT = eltype(Qhypervisc_grad)
        nstate = num_state(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)
        nauxstate = num_aux(bl, FT)
        ngradtransformstate = nstate

        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq

        l_grad_lap = MArray{Tuple{3, ngradlapstate}, FT}(undef)
        l_Qhypervisc = MArray{Tuple{nhyperviscstate}, FT}(undef)
    end

    s_lap = @localmem FT (Nq, Nq, Nqk, ngradlapstate)
    s_D = @localmem FT (Nq, Nq)
    s_ω = @localmem FT (Nq,)
    l_Q = @private FT (ngradtransformstate,)
    l_aux = @private FT (nauxstate,)

    e = @index(Group, Linear)
    ijk = @index(Local, Linear)
    i, j, k = @index(Local, NTuple)

    @inbounds @views begin
        s_ω[j] = ω[j]
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradtransformstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        @unroll for s in 1:ngradlapstate
            s_lap[i, j, k, s] = Qhypervisc_div[ijk, s, e]
        end
        @synchronize

        ξ1x1, ξ1x2, ξ1x3 =
            vgeo[ijk, _ξ1x1, e], vgeo[ijk, _ξ1x2, e], vgeo[ijk, _ξ1x3, e]
        if dim == 3 || (dim == 2 && direction isa EveryDirection)
            ξ2x1, ξ2x2, ξ2x3 =
                vgeo[ijk, _ξ2x1, e], vgeo[ijk, _ξ2x2, e], vgeo[ijk, _ξ2x3, e]
        end
        if dim == 3 && direction isa EveryDirection
            ξ3x1, ξ3x2, ξ3x3 =
                vgeo[ijk, _ξ3x1, e], vgeo[ijk, _ξ3x2, e], vgeo[ijk, _ξ3x3, e]
        end
        @unroll for s in 1:ngradlapstate
            lap_ξ1 = lap_ξ2 = lap_ξ3 = zero(FT)
            @unroll for n in 1:Nq
                njk = n + Nq * ((j - 1) + Nq * (k - 1))
                Dni = s_D[n, i] * s_ω[n] / s_ω[i]
                lap_njk = s_lap[n, j, k, s]
                lap_ξ1 += Dni * lap_njk
                if dim == 3 || (dim == 2 && direction isa EveryDirection)
                    ink = i + Nq * ((n - 1) + Nq * (k - 1))
                    Dnj = s_D[n, j] * s_ω[n] / s_ω[j]
                    lap_ink = s_lap[i, n, k, s]
                    lap_ξ2 += Dnj * lap_ink
                end
                if dim == 3 && direction isa EveryDirection
                    ijn = i + Nq * ((j - 1) + Nq * (n - 1))
                    Dnk = s_D[n, k] * s_ω[n] / s_ω[k]
                    lap_ijn = s_lap[i, j, n, s]
                    lap_ξ3 += Dnk * lap_ijn
                end
            end

            l_grad_lap[1, s] = -ξ1x1 * lap_ξ1
            l_grad_lap[2, s] = -ξ1x2 * lap_ξ1
            l_grad_lap[3, s] = -ξ1x3 * lap_ξ1

            if dim == 3 || (dim == 2 && direction isa EveryDirection)
                l_grad_lap[1, s] -= ξ2x1 * lap_ξ2
                l_grad_lap[2, s] -= ξ2x2 * lap_ξ2
                l_grad_lap[3, s] -= ξ2x3 * lap_ξ2
            end

            if dim == 3 && direction isa EveryDirection
                l_grad_lap[1, s] -= ξ3x1 * lap_ξ3
                l_grad_lap[2, s] -= ξ3x2 * lap_ξ3
                l_grad_lap[3, s] -= ξ3x3 * lap_ξ3
            end
        end

        fill!(l_Qhypervisc, -zero(eltype(l_Qhypervisc)))
        hyperdiffusive!(
            bl,
            Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
            Grad{vars_gradient_laplacian(bl, FT)}(l_grad_lap),
            Vars{vars_state(bl, FT)}(l_Q[:]),
            Vars{vars_aux(bl, FT)}(l_aux[:]),
            t,
        )
        @unroll for s in 1:nhyperviscstate
            Qhypervisc_grad[ijk, s, e] = l_Qhypervisc[s]
        end
        @synchronize
    end
end

@kernel function volumehyperviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    ::VerticalDirection,
    Qhypervisc_grad,
    Qhypervisc_div,
    Q,
    auxstate,
    vgeo,
    ω,
    D,
    elems,
    t,
) where {dim, polyorder}
    @uniform begin
        N = polyorder

        FT = eltype(Qhypervisc_grad)
        nstate = num_state(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)
        nauxstate = num_aux(bl, FT)
        ngradtransformstate = nstate

        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq

        l_grad_lap = MArray{Tuple{3, ngradlapstate}, FT}(undef)
        l_Qhypervisc = MArray{Tuple{nhyperviscstate}, FT}(undef)
    end

    s_lap = @localmem FT (Nq, Nq, Nqk, ngradlapstate)
    s_D = @localmem FT (Nq, Nq)
    s_ω = @localmem FT (Nq,)
    l_Q = @private FT (ngradtransformstate,)
    l_aux = @private FT (nauxstate,)

    e = @index(Group, Linear)
    ijk = @index(Local, Linear)
    i, j, k = @index(Local, NTuple)

    @inbounds @views begin
        s_ω[j] = ω[j]
        s_D[i, j] = D[i, j]

        @unroll for s in 1:ngradtransformstate
            l_Q[s] = Q[ijk, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[ijk, s, e]
        end

        @unroll for s in 1:ngradlapstate
            s_lap[i, j, k, s] = Qhypervisc_div[ijk, s, e]
        end
        @synchronize

        if dim == 2
            ξvx1, ξvx2, ξvx3 =
                vgeo[ijk, _ξ2x1, e], vgeo[ijk, _ξ2x2, e], vgeo[ijk, _ξ2x3, e]
        else
            ξvx1, ξvx2, ξvx3 =
                vgeo[ijk, _ξ3x1, e], vgeo[ijk, _ξ3x2, e], vgeo[ijk, _ξ3x3, e]
        end
        @unroll for s in 1:ngradlapstate
            lap_ξv = zero(FT)
            @unroll for n in 1:Nq
                if dim == 2
                    ink = i + Nq * ((n - 1) + Nq * (k - 1))
                    Dnj = s_D[n, j] * s_ω[n] / s_ω[j]
                    lap_ink = s_lap[i, n, k, s]
                    lap_ξv += Dnj * lap_ink
                else
                    ijn = i + Nq * ((j - 1) + Nq * (n - 1))
                    Dnk = s_D[n, k] * s_ω[n] / s_ω[k]
                    lap_ijn = s_lap[i, j, n, s]
                    lap_ξv += Dnk * lap_ijn
                end
            end

            l_grad_lap[1, s] = -ξvx1 * lap_ξv
            l_grad_lap[2, s] = -ξvx2 * lap_ξv
            l_grad_lap[3, s] = -ξvx3 * lap_ξv
        end

        fill!(l_Qhypervisc, -zero(eltype(l_Qhypervisc)))
        hyperdiffusive!(
            bl,
            Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
            Grad{vars_gradient_laplacian(bl, FT)}(l_grad_lap),
            Vars{vars_state(bl, FT)}(l_Q[:]),
            Vars{vars_aux(bl, FT)}(l_aux[:]),
            t,
        )
        @unroll for s in 1:nhyperviscstate
            Qhypervisc_grad[ijk, s, e] = l_Qhypervisc[s]
        end
        @synchronize
    end
end

@kernel function facehyperviscterms!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{polyorder},
    direction,
    hyperviscnumflux,
    Qhypervisc_grad,
    Qhypervisc_div,
    Q,
    auxstate,
    vgeo,
    sgeo,
    vmap⁻,
    vmap⁺,
    elemtobndy,
    elems,
    t,
) where {dim, polyorder}
    @uniform begin
        N = polyorder
        FT = eltype(Qhypervisc_grad)
        nstate = num_state(bl, FT)
        ngradlapstate = num_gradient_laplacian(bl, FT)
        nhyperviscstate = num_hyperdiffusive(bl, FT)
        nauxstate = num_aux(bl, FT)
        ngradtransformstate = nstate

        if dim == 1
            Np = (N + 1)
            Nfp = 1
            nface = 2
        elseif dim == 2
            Np = (N + 1) * (N + 1)
            Nfp = (N + 1)
            nface = 4
        elseif dim == 3
            Np = (N + 1) * (N + 1) * (N + 1)
            Nfp = (N + 1) * (N + 1)
            nface = 6
        end

        faces = 1:nface
        if direction isa VerticalDirection
            faces = (nface - 1):nface
        elseif direction isa HorizontalDirection
            faces = 1:(nface - 2)
        end

        Nqk = dim == 2 ? 1 : N + 1

        l_lap⁻ = MArray{Tuple{ngradlapstate}, FT}(undef)
        l_lap⁺ = MArray{Tuple{ngradlapstate}, FT}(undef)
        l_Qhypervisc = MArray{Tuple{nhyperviscstate}, FT}(undef)

        l_Q⁻ = MArray{Tuple{ngradtransformstate}, FT}(undef)
        l_aux⁻ = MArray{Tuple{nauxstate}, FT}(undef)

        l_Q⁺ = MArray{Tuple{ngradtransformstate}, FT}(undef)
        l_aux⁺ = MArray{Tuple{nauxstate}, FT}(undef)
    end

    eI = @index(Group, Linear)
    n = @index(Local, Linear)

    e = @private Int (1,)
    @inbounds e[1] = elems[eI]

    @inbounds for f in faces
        e⁻ = e[1]
        n⁻ = SVector(
            sgeo[_n1, n, f, e⁻],
            sgeo[_n2, n, f, e⁻],
            sgeo[_n3, n, f, e⁻],
        )
        sM, vMI = sgeo[_sM, n, f, e⁻], sgeo[_vMI, n, f, e⁻]
        id⁻, id⁺ = vmap⁻[n, f, e⁻], vmap⁺[n, f, e⁻]
        e⁺ = ((id⁺ - 1) ÷ Np) + 1

        vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

        # Load minus side data
        @unroll for s in 1:ngradtransformstate
            l_Q⁻[s] = Q[vid⁻, s, e⁻]
        end

        @unroll for s in 1:nauxstate
            l_aux⁻[s] = auxstate[vid⁻, s, e⁻]
        end

        @unroll for s in 1:ngradlapstate
            l_lap⁻[s] = Qhypervisc_div[vid⁻, s, e⁻]
        end

        # Load plus side data
        @unroll for s in 1:ngradtransformstate
            l_Q⁺[s] = Q[vid⁺, s, e⁺]
        end

        @unroll for s in 1:nauxstate
            l_aux⁺[s] = auxstate[vid⁺, s, e⁺]
        end

        @unroll for s in 1:ngradlapstate
            l_lap⁺[s] = Qhypervisc_div[vid⁺, s, e⁺]
        end

        bctype = elemtobndy[f, e⁻]
        if bctype == 0
            numerical_flux_hyperdiffusive!(
                hyperviscnumflux,
                bl,
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
                n⁻,
                Vars{vars_gradient_laplacian(bl, FT)}(l_lap⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_gradient_laplacian(bl, FT)}(l_lap⁺),
                Vars{vars_state(bl, FT)}(l_Q⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺),
                t,
            )
        else
            numerical_boundary_flux_hyperdiffusive!(
                hyperviscnumflux,
                bl,
                Vars{vars_hyperdiffusive(bl, FT)}(l_Qhypervisc),
                n⁻,
                Vars{vars_gradient_laplacian(bl, FT)}(l_lap⁻),
                Vars{vars_state(bl, FT)}(l_Q⁻),
                Vars{vars_aux(bl, FT)}(l_aux⁻),
                Vars{vars_gradient_laplacian(bl, FT)}(l_lap⁺),
                Vars{vars_state(bl, FT)}(l_Q⁺),
                Vars{vars_aux(bl, FT)}(l_aux⁺),
                bctype,
                t,
            )
        end

        @unroll for s in 1:nhyperviscstate
            Qhypervisc_grad[vid⁻, s, e⁻] += vMI * sM * l_Qhypervisc[s]
        end
        # Need to wait after even faces to avoid race conditions
        @synchronize(f % 2 == 0)
    end
end

@kernel function knl_local_courant!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{N},
    pointwise_courant,
    local_courant,
    Q,
    auxstate,
    diffstate,
    elems,
    Δt,
    simtime,
    direction,
) where {dim, N}
    @uniform begin
        FT = eltype(Q)
        nstate = num_state(bl, FT)
        nviscstate = num_diffusive(bl, FT)
        nauxstate = num_aux(bl, FT)

        Nq = N + 1

        Nqk = dim == 2 ? 1 : Nq

        Np = Nq * Nq * Nqk

        l_Q = MArray{Tuple{nstate}, FT}(undef)
        l_aux = MArray{Tuple{nauxstate}, FT}(undef)
        l_diff = MArray{Tuple{nviscstate}, FT}(undef)
    end

    e = @index(Group, Linear)
    n = @index(Local, Linear)
    @inbounds begin
        @unroll for s in 1:nstate
            l_Q[s] = Q[n, s, e]
        end

        @unroll for s in 1:nauxstate
            l_aux[s] = auxstate[n, s, e]
        end

        @unroll for s in 1:nviscstate
            l_diff[s] = diffstate[n, s, e]
        end

        Δx = pointwise_courant[n, e]
        c = local_courant(
            bl,
            Vars{vars_state(bl, FT)}(l_Q),
            Vars{vars_aux(bl, FT)}(l_aux),
            Vars{vars_diffusive(bl, FT)}(l_diff),
            Δx,
            Δt,
            simtime,
            direction,
        )

        pointwise_courant[n, e] = c
    end
end

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
################################################################################
#
#  nmod_mat.jl: flint nmod_mat types in julia
#
################################################################################

export nmod_mat, NmodMatSpace, getindex, setindex!, set_entry!, deepcopy, rows,
       cols, parent, base_ring, zero, one, issquare, show, transpose,
       transpose!, rref, rref!, trace, det, rank, inv, solve, lufact,
       sub, hcat, vcat, Array, lift, lift!, MatrixSpace, check_parent,
       howell_form, howell_form!, strong_echelon_form, strong_echelon_form!

################################################################################
#
#  Data type and parent object methods
#
################################################################################

parent_type(::Type{nmod_mat}) = NmodMatSpace

elem_type(::Type{NmodMatSpace}) = nmod_mat

function check_parent(x::nmod_mat, y::nmod_mat)
  base_ring(x) != base_ring(y) && error("Residue rings must be equal")
  (cols(x) != cols(y)) && (rows(x) != rows(y)) &&
          error("Matrices have wrong dimensions")
  return nothing
end

size(x::nmod_mat) = tuple(x.r, x.c)

size(t::nmod_mat, d) = d <= 2 ? size(t)[d] : 1

issquare(a::nmod_mat) = (rows(a) == cols(a))

###############################################################################
#
#   Similar
#
###############################################################################

function similar(x::nmod_mat)
   z = nmod_mat(rows(x), cols(x), x.n)
   z.base_ring = x.base_ring
   return z
end

function similar(x::nmod_mat, r::Int, c::Int)
   z = nmod_mat(r, c, x.n)
   z.base_ring = x.base_ring
   return z
end

################################################################################
#
#  Manipulation
#
################################################################################

@inline function getindex(a::nmod_mat, i::Int, j::Int)
  @boundscheck Generic._checkbounds(a, i, j)
  u = ccall((:nmod_mat_get_entry, :libflint), UInt,
              (Ref{nmod_mat}, Int, Int), a, i - 1 , j - 1)
  return base_ring(a)(u)
end

#as above, but as a plain UInt
function getindex_raw(a::nmod_mat, i::Int, j::Int)
  return ccall((:nmod_mat_get_entry, :libflint), UInt,
                 (Ref{nmod_mat}, Int, Int), a, i - 1, j - 1)
end

@inline function setindex!(a::nmod_mat, u::UInt, i::Int, j::Int)
  @boundscheck Generic._checkbounds(a, i, j)
  set_entry!(a, i, j, u)
end

@inline function setindex!(a::nmod_mat, u::fmpz, i::Int, j::Int)
  @boundscheck Generic._checkbounds(a, i, j)
  set_entry!(a, i, j, u)
end

@inline function setindex!(a::nmod_mat, u::nmod, i::Int, j::Int)
  @boundscheck Generic._checkbounds(a, i, j)
  (base_ring(a) != parent(u)) && error("Parent objects must coincide")
  set_entry!(a, i, j, u.data)
end

setindex!(a::nmod_mat, u::Integer, i::Int, j::Int) =
        setindex!(a, fmpz(u), i, j)

setindex_t!(a::nmod_mat, u::T, i::Int, j::Int) where {T<:Union{RingElem, Integer}} =
  setindex!(a, u, j, i)

function set_entry!(a::nmod_mat, i::Int, j::Int, u::UInt)
  ccall((:nmod_mat_set_entry, :libflint), Void,
          (Ref{nmod_mat}, Int, Int, UInt), a, i - 1, j - 1, u)
end

function set_entry!(a::nmod_mat, i::Int, j::Int, u::fmpz)
  t = fmpz()
  ccall((:fmpz_mod_ui, :libflint), UInt,
          (Ref{fmpz}, Ref{fmpz}, UInt), t, u, a.n)
  tt = ccall((:fmpz_get_ui, :libflint), UInt, (Ref{fmpz}, ), t)
  set_entry!(a, i, j, tt)
end

set_entry!(a::nmod_mat, i::Int, j::Int, u::nmod) =
        set_entry!(a, i, j, u.data)

set_entry_t!(a::nmod_mat, i::Int, j::Int, u::T) where {T<:Union{RingElem, Integer}} =
  set_entry!(a, j, i, u)

function deepcopy_internal(a::nmod_mat, dict::ObjectIdDict)
  z = nmod_mat(rows(a), cols(a), a.n)
  if isdefined(a, :base_ring)
    z.base_ring = a.base_ring
  end
  ccall((:nmod_mat_set, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}), z, a)
  return z
end

rows(a::nmod_mat) = a.r

cols(a::nmod_mat) = a.c

parent(a::nmod_mat, cached::Bool = true) = MatrixSpace(base_ring(a), rows(a), cols(a), cached)

base_ring(a::NmodMatSpace) = a.base_ring

base_ring(a::nmod_mat) = a.base_ring

zero(a::NmodMatSpace) = a()

function one(a::NmodMatSpace)
  (a.rows != a.cols) && error("Matrices must be quadratic")
  z = a()
  ccall((:nmod_mat_one, :libflint), Void, (Ref{nmod_mat}, ), z)
  return z
end

function iszero(a::nmod_mat)
  r = ccall((:nmod_mat_is_zero, :libflint), Cint, (Ref{nmod_mat}, ), a)
  return Bool(r)
end

################################################################################
#
#  AbstractString I/O
#
################################################################################

function show(io::IO, a::NmodMatSpace)
   print(io, "Matrix Space of ")
   print(io, a.rows, " rows and ", a.cols, " columns over ")
   print(io, a.base_ring)
end

function show(io::IO, a::nmod_mat)
   rows = a.r
   cols = a.c
   for i = 1:rows
      print(io, "[")
      for j = 1:cols
         print(io, a[i, j])
         if j != cols
            print(io, " ")
         end
      end
      print(io, "]")
      if i != rows
         println(io, "")
      end
   end
end

################################################################################
#
#  Comparison
#
################################################################################

==(a::nmod_mat, b::nmod_mat) = (a.base_ring == b.base_ring) &&
        Bool(ccall((:nmod_mat_equal, :libflint), Cint,
                (Ref{nmod_mat}, Ref{nmod_mat}), a, b))

################################################################################
#
#  Transpose
#
################################################################################

function transpose(a::nmod_mat)
  z = NmodMatSpace(base_ring(a), cols(a), rows(a))()
  ccall((:nmod_mat_transpose, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}), z, a)
  return z
end

function transpose!(a::nmod_mat)
  !issquare(a) && error("Matrix must be a square matrix")
  ccall((:nmod_mat_transpose, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}), a, a)
end

################################################################################
#
#  Unary operators
#
################################################################################

function -(x::nmod_mat)
  z = similar(x)
  ccall((:nmod_mat_neg, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}), z, x)
  return z
end

################################################################################
#
#  Binary operators
#
################################################################################

function +(x::nmod_mat, y::nmod_mat)
  check_parent(x,y)
  z = similar(x)
  ccall((:nmod_mat_add, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  return z
end

function -(x::nmod_mat, y::nmod_mat)
  check_parent(x,y)
  z = similar(x)
  ccall((:nmod_mat_sub, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  return z
end

function *(x::nmod_mat, y::nmod_mat)
  (base_ring(x) != base_ring(y)) && error("Base ring must be equal")
  (cols(x) != rows(y)) && error("Dimensions are wrong")
  z = similar(x, rows(x), cols(y))
  ccall((:nmod_mat_mul, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  return z
end

################################################################################
#
#  Unsafe operations
#
################################################################################

function mul!(a::nmod_mat, b::nmod_mat, c::nmod_mat)
  ccall((:nmod_mat_mul, :libflint), Void, (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), a, b, c)
  return a
end

function add!(a::nmod_mat, b::nmod_mat, c::nmod_mat)
  ccall((:nmod_mat_add, :libflint), Void, (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), a, b, c)
  return a
end

function zero!(a::nmod_mat)
  ccall((:nmod_mat_zero, :libflint), Void, (Ref{nmod_mat}, ), a)
  return a
end

################################################################################
#
#  Ad hoc binary operators
#
################################################################################

function *(x::nmod_mat, y::UInt)
  z = similar(x)
  ccall((:nmod_mat_scalar_mul, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, UInt), z, x, y)
  return z
end

*(x::UInt, y::nmod_mat) = y*x

function *(x::nmod_mat, y::fmpz)
  t = fmpz()
  ccall((:fmpz_mod_ui, :libflint), UInt,
          (Ref{fmpz}, Ref{fmpz}, UInt), t, y, x.n)
  tt = ccall((:fmpz_get_ui, :libflint), UInt, (Ref{fmpz}, ), t)
  return x*tt
end

*(x::fmpz, y::nmod_mat) = y*x

function *(x::nmod_mat, y::Integer)
  return x*fmpz(y)
end

*(x::Integer, y::nmod_mat) = y*x

function *(x::nmod_mat, y::nmod)
  (base_ring(x) != parent(y)) && error("Parent objects must coincide")
  return x*y.data
end

*(x::Generic.Res{fmpz}, y::nmod_mat) = y*x

################################################################################
#
#  Powering
#
################################################################################

function ^(x::nmod_mat, y::UInt)
  z = similar(x)
  ccall((:nmod_mat_pow, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, UInt), z, x, y)
  return z
end

function ^(x::nmod_mat, y::Int)
  ( y < 0 ) && error("Exponent must be positive")
  return x^UInt(y)
end

function ^(x::nmod_mat, y::fmpz)
  ( y < 0 ) && error("Exponent must be positive")
  ( y > fmpz(typemax(UInt))) &&
          error("Exponent must be smaller then ", fmpz(typemax(UInt)))
  return x^(UInt(y))
end

################################################################################
#
#  Row echelon form
#
################################################################################

function rref(a::nmod_mat)
  z = deepcopy(a)
  ccall((:nmod_mat_rref, :libflint), Void, (Ref{nmod_mat}, ), z)
  return z
end

function rref!(a::nmod_mat)
  ccall((:nmod_mat_rref, :libflint), Void, (Ref{nmod_mat}, ), a)
  return a
end

################################################################################
#
#  Strong echelon form and Howell form
#
################################################################################

function strong_echelon_form!(a::nmod_mat)
  ccall((:nmod_mat_strong_echelon_form, :libflint), Void, (Ref{nmod_mat}, ), a)
end

doc"""
    strong_echelon_form(a::nmod_mat)
> Return the strong echeleon form of $a$. The matrix $a$ must have at least as
> many rows as columns.
"""
function strong_echelon_form(a::nmod_mat)
  (rows(a) < cols(a)) &&
              error("Matrix must have at least as many rows as columns")
  z = deepcopy(a)
  strong_echelon_form!(z)
  return z
end

function howell_form!(a::nmod_mat)
  ccall((:nmod_mat_howell_form, :libflint), Void, (Ref{nmod_mat}, ), a)
end

doc"""
    howell_form(a::nmod_mat)
> Return the Howell normal form of $a$. The matrix $a$ must have at least as
> many rows as columns.
"""
function howell_form(a::nmod_mat)
  (rows(a) < cols(a)) &&
              error("Matrix must have at least as many rows as columns")

  z = deepcopy(a)
  howell_form!(z)
  return z
end

################################################################################
#
#  Trace
#
################################################################################

function trace(a::nmod_mat)
  !issquare(a) && error("Matrix must be a square matrix")
  r = ccall((:nmod_mat_trace, :libflint), UInt, (Ref{nmod_mat}, ), a)
  return base_ring(a)(r)
end

################################################################################
#
#  Determinant
#
################################################################################

function det(a::nmod_mat)
  !issquare(a) && error("Matrix must be a square matrix")
  if is_prime(a.n)
     r = ccall((:nmod_mat_det, :libflint), UInt, (Ref{nmod_mat}, ), a)
     return base_ring(a)(r)
  else
     try
        return det_fflu(a)
     catch
        return det_df(a)
     end
  end
end

################################################################################
#
#  Rank
#
################################################################################

function rank(a::nmod_mat)
  r = ccall((:nmod_mat_rank, :libflint), Int, (Ref{nmod_mat}, ), a)
  return r
end

################################################################################
#
#  Inverse
#
################################################################################

function inv(a::nmod_mat)
  !issquare(a) && error("Matrix must be a square matrix")
  z = similar(a)
  r = ccall((:nmod_mat_inv, :libflint), Int,
          (Ref{nmod_mat}, Ref{nmod_mat}), z, a)
  !Bool(r) && error("Matrix not invertible")
  return z
end

################################################################################
#
#  Linear solving
#
################################################################################

function solve(x::nmod_mat, y::nmod_mat)
  (base_ring(x) != base_ring(y)) && error("Matrices must have same base ring")
  !issquare(x)&& error("First argument not a square matrix in solve")
  (y.r != x.r) || y.c != 1 && ("Not a column vector in solve")
  z = similar(y)
  r = ccall((:nmod_mat_solve, :libflint), Int,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  !Bool(r) && error("Singular matrix in solve")
  return z
end

################################################################################
#
#  LU decomposition
#
################################################################################

function lufact!(P::Generic.perm, x::nmod_mat)
  rank = ccall((:nmod_mat_lu, :libflint), Cint, (Ref{Int}, Ref{nmod_mat}, Cint),
           P.d, x, 0)

  for i in 1:length(P.d)
    P.d[i] += 1
  end

  return rank
end

function lufact(x::nmod_mat, P = PermGroup(rows(x)))
  m = rows(x)
  n = cols(x)
  P.n != m && error("Permutation does not match matrix")
  p = P()
  R = base_ring(x)
  U = deepcopy(x)

  L = similar(x, m, m)

  rank = lufact!(p, U)

  for i = 1:m
    for j = 1:n
      if i > j
        L[i, j] = U[i, j]
        U[i, j] = R()
      elseif i == j
        L[i, j] = R(1)
      elseif j <= m
        L[i, j] = R()
      end
    end
  end
  return rank, p, L, U
end

################################################################################
#
#  Windowing
#
################################################################################

function Base.view(x::nmod_mat, r1::Int, c1::Int, r2::Int, c2::Int)
  Generic._checkbounds(x, r1, c1)
  Generic._checkbounds(x, r2, c2)
  (r1 > r2 || c1 > c2) && error("Invalid parameters")
  temp = similar(x, r2 - r1 + 1, c2 - c1 + 1)
  ccall((:nmod_mat_window_init, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Int, Int, Int, Int),
          temp, x, r1-1, c1-1, r2, c2)
  z = deepcopy(temp)
  ccall((:nmod_mat_window_clear, :libflint), Void, (Ref{nmod_mat}, ), temp)
  return z
end

function Base.view(x::nmod_mat, r::UnitRange{Int}, c::UnitRange{Int})
  return Base.view(x, r.start, c.start, r.stop, c.stop)
end

sub(x::nmod_mat, r1::Int, c1::Int, r2::Int, c2::Int) =
        Base.view(x, r1, c1, r2, c2)

sub(x::nmod_mat, r::UnitRange{Int}, c::UnitRange{Int}) = Base.view(x, r, c)

################################################################################
#
#  Concatenation
#
################################################################################

function hcat(x::nmod_mat, y::nmod_mat)
  (base_ring(x) != base_ring(y)) && error("Matrices must have same base ring")
  (x.r != y.r) && error("Matrices must have same number of rows")
  z = similar(x, rows(x), cols(x) + cols(y))
  ccall((:nmod_mat_concat_horizontal, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  return z
end

function vcat(x::nmod_mat, y::nmod_mat)
  (base_ring(x) != base_ring(y)) && error("Matrices must have same base ring")
  (x.c != y.c) && error("Matrices must have same number of columns")
  z = similar(x, rows(x) + rows(y), cols(x))
  ccall((:nmod_mat_concat_vertical, :libflint), Void,
          (Ref{nmod_mat}, Ref{nmod_mat}, Ref{nmod_mat}), z, x, y)
  return z
end

################################################################################
#
#  Conversion
#
################################################################################

function Array(b::nmod_mat)
  a = Array{nmod}(b.r, b.c)
  for i = 1:b.r
    for j = 1:b.c
      a[i,j] = b[i,j]
    end
  end
  return a
end

################################################################################
#
#  Lifting
#
################################################################################

doc"""
    lift(a::nmod_mat)
> Return a lift of the matrix $a$ to a matrix over $\mathbb{Z}$, i.e. where the
> entries of the returned matrix are those of $a$ lifted to $\mathbb{Z}$.
"""
function lift(a::nmod_mat)
  z = fmpz_mat(rows(a), cols(a))
  z.base_ring = FlintIntegerRing()
  ccall((:fmpz_mat_set_nmod_mat, :libflint), Void,
          (Ref{fmpz_mat}, Ref{nmod_mat}), z, a)
  return z
end

function lift!(z::fmpz_mat, a::nmod_mat)
  ccall((:fmpz_mat_set_nmod_mat, :libflint), Void,
          (Ref{fmpz_mat}, Ref{nmod_mat}), z, a)
  return z
end

################################################################################
#
#  Characteristic polynomial
#
################################################################################

function charpoly(R::NmodPolyRing, a::nmod_mat)
  m = deepcopy(a)
  p = R()
  ccall((:nmod_mat_charpoly, :libflint), Void,
          (Ref{nmod_poly}, Ref{nmod_mat}), p, m)
  return p
end

################################################################################
#
#  Minimal polynomial
#
################################################################################

function minpoly(R::NmodPolyRing, a::nmod_mat)
  p = R()
  ccall((:nmod_mat_minpoly, :libflint), Void,
          (Ref{nmod_poly}, Ref{nmod_mat}), p, a)
  return p
end

###############################################################################
#
#   Promotion rules
#
###############################################################################

promote_rule(::Type{nmod_mat}, ::Type{V}) where {V <: Integer} = nmod_mat

promote_rule(::Type{nmod_mat}, ::Type{nmod}) = nmod_mat

promote_rule(::Type{nmod_mat}, ::Type{fmpz}) = nmod_mat

################################################################################
#
#  Parent object overloading
#
################################################################################

function (a::NmodMatSpace)()
  z = nmod_mat(a.rows, a.cols, a.n)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(b::Integer)
   M = a()
   for i = 1:a.rows
      for j = 1:a.cols
         if i != j
            M[i, j] = zero(base_ring(a))
         else
            M[i, j] = base_ring(a)(b)
         end
      end
   end
   return M
end

function (a::NmodMatSpace)(b::fmpz)
   M = a()
   for i = 1:a.rows
      for j = 1:a.cols
         if i != j
            M[i, j] = zero(base_ring(a))
         else
            M[i, j] = base_ring(a)(b)
         end
      end
   end
   return M
end

function (a::NmodMatSpace)(b::nmod)
   parent(b) != base_ring(a) && error("Unable to coerce to matrix")
   M = a()
   for i = 1:a.rows
      for j = 1:a.cols
         if i != j
            M[i, j] = zero(base_ring(a))
         else
            M[i, j] = deepcopy(b)
         end
      end
   end
   return M
end

function (a::NmodMatSpace)(arr::Array{BigInt, 2}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr, transpose)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{BigInt, 1}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{fmpz, 2}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr, transpose)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{fmpz, 1}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{Int, 2}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr, transpose)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{Int, 1}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr)
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{nmod, 2}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr, transpose)
  (length(arr) > 0 && (base_ring(a) != parent(arr[1]))) && error("Elements must have same base ring")
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(arr::Array{nmod, 1}, transpose::Bool = false)
  _check_dim(a.rows, a.cols, arr)
  (length(arr) > 0 && (base_ring(a) != parent(arr[1]))) && error("Elements must have same base ring")
  z = nmod_mat(a.rows, a.cols, a.n, arr, transpose)
  z.base_ring = a.base_ring
  return z
end

function (a::NmodMatSpace)(b::fmpz_mat)
  (a.cols != b.c || a.rows != b.r) && error("Dimensions do not fit")
  z = nmod_mat(a.n, b)
  z.base_ring = a.base_ring
  return z
end

###############################################################################
#
#   Matrix constructor
#
###############################################################################

function matrix(R::NmodRing, arr::Array{<: Union{nmod, fmpz, Integer}, 2})
   z = nmod_mat(size(arr, 1), size(arr, 2), R.n, arr)
   z.base_ring = R
   return z
end

function matrix(R::NmodRing, r::Int, c::Int, arr::Array{<: Union{nmod, fmpz, Integer}, 1})
   _check_dim(r, c, arr)
   z = nmod_mat(r, c, R.n, arr)
   z.base_ring = R
   return z
end

###############################################################################
#
#  Zero matrix
#
###############################################################################

function zero_matrix(R::NmodRing, r::Int, c::Int)
   z = nmod_mat(r, c, R.n)
   z.base_ring = R
   return z
end

###############################################################################
#
#  Identity matrix
#
###############################################################################

function identity_matrix(R::NmodRing, n::Int)
   z = zero_matrix(R, n, n)
   for i in 1:n
      z[i, i] = one(R)
   end
   z.base_ring = R
   return z
end

################################################################################
#
#  Matrix space constructor
#
################################################################################

function MatrixSpace(R::NmodRing, r::Int, c::Int, cached::Bool = true)
  NmodMatSpace(R, r, c, cached)
end

function MatrixSpace(R::Generic.ResRing{fmpz}, r::Int, c::Int, cached::Bool = true)
  T = elem_type(R)
  return Generic.MatSpace{T}(R, r, c, cached)
end

ext: jl


text:
function __solve(prob::OptimizationProblem, opt::NLopt.Opt;
                 maxiters = nothing, nstart = 1,
                 local_method = nothing,
                 progress = false, kwargs...)
    local x

    if !(isnothing(maxiters)) && maxiters <= 0.0
        error("The number of maxiters has to be a non-negative and non-zero number.")
    elseif !(isnothing(maxiters))
        maxiters = convert(Int, maxiters)
    end

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    _loss = function(θ)
        x = f.f(θ, prob.p)
        return x[1]
    end

    fg! = function (θ,G)
        if length(G) > 0
            f.grad(G, θ)
        end

        return _loss(θ)
    end

    NLopt.min_objective!(opt, fg!)

    if prob.ub !== nothing
        NLopt.upper_bounds!(opt, prob.ub)
    end
    if prob.lb !== nothing
        NLopt.lower_bounds!(opt, prob.lb)
    end
    if !(isnothing(maxiters))
        NLopt.maxeval!(opt, maxiters)
    end
    if nstart > 1 && local_method !== nothing
        NLopt.local_optimizer!(opt, local_method)
        if !(isnothing(maxiters))
            NLopt.maxeval!(opt, nstart * maxiters)
        end
    end

    t0 = time()
    (minf,minx,ret) = NLopt.optimize(opt, prob.u0)
    _time = time()

    SciMLBase.build_solution(prob, opt, minx, minf; original=nothing)
end


text:
# AUTO GENERATED FILE - DO NOT EDIT

export datepicker

"""
    datepicker(;kwargs...)

A DatePicker component.
Capture date input from user
Keyword arguments:
- `id` (String; optional): The ID of this component, used to identify dash components in callbacks
- `allowFreeInput` (Bool; optional): Allow free input
- `allowLevelChange` (Bool; optional): Allow to change level (date – month – year)
- `amountOfMonths` (Real; optional): Amount of displayed months
- `className` (String; optional): Often used with CSS to style elements with common properties
- `clearable` (Bool; optional): Allow to clear value
- `closeCalendarOnChange` (Bool; optional): Set to false to force dropdown to stay open after date was selected
- `closeDropdownOnScroll` (Bool; optional): Set to true to disable dropdown closing on scroll
- `date` (String; optional): Selected date
- `description` (String; optional): Input description, displayed after label
- `disableOutsideEvents` (Bool; optional): When true dates that are outside of given month cannot be clicked or focused
- `disabled` (Bool; optional): A Datepicker can show it is currently unable to be interacted with
- `dropdownType` (a value equal to: "modal", "popover"; optional): Where to show calendar in modal or popover
- `error` (String; optional): Displays error message after input
- `firstDayOfWeek` (a value equal to: "sunday", "monday"; optional): Set first day of the week
- `fixOnBlur` (Bool; optional): call onChange with last valid value onBlur
- `focusable` (Bool; optional): Should focusable days have tabIndex={0}?
- `format` (String; optional): DatePicker display format
- `fullWidth` (Bool; optional): Set to true to make calendar take 100% of container width
- `hideWeekdays` (Bool; optional): Set to false to remove weekdays row
- `initialLevel` (a value equal to: "date", "month", "year"; optional): Initial date selection level
- `initialMonth` (String; optional): Initial selected month
- `initiallyOpened` (Bool; optional): Control initial dropdown opened state
- `label` (String; optional): Input label, displayed before input
- `locale` (String; optional): Locale used for all labels formatting
- `maxDate` (String; optional): Maximum possible date
- `minDate` (String; optional): Minimum possible date
- `multiline` (Bool; optional): Will input have multiple lines?
- `placeholder` (String; optional): Placeholder, displayed when date is not selected
- `preventFocus` (Bool; optional): Prevent focusing upon clicking
- `radius` (a value equal to: "xs", "sm", "md", "lg", "xl" | Real; optional): Input border-radius from theme or number to set border-radius in px
- `required` (Bool; optional): Adds red asterisk on the right side of label
- `size` (a value equal to: "xs", "sm", "md", "lg", "xl"; optional): Input size
- `style` (Dict; optional): Inline style override
- `variant` (a value equal to: "default", "filled", "unstyled", "headless"; optional): Defines input appearance, defaults to default in light color scheme and filled in dark
- `withinPortal` (Bool; optional): Whether to render the dropdown in a Portal
- `zIndex` (Real; optional): Popper zIndex
"""
function datepicker(; kwargs...)
        available_props = Symbol[:id, :allowFreeInput, :allowLevelChange, :amountOfMonths, :className, :clearable, :closeCalendarOnChange, :closeDropdownOnScroll, :date, :description, :disableOutsideEvents, :disabled, :dropdownType, :error, :firstDayOfWeek, :fixOnBlur, :focusable, :format, :fullWidth, :hideWeekdays, :initialLevel, :initialMonth, :initiallyOpened, :label, :locale, :maxDate, :minDate, :multiline, :placeholder, :preventFocus, :radius, :required, :size, :style, :variant, :withinPortal, :zIndex]
        wild_props = Symbol[]
        return Component("datepicker", "DatePicker", "dash_mantine_components", available_props, wild_props; kwargs...)
end


text:
function ASIGetNumOfConnectedCameras()
    ccall((:ASIGetNumOfConnectedCameras, libASICamera2), Cint, ())
end

function ASIGetProductIDs(pPIDs)
    ccall((:ASIGetProductIDs, libASICamera2), Cint, (Ref{Cint},), pPIDs)
end

function ASIGetCameraProperty(pASICameraInfo, iCameraIndex)
    ccall((:ASIGetCameraProperty, libASICamera2), ASI_ERROR_CODE, (Ref{_ASI_CAMERA_INFO}, Cint), pASICameraInfo, iCameraIndex)
end

function ASIGetCameraPropertyByID(iCameraID, pASICameraInfo)
    ccall((:ASIGetCameraPropertyByID, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{ASI_CAMERA_INFO}), iCameraID, pASICameraInfo)
end

function ASIOpenCamera(iCameraID)
    ccall((:ASIOpenCamera, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASIInitCamera(iCameraID)
    ccall((:ASIInitCamera, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASICloseCamera(iCameraID)
    ccall((:ASICloseCamera, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

# ASI_ERROR_CODE ASIGetNumOfControls(int iCameraID, int * piNumberOfControls);
function ASIGetNumOfControls(iCameraID, piNumberOfControls::Ref{Cint})
    ccall((:ASIGetNumOfControls, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cint}), iCameraID, piNumberOfControls)
end

function ASIGetControlCaps(iCameraID, iControlIndex, pControlCaps)
    ccall((:ASIGetControlCaps, libASICamera2), ASI_ERROR_CODE, (Cint, Cint, Ref{_ASI_CONTROL_CAPS}), iCameraID, iControlIndex, pControlCaps)
end

function ASIGetControlValue(iCameraID, ControlType, plValue::Ref{Clong}, pbAuto::Ref{ASI_BOOL})
    ccall((:ASIGetControlValue, libASICamera2), ASI_ERROR_CODE, (Cint, Cint, Ref{Clong}, Ref{ASI_BOOL}), iCameraID, ControlType, plValue, pbAuto)
end

function ASISetControlValue(iCameraID, ControlType, lValue, bAuto)
    ccall((:ASISetControlValue, libASICamera2), ASI_ERROR_CODE, (Cint, Cint, Clong, Cint), iCameraID, ControlType, lValue, bAuto)
end

function ASISetROIFormat(iCameraID, iWidth, iHeight, iBin, Img_type::ASI_IMG_TYPE)
    ccall((:ASISetROIFormat, libASICamera2), ASI_ERROR_CODE, (Cint, Cint, Cint, Cint, ASI_IMG_TYPE), iCameraID, iWidth, iHeight, iBin, Img_type)
end

function ASIGetROIFormat(iCameraID, piWidth::Ref{Cint}, piHeight::Ref{Cint}, piBin::Ref{Cint}, pImg_type::Ref{ASI_IMG_TYPE})
    ccall((:ASIGetROIFormat, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cint}, Ref{Cint}, Ref{Cint}, Ref{ASI_IMG_TYPE}), iCameraID, piWidth, piHeight, piBin, pImg_type)
end

function ASISetStartPos(iCameraID, iStartX, iStartY)
    ccall((:ASISetStartPos, libASICamera2), ASI_ERROR_CODE, (Cint, Cint, Cint), iCameraID, iStartX, iStartY)
end

function ASIGetStartPos(iCameraID, piStartX::Ref{Cint}, piStartY::Ref{Cint})
    ccall((:ASIGetStartPos, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cint}, Ref{Cint}), iCameraID, piStartX, piStartY)
end

function ASIGetDroppedFrames(iCameraID, piDropFrames)
    ccall((:ASIGetDroppedFrames, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cint}), iCameraID, piDropFrames)
end

function ASIEnableDarkSubtract(iCameraID, pcBMPPath)
    ccall((:ASIEnableDarkSubtract, libASICamera2), ASI_ERROR_CODE, (Cint, Cstring), iCameraID, pcBMPPath)
end

function ASIDisableDarkSubtract(iCameraID)
    ccall((:ASIDisableDarkSubtract, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASIStartVideoCapture(iCameraID)
    ccall((:ASIStartVideoCapture, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASIStopVideoCapture(iCameraID)
    ccall((:ASIStopVideoCapture, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASIGetVideoData(iCameraID, pBuffer::T, lBuffSize, iWaitms) where T
    if T == Matrix{UInt8}
        ccall((:ASIGetVideoData, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cuchar}, Clong, Cint), iCameraID, pBuffer, lBuffSize, iWaitms)
    elseif T == Matrix{UInt16}  # explicitly handle unsafe conversion from uint16 to uint8
        ccall((:ASIGetVideoData, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cuchar}, Clong, Cint), iCameraID, Base.unsafe_convert(Ref{Cuchar}, pBuffer), lBuffSize, iWaitms)
    end
end

function ASIPulseGuideOn(iCameraID, direction)
    ccall((:ASIPulseGuideOn, libASICamera2), ASI_ERROR_CODE, (Cint, Cint), iCameraID, direction)
end

function ASIPulseGuideOff(iCameraID, direction)
    ccall((:ASIPulseGuideOff, libASICamera2), ASI_ERROR_CODE, (Cint, Cint), iCameraID, direction)
end

function ASIStartExposure(iCameraID, bIsDark)
    ccall((:ASIStartExposure, libASICamera2), ASI_ERROR_CODE, (Cint, Cint), iCameraID, bIsDark)
end

function ASIStopExposure(iCameraID)
    ccall((:ASIStopExposure, libASICamera2), ASI_ERROR_CODE, (Cint,), iCameraID)
end

function ASIGetExpStatus(iCameraID, pExpStatus)
    ccall((:ASIGetExpStatus, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{ASI_EXPOSURE_STATUS}), iCameraID, pExpStatus)
end

function ASIGetDataAfterExp(iCameraID, pBuffer::T, lBuffSize) where T
    if T == Matrix{UInt8}
        ccall((:ASIGetDataAfterExp, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cuchar}, Clong), iCameraID, pBuffer, lBuffSize)
    elseif T == Matrix{UInt16}  # explicitly handle unsafe conversion from uint16 to uint8
        ccall((:ASIGetDataAfterExp, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cuchar}, Clong), iCameraID, Base.unsafe_convert(Ref{Cuchar}, pBuffer), lBuffSize)
    end
end

function ASIGetID(iCameraID, pID)
    ccall((:ASIGetID, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{ASI_ID}), iCameraID, pID)
end

function ASISetID(iCameraID, ID)
    ccall((:ASISetID, libASICamera2), ASI_ERROR_CODE, (Cint, ASI_ID), iCameraID, ID)
end

function ASIGetGainOffset(iCameraID, pOffset_HighestDR, pOffset_UnityGain, pGain_LowestRN, pOffset_LowestRN)
    ccall((:ASIGetGainOffset, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{Cint}, Ref{Cint}, Ref{Cint}, Ref{Cint}), iCameraID, pOffset_HighestDR, pOffset_UnityGain, pGain_LowestRN, pOffset_LowestRN)
end

function ASIGetSDKVersion()
    ccall((:ASIGetSDKVersion, libASICamera2), Cstring, ())
end

function ASIGetCameraSupportMode(iCameraID, pSupportedMode)
    ccall((:ASIGetCameraSupportMode, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{_ASI_SUPPORTED_MODE}), iCameraID, pSupportedMode)
end

function ASIGetCameraMode(iCameraID, mode)
    ccall((:ASIGetCameraMode, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{ASI_CAMERA_MODE}), iCameraID, mode)
end

function ASISetCameraMode(iCameraID, mode)
    ccall((:ASISetCameraMode, libASICamera2), ASI_ERROR_CODE, (Cint, ASI_CAMERA_MODE), iCameraID, mode)
end

function ASISendSoftTrigger(iCameraID, bStart)
    ccall((:ASISendSoftTrigger, libASICamera2), ASI_ERROR_CODE, (Cint, Cint), iCameraID, bStart)
end

function ASIGetSerialNumber(iCameraID, pSN)
    ccall((:ASIGetSerialNumber, libASICamera2), ASI_ERROR_CODE, (Cint, Ref{ASI_SN}), iCameraID, pSN)
end

function ASISetTriggerOutputIOConf(iCameraID, pin, bPinHigh, lDelay, lDuration)
    ccall((:ASISetTriggerOutputIOConf, libASICamera2), ASI_ERROR_CODE, (Cint, ASI_TRIG_OUTPUT_PIN, Cint, Clong, Clong), iCameraID, pin, bPinHigh, lDelay, lDuration)
end

function ASIGetTriggerOutputIOConf(iCameraID, pin, bPinHigh, lDelay, lDuration)
    ccall((:ASIGetTriggerOutputIOConf, libASICamera2), ASI_ERROR_CODE, (Cint, ASI_TRIG_OUTPUT_PIN, Ref{Cint}, Ref{Clong}, Ref{Clong}), iCameraID, pin, bPinHigh, lDelay, lDuration)
end


text:
module NVTX

using CUDAapi

using CUDAdrv: CUstream, CUdevice, CUcontext, CUevent

using CEnum

using ExprTools

using ..CUDAnative
using ..CUDAnative: libnvtx

# core library
initialize_api() = return
include("libnvtx_common.jl")
include("libnvtx.jl")

include("highlevel.jl")

end