From a0d67cdb595b79b9ba2f6047c6617d160aed0d03 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Mon, 11 May 2026 02:23:04 +0300 Subject: [PATCH] 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. --- src/assemblers/dof_based/README.md | 32 ++++- .../dof_based/ka_column_homogeneity.jl | 99 +++++++++++++++ src/assemblers/dof_based/kernel_column.jl | 118 ++++++++++++++++++ 3 files changed, 248 insertions(+), 1 deletion(-) create mode 100644 src/assemblers/dof_based/ka_column_homogeneity.jl create mode 100644 src/assemblers/dof_based/kernel_column.jl diff --git a/src/assemblers/dof_based/README.md b/src/assemblers/dof_based/README.md index 8d4a435..1d55ba6 100644 --- a/src/assemblers/dof_based/README.md +++ b/src/assemblers/dof_based/README.md @@ -1,3 +1,8 @@ + + # 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`). diff --git a/src/assemblers/dof_based/ka_column_homogeneity.jl b/src/assemblers/dof_based/ka_column_homogeneity.jl new file mode 100644 index 0000000..b1e2772 --- /dev/null +++ b/src/assemblers/dof_based/ka_column_homogeneity.jl @@ -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 diff --git a/src/assemblers/dof_based/kernel_column.jl b/src/assemblers/dof_based/kernel_column.jl new file mode 100644 index 0000000..2e7046c --- /dev/null +++ b/src/assemblers/dof_based/kernel_column.jl @@ -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