Glassvision/JuliaFEM.jl
1
0
Fork 0
mirror of https://github.com/JuliaFEM/JuliaFEM.jl.git synced 2026-08-05 20:11:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(assemblers): abstract cache, element cache, microkernel, partitioned matvec

Browse Source 
Tighten assembler/kernel contracts, material workspace hooks, and redundant-kernel depwarns on partitioned halo helpers.
This commit is contained in:
Jukka Aho
2026-05-11 02:23:05 +03:00
parent c430babf85
commit 6968cbf9b3
4 changed files with 275 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/assemblers/abstract.jl
+6 -6
View File
@@ -1,5 +1,5 @@
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
# SPDX-License-Identifier: MIT
"""
Abstract assembler type hierarchy.
@@ -118,10 +118,10 @@ multi-thread row-parallel CPU is planned.
# Usage
```julia
elements = create_elements!(mesh, kernel, dof_handler) # Vector{Element}
elements, handler = create_elements!(mesh, ET)
assembler = DOFBasedCOOAssembler()
cache = create_cache(assembler, elements, dof_handler, mesh, kernel)
assemble!(cache, assembler, kernel, mesh) # DOF-wise traversal!
cache = create_cache(assembler, elements, handler, mesh, kernel)
assemble!(cache, assembler, mesh) # DOF-wise traversal (kernel lives in cache)
K, f = extract_system(cache)
```
"""
@@ -252,7 +252,7 @@ Trait answering whether the matrix-free stiffness operator built around
`kernel` is symmetric positive-definite. The default is `true` (single-field
elliptic problems such as continuum displacement, heat). Mixed / saddle-point
kernels (`MixedUPKernel`, `StokesMixedKernel`, `HellingerReissnerKernel`,
`HuWashizuKernel`) override this to `false` so Krylov stacks (CG, etc.) do
`HuWashizuKernel`, `BiotPoroelasticKernel`, `ThermoPoroelasticKernel`) override this to `false` so Krylov stacks (CG, etc.) do
not pick the SPD branch.
"""
@inline operator_is_posdef(::AbstractKernel) = true
src/assemblers/caches/element_cache.jl
+2 -2
View File
@@ -1,5 +1,5 @@
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
# SPDX-License-Identifier: MIT
"""
Element and node cache implementations for zero-allocation assembly.
src/assemblers/microkernel.jl
+71 -8
View File
@@ -1,12 +1,13 @@
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
# SPDX-License-Identifier: MIT
#=
Microkernel contract for the DOF-based assembler.
The DOF-based assembler walks one DOF row at a time and asks the kernel
for a single scalar `K[i, j]`. To keep the assembler kernel-agnostic and
zero-allocation, every kernel must opt in by implementing three pieces:
zero-allocation, every kernel must opt in by implementing three pieces
plus one Pass~1 material hook:
1. `qpoint_buffer_eltype(kernel)` — what type of value the kernel needs
stored once per quadrature point per element. For continuum mechanics
@@ -20,17 +21,32 @@ zero-allocation, every kernel must opt in by implementing three pieces:
use an extra `eid` argument only on the internal
`_dof_based_fill_qpoint_buffer!` dispatch path in `dof_based_coo.jl`.
3. `evaluate_entry(kernel, geometry_cache, qpoint_buffer, layout_i, layout_j, elem_id)`
3. `prepare_dof_based_material_workspace!(kernel, material_workspace,
geometry_cache, element_cache, eid, configuration, global_material_cache,
Δt, ::Type{E})` — fill each integration point of `material_workspace`
for this element before `update_qpoint_buffer!`. The default seeds every
IP from [`reference_fields`](@ref)`(kernel)`; continuum mechanics overrides
this for strain- and state-dependent materials (see
`domains/continuum/dof_based_pass1.jl`).
4. `evaluate_entry(kernel, geometry_cache, qpoint_buffer, layout_i, layout_j, elem_id)`
— the actual microkernel. Returns the single scalar `K[i, j]` for the
local DOF pair `(i, j)` described by two `DOFLayoutEntry` values.
`elem_id` is the volume element index (needed for facet-oriented kernels).
Called inside Pass 2 of `assemble!`, in a hot loop, so it must also
be allocation-free.
Together these three methods let the DOF-based assembler dispatch on any
`AbstractKernel` without baking in continuum-specific assumptions, while
keeping the inner loop fully type-stable thanks to the compile-time
`local_dof_layout(E)` table that produces the `DOFLayoutEntry` arguments.
Together these methods let the DOF-based assembler dispatch on any
`AbstractKernel` without baking in domain-specific Pass~1 logic in the
driver, while keeping the inner loop fully type-stable thanks to the
compile-time `local_dof_layout(E)` table that produces the `DOFLayoutEntry`
arguments.
The KernelAbstractions stiffness / mass matvecs (`dof_based_coo_ka.jl`) call
the same `evaluate_entry` / `evaluate_mass_entry` with one prototype kernel per
launch; per-element variation must live in `qpoint_buffer` columns unless the
volume kernel column is uniform. For `PerElementKernelColumn`, see
`ka_per_element_kernel_column_supported` in `dof_based/kernel_column.jl`.
=#
"""
@@ -142,3 +158,50 @@ A kernel must define this method; there is no default — the previous
assembler accidentally continuum-only.
"""
function reference_fields end
"""
prepare_dof_based_material_workspace!(
kernel,
material_workspace,
geometry_cache,
element_cache,
eid::Int,
configuration::Union{Nothing,AbstractVector{Float64}},
global_material_cache::Union{Nothing,GlobalMaterialCache},
Δt::Float64,
::Type{E},
) -> Nothing
DOF-based Pass~1 (element loop): populate every integration point of
`material_workspace` for this element before [`update_qpoint_buffer!`](@ref).
The default implementation copies [`reference_fields`](@ref)`(kernel)` to
all IPs. Kernels that need configuration-dependent or stateful constitutive
updates (e.g. [`ContinuumKernel`](@ref)) should override this method.
Must be allocation-free on the hot path.
"""
function prepare_dof_based_material_workspace! end
@inline function prepare_dof_based_material_workspace!(
kernel::AbstractKernel,
material_workspace::AssemblyMaterialWorkspace,
::GeometryCache,
element_cache::ElementCache,
::Int,
::Union{Nothing,AbstractVector{Float64}},
::Union{Nothing,GlobalMaterialCache},
::Float64,
::Type{<:AbstractElement},
)
fields_ref_e, empty_state_e = reference_fields(kernel)
fields_mw = getfield(material_workspace, 1)
states_mw = getfield(material_workspace, 2)
ips_ec = getfield(element_cache, :ips)
nips = length(ips_ec)
@inbounds for q in 1:nips
fields_mw[q] = fields_ref_e
states_mw[q] = empty_state_e
end
return nothing
end
src/assemblers/partitioned_matvec.jl
+196 -13
View File
@@ -1,8 +1,8 @@
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
# SPDX-License-Identifier: MIT
# Single-partition orchestration for matrix-free `K*x` on owned rows using a
# packed patch + matvec halo exchange (MPI-ready).
# Single-partition orchestration for matrix-free `K*x` and owned-row internal
# force using a packed patch + matvec halo exchange (MPI-ready).
"""
allocate_halo_recv_buffers(exchange::RankHaloExchange) -> Vector{Vector{Float64}}
@@ -107,10 +107,17 @@ function simulate_halo_recv_from_global!(
end
"""
partitioned_owned_matvec!(
y_contrib, x_global, packed, work, recv_vals,
layout, exchange, cache, assembler, mesh;
fill_recv_from_global = true,
configuration = nothing, global_material_cache = nothing, Δt = 0.0,
) -> y_contrib
partitioned_owned_matvec!(
y_contrib, x_global, packed, work, recv_vals,
layout, exchange, cache, assembler, kernel, mesh;
fill_recv_from_global = true,
configuration = nothing, global_material_cache = nothing, Δt = 0.0,
) -> y_contrib
One partition’s additive contribution to `K*x` on **owned rows** (zeros elsewhere
@@ -121,7 +128,10 @@ metadata:
2. If `fill_recv_from_global` — [`simulate_halo_recv_from_global!`](@ref)`(recv_vals, x_global, exchange)` (skip when MPI fills `recv_vals`)
3. [`unpack_halo_recv_to_packed!`](@ref)`(packed, recv_vals, exchange, layout)`
4. [`fill!`](@ref)`(work, 0)` then [`expand_packed_to_global!`](@ref)`(work, packed, layout)`
5. [`apply_K_owned_rows!`](@ref)`(y_contrib, layout.owned_rows, ...)`
5. [`apply_K_owned_rows!`](@ref)`(y_contrib, layout.owned_rows, …; configuration, …)`
Optional keywords `configuration`, `global_material_cache`, and `Δt` match
[`apply_K!`](@ref) / [`apply_K_owned_rows!`](@ref) (nonlinear continuum Pass~1).
Requires `exchange.part == layout.part`, `length(y_contrib) == cache.ndofs`, and
matching `work` / `packed` sizes from [`partitioned_matvec_workspace`](@ref).
@@ -130,6 +140,8 @@ Summing `y_contrib` over disjoint owned partitions recovers global [`apply_K!`](
when `x_global` is the full vector (serial replica).
Allocation-free after workspace setup when `fill_recv_from_global` follows the same pattern.
The form with `kernel` ignores it (backward compatibility).
"""
function partitioned_owned_matvec!(
y_contrib::AbstractVector{Float64},
@@ -141,9 +153,11 @@ function partitioned_owned_matvec!(
exchange::RankHaloExchange,
cache::DOFBasedCOOCache,
assembler::DOFBasedCOOAssembler,
kernel::AbstractKernel,
mesh::AbstractMesh;
fill_recv_from_global::Bool = true,
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
)
exchange.part == layout.part ||
throw(ArgumentError("exchange.part $(exchange.part) != layout.part $(layout.part)"))
@@ -164,11 +178,149 @@ function partitioned_owned_matvec!(
unpack_halo_recv_to_packed!(packed, recv_vals, exchange, layout)
fill!(work, 0.0)
expand_packed_to_global!(work, packed, layout)
apply_K_owned_rows!(y_contrib, layout.owned_rows, cache, assembler, kernel, mesh, work)
apply_K_owned_rows!(
y_contrib, layout.owned_rows, cache, assembler, mesh, work;
configuration = configuration,
global_material_cache = global_material_cache,
Δt = Δt,
)
return y_contrib
end
function partitioned_owned_matvec!(
y_contrib::AbstractVector{Float64},
x_global::AbstractVector{Float64},
packed::AbstractVector{Float64},
work::AbstractVector{Float64},
recv_vals::Vector{Vector{Float64}},
layout::PartitionPackedLayout,
exchange::RankHaloExchange,
cache::DOFBasedCOOCache,
assembler::DOFBasedCOOAssembler,
::AbstractKernel,
mesh::AbstractMesh;
fill_recv_from_global::Bool = true,
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
)
_depwarn_redundant_kernel_arg!(:partitioned_owned_matvec!)
return partitioned_owned_matvec!(
y_contrib, x_global, packed, work, recv_vals, layout, exchange, cache, assembler, mesh;
fill_recv_from_global = fill_recv_from_global,
configuration = configuration,
global_material_cache = global_material_cache,
Δt = Δt,
)
end
"""
partitioned_owned_internal_force!(
y_contrib, u_global, packed, work, recv_vals,
layout, exchange, cache, assembler, mesh;
fill_recv_from_global = true,
configuration = nothing, global_material_cache = nothing, Δt = 0.0,
) -> y_contrib
partitioned_owned_internal_force!(
y_contrib, u_global, packed, work, recv_vals,
layout, exchange, cache, assembler, kernel, mesh;
fill_recv_from_global = true,
configuration = nothing, global_material_cache = nothing, Δt = 0.0,
) -> y_contrib
One partition’s additive contribution to the nonlinear internal force
[`assemble_internal_force!`](@ref) on **owned rows** (zeros elsewhere in
`y_contrib`), using the same packed patch and halo metadata as
[`partitioned_owned_matvec!`](@ref):
1. [`gather_owned_from_global_to_packed!`](@ref)`(packed, u_global, layout)`
2. If `fill_recv_from_global` — [`simulate_halo_recv_from_global!`](@ref)`(recv_vals, u_global, exchange)`
3. [`unpack_halo_recv_to_packed!`](@ref)
4. [`fill!`](@ref)`(work, 0)` then [`expand_packed_to_global!`](@ref)`(work, packed, layout)`
5. [`apply_f_int_owned_rows!`](@ref)`(y_contrib, layout.owned_rows, …; configuration, …)`
When `configuration === nothing`, Pass~1 uses the expanded patch displacement
`work`. Otherwise `configuration` is forwarded to [`_prepare_caches!`](@ref)
explicitly (length `cache.ndofs`).
The form with `kernel` ignores it (backward compatibility).
"""
function partitioned_owned_internal_force!(
y_contrib::AbstractVector{Float64},
u_global::AbstractVector{Float64},
packed::AbstractVector{Float64},
work::AbstractVector{Float64},
recv_vals::Vector{Vector{Float64}},
layout::PartitionPackedLayout,
exchange::RankHaloExchange,
cache::DOFBasedCOOCache,
assembler::DOFBasedCOOAssembler,
mesh::AbstractMesh;
fill_recv_from_global::Bool = true,
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
)
exchange.part == layout.part ||
throw(ArgumentError("exchange.part $(exchange.part) != layout.part $(layout.part)"))
nd = cache.ndofs
length(y_contrib) == nd ||
throw(DimensionMismatch("y_contrib length $(length(y_contrib)), ndofs $nd"))
length(u_global) == nd ||
throw(DimensionMismatch("u_global length $(length(u_global)), ndofs $nd"))
length(work) == nd ||
throw(DimensionMismatch("work length $(length(work)), ndofs $nd"))
layout.ndofs_global == nd ||
throw(DimensionMismatch("layout.ndofs_global $(layout.ndofs_global), cache.ndofs $nd"))
gather_owned_from_global_to_packed!(packed, u_global, layout)
if fill_recv_from_global
simulate_halo_recv_from_global!(recv_vals, u_global, exchange)
end
unpack_halo_recv_to_packed!(packed, recv_vals, exchange, layout)
fill!(work, 0.0)
expand_packed_to_global!(work, packed, layout)
prep_cfg = configuration === nothing ? work : configuration
apply_f_int_owned_rows!(
y_contrib, layout.owned_rows, cache, assembler, mesh;
configuration = prep_cfg,
global_material_cache = global_material_cache,
Δt = Δt,
)
return y_contrib
end
function partitioned_owned_internal_force!(
y_contrib::AbstractVector{Float64},
u_global::AbstractVector{Float64},
packed::AbstractVector{Float64},
work::AbstractVector{Float64},
recv_vals::Vector{Vector{Float64}},
layout::PartitionPackedLayout,
exchange::RankHaloExchange,
cache::DOFBasedCOOCache,
assembler::DOFBasedCOOAssembler,
::AbstractKernel,
mesh::AbstractMesh;
fill_recv_from_global::Bool = true,
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
)
_depwarn_redundant_kernel_arg!(:partitioned_owned_internal_force!)
return partitioned_owned_internal_force!(
y_contrib, u_global, packed, work, recv_vals, layout, exchange, cache, assembler, mesh;
fill_recv_from_global = fill_recv_from_global,
configuration = configuration,
global_material_cache = global_material_cache,
Δt = Δt,
)
end
"""
apply_K_owned_rows_from_packed!(
Ap_owned, packed, layout, cache, assembler, mesh,
) -> Ap_owned
apply_K_owned_rows_from_packed!(
Ap_owned, packed, layout, cache, assembler, kernel, mesh,
) -> Ap_owned
@@ -183,16 +335,20 @@ Requires every stencil neighbor `j` of those rows to satisfy `layout.global_to_p
`length(Ap_owned) == layout.n_owned`, `length(packed) ≥ layout.n_packed`, `layout.ndofs_global == cache.ndofs`.
Allocation-free after warmup.
The form with `kernel` ignores it (backward compatibility).
"""
function apply_K_owned_rows_from_packed!(
Ap_owned::AbstractVector{Float64},
packed::AbstractVector{Float64},
layout::PartitionPackedLayout,
cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType},
cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType,KS},
assembler::DOFBasedCOOAssembler,
kernel::AbstractKernel,
mesh::AbstractMesh,
) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType}
mesh::AbstractMesh;
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType,KS}
ndofs = cache.ndofs
layout.ndofs_global == ndofs ||
throw(DimensionMismatch(
@@ -212,7 +368,12 @@ function apply_K_owned_rows_from_packed!(
g2p = layout.global_to_packed
p2g = layout.packed_to_global
_prepare_caches!(cache, kernel, mesh)
_prepare_caches!(
cache, mesh;
configuration = configuration,
global_material_cache = global_material_cache,
Δt = Δt,
)
elements = cache.elements
element_caches = cache.element_caches
@@ -246,13 +407,14 @@ function apply_K_owned_rows_from_packed!(
entry_i = loc_layout[local_i]
dofs_elem = element_cache.dofs
k_e = kernel_at(cache, Int(elem_id_val))
@inbounds for local_j in 1:ndofs_elem
dof_j_global = Int(dofs_elem[local_j])
entry_j = loc_layout[local_j]
K_ij = evaluate_entry(
kernel,
k_e,
geometry_cache,
qp_buffer,
entry_i,
@@ -274,3 +436,24 @@ function apply_K_owned_rows_from_packed!(
return Ap_owned
end
function apply_K_owned_rows_from_packed!(
Ap_owned::AbstractVector{Float64},
packed::AbstractVector{Float64},
layout::PartitionPackedLayout,
cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType,KS},
assembler::DOFBasedCOOAssembler,
::AbstractKernel,
mesh::AbstractMesh;
configuration::Union{Nothing,AbstractVector{Float64}} = nothing,
global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing,
Δt::Float64 = 0.0,
) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType,KS}
_depwarn_redundant_kernel_arg!(:apply_K_owned_rows_from_packed!)
return apply_K_owned_rows_from_packed!(
Ap_owned, packed, layout, cache, assembler, mesh;
configuration = configuration,
global_material_cache = global_material_cache,
Δt = Δt,
)
end