mirror of
https://github.com/JuliaFEM/JuliaFEM.jl.git
synced 2026-09-17 17:22:10 +00:00
feat(assemblers): Add parametric COOCache for zero-allocation assembly
New file: src/assemblers/coo_cache.jl (184 lines)
Features:
- Parametric struct COOCache{EC<:ElementCache, MC<:MaterialStateCache}
- Eliminates type instability from cache field accesses
- Stores triplets (I, J, V) for sparse matrix construction
- Includes reset! and extract_system functions
Performance impact:
- Enables zero allocations in assembly loop
- Required for achieving 500K elem/s throughput
- Critical optimization for type stability
Documentation includes:
- COO format explanation
- Performance characteristics
- Use cases and trade-offs
This commit is contained in:
@@ -0,0 +1,184 @@
|
||||
# This file is a part of JuliaFEM.
|
||||
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
|
||||
|
||||
"""
|
||||
COO (Coordinate format) cache for element-based assembly.
|
||||
|
||||
COO format stores sparse matrices as triplets (I, J, V) where:
|
||||
- I[k] = row index of k-th entry
|
||||
- J[k] = column index of k-th entry
|
||||
- V[k] = value of k-th entry
|
||||
|
||||
After assembly, triplets are converted to sparse matrix using `sparse(I, J, V)`.
|
||||
|
||||
# Performance
|
||||
- Fast assembly (no structure lookups)
|
||||
- Slow sparse matrix construction (O(nnz log nnz) for sorting)
|
||||
- Memory overhead (stores all triplets including duplicates)
|
||||
|
||||
# Use case
|
||||
Good for problems where sparsity pattern changes (e.g., contact, topology optimization).
|
||||
"""
|
||||
|
||||
using SparseArrays
|
||||
|
||||
"""
|
||||
COOCache <: AbstractAssemblerCache
|
||||
|
||||
Cache for COO (coordinate format) assembly.
|
||||
|
||||
Pre-allocates triplet vectors `(I, J, V)` and workspace for element assembly.
|
||||
After assembly, triplets are converted to sparse matrix using `sparse(I, J, V)`.
|
||||
|
||||
# Fields
|
||||
- `I::Vector{Int}`: Row indices (pre-allocated, max capacity)
|
||||
- `J::Vector{Int}`: Column indices (pre-allocated, max capacity)
|
||||
- `V::Vector{Float64}`: Values (pre-allocated, max capacity)
|
||||
- `f::Vector{Float64}`: Global force vector
|
||||
- `element_cache::ElementCache`: Per-element workspace
|
||||
- `geometry_cache::GeometryCache`: Geometry workspace
|
||||
- `material_cache::MaterialStateCache`: Material state workspace
|
||||
- `counter::Ref{Int}`: Current position in triplet arrays
|
||||
- `capacity::Int`: Maximum triplet capacity
|
||||
- `ndofs::Int`: Total number of DOFs
|
||||
|
||||
# Zero-Allocation Usage
|
||||
|
||||
```julia
|
||||
cache = COOCache(mesh, kernel)
|
||||
fill!(cache) # Reset counter, zero arrays
|
||||
assemble!(cache, assembler, kernel, mesh) # No allocations
|
||||
K, f = extract_system(cache) # Build sparse matrix
|
||||
```
|
||||
"""
|
||||
mutable struct COOCache{EC<:ElementCache,MC<:MaterialStateCache} <: AbstractAssemblerCache
|
||||
I::Vector{Int} # Row indices
|
||||
J::Vector{Int} # Column indices
|
||||
V::Vector{Float64} # Values
|
||||
f::Vector{Float64} # Force vector
|
||||
element_cache::EC # Element workspace (concrete type!)
|
||||
geometry_cache::GeometryCache # Geometry workspace
|
||||
material_cache::MC # Material state workspace (concrete type!)
|
||||
counter::Ref{Int} # Current triplet count
|
||||
capacity::Int # Maximum triplet capacity
|
||||
ndofs::Int # Total number of DOFs
|
||||
end
|
||||
|
||||
"""
|
||||
COOCache(mesh, kernel) -> COOCache
|
||||
|
||||
Create pre-allocated COO cache.
|
||||
|
||||
Estimates maximum triplet count based on mesh connectivity and DOF structure.
|
||||
Over-allocates by 20% to handle irregular meshes safely.
|
||||
|
||||
# Arguments
|
||||
- `mesh`: Finite element mesh
|
||||
- `kernel`: Domain kernel defining DOF structure
|
||||
|
||||
# Returns
|
||||
- `COOCache` with pre-allocated triplet arrays
|
||||
"""
|
||||
function COOCache(mesh::AbstractMesh, kernel::AbstractKernel)
|
||||
nelems = nelements(mesh)
|
||||
ndofs_per_node = dofs_per_node(kernel)
|
||||
nnodes_total_mesh = nnodes_total(mesh)
|
||||
ndofs = nnodes_total_mesh * ndofs_per_node
|
||||
|
||||
# Estimate triplet count: sum over elements of ndofs_elem^2
|
||||
# For uniform mesh: nelems * (nnodes_per_elem * ndofs_per_node)^2
|
||||
# Over-allocate by 20% for safety
|
||||
# For Mesh{N,T}, N is the first type parameter (nnodes_per_elem)
|
||||
MeshType = typeof(mesh)
|
||||
nnodes_elem = MeshType.parameters[1]::Int
|
||||
avg_ndofs_per_elem = Int(ceil(nnodes_elem * ndofs_per_node))
|
||||
estimated_triplets = Int(ceil(1.2 * nelems * avg_ndofs_per_elem^2))
|
||||
|
||||
I = zeros(Int, estimated_triplets)
|
||||
J = zeros(Int, estimated_triplets)
|
||||
V = zeros(Float64, estimated_triplets)
|
||||
f = zeros(Float64, ndofs)
|
||||
|
||||
# Create all caches
|
||||
element_cache = create_element_cache(mesh, kernel)
|
||||
|
||||
# Get max integration points for geometry and material caches
|
||||
max_nips = length(element_cache.ips)
|
||||
geometry_cache = create_geometry_cache(nnodes_elem, max_nips)
|
||||
material_cache = create_material_cache(kernel.material, max_nips)
|
||||
|
||||
return COOCache(I, J, V, f, element_cache, geometry_cache, material_cache,
|
||||
Ref(0), estimated_triplets, ndofs)
|
||||
end
|
||||
|
||||
"""
|
||||
reset!(cache::COOCache)
|
||||
|
||||
Reset COO cache for new assembly.
|
||||
|
||||
Zeros out triplet arrays and force vector, resets counter.
|
||||
**Zero allocations** - reuses existing arrays.
|
||||
"""
|
||||
function reset!(cache::COOCache)
|
||||
# Only zero up to current counter position (faster than fill!)
|
||||
current = cache.counter[]
|
||||
if current > 0
|
||||
@views cache.I[1:current] .= 0
|
||||
@views cache.J[1:current] .= 0
|
||||
@views cache.V[1:current] .= 0
|
||||
end
|
||||
fill!(cache.f, 0.0)
|
||||
cache.counter[] = 0
|
||||
return nothing
|
||||
end
|
||||
|
||||
"""
|
||||
extract_system(cache::COOCache) -> (K, f)
|
||||
|
||||
Extract global system from COO cache.
|
||||
|
||||
Builds sparse matrix from accumulated triplets. **Allocates** - only call
|
||||
once per assembly.
|
||||
|
||||
# Arguments
|
||||
- `cache`: COO cache after assembly
|
||||
|
||||
# Returns
|
||||
- `K`: Sparse matrix built from triplets
|
||||
- `f`: Force vector (reference, no copy)
|
||||
"""
|
||||
function extract_system(cache::COOCache)
|
||||
n = cache.counter[]
|
||||
I = @view cache.I[1:n]
|
||||
J = @view cache.J[1:n]
|
||||
V = @view cache.V[1:n]
|
||||
K = sparse(I, J, V, cache.ndofs, cache.ndofs)
|
||||
return K, cache.f
|
||||
end
|
||||
|
||||
"""
|
||||
create_cache(assembler::COOAssembler, mesh::AbstractMesh, kernel::AbstractKernel) -> COOCache
|
||||
|
||||
Create pre-allocated cache for COO assembly.
|
||||
|
||||
Convenience function that wraps `COOCache(mesh, kernel)`.
|
||||
|
||||
# Arguments
|
||||
- `assembler`: COO assembler
|
||||
- `mesh`: Finite element mesh
|
||||
- `kernel`: Domain kernel
|
||||
|
||||
# Returns
|
||||
- Pre-allocated COO cache
|
||||
|
||||
# Example
|
||||
|
||||
```julia
|
||||
cache = create_cache(COOAssembler(), mesh, kernel)
|
||||
assemble!(cache, COOAssembler(), kernel, mesh)
|
||||
K, f = extract_system(cache)
|
||||
```
|
||||
"""
|
||||
function create_cache(assembler::COOAssembler, mesh::AbstractMesh, kernel::AbstractKernel)
|
||||
return COOCache(mesh, kernel)
|
||||
end
|
||||
Reference in New Issue
Block a user