From eeecf726568f34cf622feb82d0dc2d2a565c85e2 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Thu, 20 Nov 2025 16:56:38 +0200 Subject: [PATCH] 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 --- src/assemblers/coo_cache.jl | 184 ++++++++++++++++++++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 src/assemblers/coo_cache.jl diff --git a/src/assemblers/coo_cache.jl b/src/assemblers/coo_cache.jl new file mode 100644 index 0000000..035cc43 --- /dev/null +++ b/src/assemblers/coo_cache.jl @@ -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