feat(assemblers): add kernel column and KA scalar homogeneity

Add PerElementKernelColumn accessors, ka_column_homogeneity guard, and README documenting cache.kernel_column and Pass~1 separation from assemblers.
2026-09-19 17:58:53 +00:00 · 2026-05-11 02:23:04 +03:00
parent a54c36ff5e
commit a0d67cdb59
3 changed files with 248 additions and 1 deletions
@@ -1,3 +1,8 @@
+<!--
+SPDX-FileCopyrightText: 2015-2026 Jukka Aho
+SPDX-License-Identifier: MIT
+-->
+
 # src/assemblers/dof_based/

 DOF-by-DOF assembler. The driver loops over global DOFs (rather than
@@ -8,9 +13,34 @@ every KernelAbstractions backend without code duplication.

 ## Files

+- `kernel_column.jl` — `UniformKernelColumn` (one kernel for all elements) and `PerElementKernelColumn` (one kernel per element id). `assert_homogeneous_dof_based_kernel_column!` validates compatible kernels before cache construction. `ka_per_element_kernel_column_supported` gates whether `DOFBasedCOOCacheKA` may be built for a per-element column (KA `apply_K!` passes a single prototype kernel; see below).
 - `dof_based_coo.jl` — CPU implementation. Defines `DOFBasedCOOAssembler`, `DOFBasedCOOCache`, and the matrix-free entry points (`apply_K!`, `apply_K_contributions!`, `apply_M!`, `assemble_M!`, `extract_system`). Built around the microkernel contract in `assemblers/microkernel.jl` and the DOF connectivity in `dofs/dof_connectivity.jl`.
 - `dof_based_coo_ka.jl` — Backend-agnostic GPU port via `KernelAbstractions.jl`. Defines `DOFBasedCOOCacheKA`, `sync_from_cpu!`, and the precision helpers used by the Float32 GPU pipeline. The same kernel runs on `CPU()`, `CUDABackend()`, `MetalBackend()`, `AMDGPUBackend()`, `oneAPIBackend()`; the in-tree CI validates the CPU() backend.

+## Kernel columns and KA `apply_K!`
+
+The KA matvec (`apply_K!` on `DOFBasedCOOCacheKA`) passes one prototype kernel
+object into `evaluate_entry` for every element row, together with per-element
+views of `qp_buffers` filled during Pass 1 on the CPU cache. That matches the
+CPU path only if every element’s stiffness contribution either reads material
+data from `qp_buffers` alone, or reads extra scalars that are already forced
+identical across the column (density, Biot `α` / `storage_S`, thermo `β`,
+thermo-poro coupling scalars, and so on). For a `PerElementKernelColumn{K}`,
+`ka_per_element_kernel_column_supported` is `true` when `K` is one of
+`ContinuumKernel`, `HeatKernel`, `ThermoElasticKernel`, `BiotPoroelasticKernel`,
+or `ThermoPoroelasticKernel`. Other kernel types still work on the CPU column
+path with `kernel_at(cache, eid)`; extending the KA gate requires checking that
+`evaluate_entry` does not depend on per-element fields in the kernel object
+beyond what Pass 1 copies into `qp_buffers`.
+
+`apply_M!(y, cache_ka, kernel, x)` mirrors the same layout, calling
+`evaluate_mass_entry` on the chosen backend (same per-element column gate as
+`apply_K!`).
+
+For Krylov-style use, `MatrixFreeMassOperatorKA` / `matrix_free_mass_op_ka` in
+`assemblers/matrix_free/operator.jl` wraps the same matvec with device-matched
+scratch buffers (`test/assemblers/test_matrix_free_mass_operator_ka.jl`).
+
 ## Design notes

 The hot path is intentionally a tight loop over global DOFs; a Pass 1
@@ -47,4 +77,4 @@ copy `x`), which is the extension point for ghost DOFs under MPI.
 MPI helpers (weak dependency `MPI`, extension `ext/JuliaFEMMPIExt.jl`): halo exchange
 `exchange_matvec_halos_mpi!`, owned-row product `mpi_partitioned_operator_matvec_owned!`,
 workspace `partitioned_mpi_owned_matvec_workspace`, and reusable request buffers via
-`allocate_exchange_matvec_halo_mpi_requests` / keyword `mpi_requests` (see `AGENTS.md` §3.6).
+`allocate_exchange_matvec_halo_mpi_requests` / keyword `mpi_requests` (see the MPI subsection in `AGENTS.md`).
@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
+# SPDX-License-Identifier: MIT
+
+"""
+Scalar-parameter homogeneity for per-element kernel columns on the KA matvec path.
+
+[`PerElementKernelColumn`](@ref) stores one kernel object per element. The
+KernelAbstractions launcher passes a **single prototype** kernel (the first
+element's) into device code while per-element data live in `qp_buffers`.
+Any extra scalar read from the kernel object (density, heat capacity, Biot
+`α`, …) must therefore match across the whole column. Setup-only checks live
+here; the hot path uses `kernel_at(col, eid)` on CPU only.
+"""
+
+@inline function _assert_homogeneous_ka_column_kernel_scalars!(kernels::Vector{K}) where {K}
+    length(kernels) < 2 && return nothing
+    @inbounds k1 = kernels[1]
+    return _assert_homogeneous_ka_column_kernel_scalars!(k1, kernels)
+end
+
+@inline _assert_homogeneous_ka_column_kernel_scalars!(::AbstractKernel, ::Vector) = nothing
+
+function _assert_homogeneous_ka_column_kernel_scalars!(k1::ContinuumKernel, kernels::Vector)
+    d0 = k1.density
+    for i in 2:length(kernels)
+        @inbounds ki = kernels[i]
+        ki.density == d0 || throw(ArgumentError(
+            "per-element ContinuumKernel: density must match on all elements " *
+            "(mass microkernel reads kernel.density); mismatch at element $i",
+        ))
+    end
+    return nothing
+end
+
+function _assert_homogeneous_ka_column_kernel_scalars!(k1::HeatKernel, kernels::Vector)
+    c0 = k1.heat_capacity
+    for i in 2:length(kernels)
+        @inbounds ki = kernels[i]
+        ki.heat_capacity == c0 || throw(ArgumentError(
+            "per-element HeatKernel: heat_capacity must match on all elements; mismatch at element $i",
+        ))
+    end
+    return nothing
+end
+
+function _assert_homogeneous_ka_column_kernel_scalars!(k1::ThermoElasticKernel, kernels::Vector)
+    β0 = k1.β
+    for i in 2:length(kernels)
+        @inbounds ki = kernels[i]
+        ki.β == β0 || throw(ArgumentError(
+            "per-element ThermoElasticKernel: β must match on all elements; mismatch at element $i",
+        ))
+    end
+    return nothing
+end
+
+function _assert_homogeneous_ka_column_kernel_scalars!(k1::BiotPoroelasticKernel, kernels::Vector)
+    α0 = k1.α
+    S0 = k1.storage_S
+    ρ0 = k1.density
+    for i in 2:length(kernels)
+        @inbounds ki = kernels[i]
+        ki.α == α0 ||
+            throw(ArgumentError("per-element BiotPoroelasticKernel: α mismatch at element $i"))
+        ki.storage_S == S0 ||
+            throw(ArgumentError("per-element BiotPoroelasticKernel: storage_S mismatch at element $i"))
+        ki.density == ρ0 ||
+            throw(ArgumentError("per-element BiotPoroelasticKernel: density mismatch at element $i"))
+    end
+    return nothing
+end
+
+function _assert_homogeneous_ka_column_kernel_scalars!(k1::ThermoPoroelasticKernel, kernels::Vector)
+    β0 = k1.β
+    α0 = k1.α
+    S0 = k1.storage_S
+    κ0 = k1.kappa_tp
+    ζ0 = k1.zeta_tp
+    ρcp0 = k1.heat_capacity
+    ρs0 = k1.density
+    for i in 2:length(kernels)
+        @inbounds ki = kernels[i]
+        ki.β == β0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: β mismatch at element $i"))
+        ki.α == α0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: α mismatch at element $i"))
+        ki.storage_S == S0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: storage_S mismatch at element $i"))
+        ki.kappa_tp == κ0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: kappa_tp mismatch at element $i"))
+        ki.zeta_tp == ζ0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: zeta_tp mismatch at element $i"))
+        ki.heat_capacity == ρcp0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: heat_capacity mismatch at element $i"))
+        ki.density == ρs0 ||
+            throw(ArgumentError("per-element ThermoPoroelasticKernel: density mismatch at element $i"))
+    end
+    return nothing
+end
@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
+# SPDX-License-Identifier: MIT
+
+#=
+Column of volume kernels for [`DOFBasedCOOCache`](@ref).
+
+Two storage shapes, both type-stable and allocation-free on assembly hot paths:
+
+* [`UniformKernelColumn`](@ref) — one kernel reused for every element (no extra
+  storage beyond the struct).
+* [`PerElementKernelColumn`](@ref) — one concrete kernel per element, stored in a
+  pre-allocated `Vector{K}` built at cache construction time (setup may allocate;
+  Pass 1 / Pass 2 only index this vector).
+
+All kernels in a per-element column must share the same `qpoint_buffer_eltype`,
+the same `reference_fields` / state NamedTuple types, and the same
+`dofs_per_node`. Kernels whose microkernels read additional scalar parameters
+from the kernel object (Biot `α`/`storage_S`/`density`, thermo-elastic `β`,
+thermo-poroelastic `β`/`α`/`storage_S`/`kappa_tp`/`zeta_tp`/`heat_capacity`,
+continuum `density` for mass, heat `heat_capacity`, …) must match those parameters across elements so
+that GPU / KA paths that pass a single prototype kernel remain consistent; the
+CPU path always uses `kernel_at(col, eid)`. Scalar checks are implemented in
+`ka_column_homogeneity.jl` as `_assert_homogeneous_ka_column_kernel_scalars!`. For kernels whose
+`evaluate_entry` reads only `qp_buffers` plus those matched scalars
+(`ContinuumKernel`, `HeatKernel`, `ThermoElasticKernel`, `BiotPoroelasticKernel`,
+`ThermoPoroelasticKernel`), `ka_per_element_kernel_column_supported` allows the
+KA matvec path.
+=#
+
+"""
+    UniformKernelColumn{K<:AbstractKernel}
+
+Store a single instance `kernel::K` used for every volume element. Hot-path
+lookup is a direct field read (no indexing).
+"""
+struct UniformKernelColumn{K<:AbstractKernel}
+    kernel::K
+end
+
+"""
+    PerElementKernelColumn{K<:AbstractKernel}
+
+Store `kernels[eid]` for each volume element id `eid`. The vector is owned for
+the lifetime of the cache and filled at construction (no per-assembly
+allocation).
+"""
+struct PerElementKernelColumn{K<:AbstractKernel}
+    kernels::Vector{K}
+end
+
+@inline kernel_at(col::UniformKernelColumn, ::Int) = col.kernel
+@inline kernel_at(col::PerElementKernelColumn, eid::Int) = @inbounds col.kernels[eid]
+
+@inline prototype_kernel(col::UniformKernelColumn) = col.kernel
+@inline prototype_kernel(col::PerElementKernelColumn) = @inbounds col.kernels[1]
+
+"""
+    assert_homogeneous_dof_based_kernel_column!(kernels::Vector{K}) where {K}
+
+Validate a vector of kernels before wrapping it in [`PerElementKernelColumn`](@ref).
+Called only from cache construction (setup tier), not from assembly hot paths.
+"""
+function assert_homogeneous_dof_based_kernel_column!(kernels::Vector{K}) where {K}
+    n = length(kernels)
+    n ≥ 1 || throw(ArgumentError("per-element kernels: empty vector"))
+    @inbounds k1 = kernels[1]
+    if k1 isa HeatKernel && k1.material isa ElementWiseScalarDiffusion
+        throw(ArgumentError(
+            "per-element kernel column is not supported with ElementWiseScalarDiffusion; " *
+            "use a single HeatKernel whose material carries λ_by_elem[elem_id] instead.",
+        ))
+    end
+    fr1, st1 = reference_fields(k1)
+    Buf1 = qpoint_buffer_eltype(k1)
+    dpn1 = dofs_per_node(k1)
+    for i in 2:n
+        @inbounds ki = kernels[i]
+        dofs_per_node(ki) == dpn1 || throw(ArgumentError(
+            "per-element kernels: dofs_per_node mismatch at element $i (" *
+            "$(dofs_per_node(ki)) vs $dpn1)",
+        ))
+        qpoint_buffer_eltype(ki) === Buf1 || throw(ArgumentError(
+            "per-element kernels: qpoint_buffer_eltype mismatch at element $i",
+        ))
+        fri, sti = reference_fields(ki)
+        typeof(fri) === typeof(fr1) || throw(ArgumentError(
+            "per-element kernels: reference_fields tuple type mismatch at element $i",
+        ))
+        typeof(sti) === typeof(st1) || throw(ArgumentError(
+            "per-element kernels: material state type mismatch at element $i",
+        ))
+    end
+    _assert_homogeneous_ka_column_kernel_scalars!(kernels)
+    return nothing
+end
+
+"""
+    ka_per_element_kernel_column_supported(col::PerElementKernelColumn) -> Bool
+
+`true` when the KA GPU matvec may use a single prototype kernel object (the
+first element's kernel) together with per-element `qp_buffers` filled on CPU
+Pass 1. Allowed when every element's `evaluate_entry` either reads material
+data only from `qp_buffers`, or reads additional scalars that
+[`assert_homogeneous_dof_based_kernel_column!`](@ref) already enforces across
+the column (`ContinuumKernel`, `HeatKernel`, `ThermoElasticKernel`,
+`BiotPoroelasticKernel`, `ThermoPoroelasticKernel`).
+"""
+@inline function ka_per_element_kernel_column_supported(::PerElementKernelColumn{K}) where {K}
+    return (
+        K <: ContinuumKernel ||
+        K <: HeatKernel ||
+        K <: ThermoElasticKernel ||
+        K <: BiotPoroelasticKernel ||
+        K <: ThermoPoroelasticKernel
+    )
+end
+
+@inline ka_per_element_kernel_column_supported(::UniformKernelColumn) = true