From 08f940f3dbaafdd3aa327b7ba07e2f18f20e8311 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Thu, 20 Nov 2025 16:56:39 +0200 Subject: [PATCH] feat(assemblers): Add node-based COO assembly implementation New file: src/assemblers/node_based_coo.jl Features: - assemble! implementation for nodal assembly - Assembles contributions node-by-node instead of element-by-element - Uses node_to_elements connectivity - Accumulates blocks for all elements touching each node Architecture: - Outer loop over nodes (not elements) - Inner loop over elements containing each node - Natural for contact mechanics (contact is nodal) Status: Experimental, proof-of-concept implementation. Not yet optimized like element-based assembly. --- src/assemblers/node_based_coo.jl | 484 +++++++++++++++++++++++++++++++ 1 file changed, 484 insertions(+) create mode 100644 src/assemblers/node_based_coo.jl diff --git a/src/assemblers/node_based_coo.jl b/src/assemblers/node_based_coo.jl new file mode 100644 index 0000000..62557bf --- /dev/null +++ b/src/assemblers/node_based_coo.jl @@ -0,0 +1,484 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Node-based COO assembly using block integration. + +**NODAL ASSEMBLY PARADIGM**: Loop over nodes, not elements! + +Each node: +1. Finds all elements touching it (via inverse connectivity) +2. For each touching element: + - Prepares element geometry once (PreparedElement) + - Computes only needed 3×3 blocks (compute_block!) +3. Scatters blocks to COO triplets + +# Key Differences from Element-Based Assembly + +**Element-Based (traditional):** +```julia +for element in elements + K_e = compute_element_stiffness(element) # Full N×N matrix of 3×3 blocks + scatter(K_e) # Scatter all entries +end +``` + +**Node-Based (this file):** +```julia +for node_i in nodes + for element in elements_touching(node_i) + prepared = prepare_element(element) # Geometry preprocessing + for node_j in element.nodes + K_ij = compute_block!(prepared, i, j) # Single 3×3 block + scatter(K_ij, i, j) # Scatter one block + end + end +end +``` + +# Advantages + +1. **GPU-friendly**: One thread per node, no race conditions +2. **Contact-ready**: Contact is naturally node-based +3. **Matrix-free ready**: Can compute K*v without forming K +4. **Cache-friendly**: Reuses PreparedElement for multiple blocks +5. **Adaptive-ready**: Easy to refine/coarsen at node level + +# Performance Expectations + +- **CPU Single-thread**: ~1.5-2x slower than element-based (more kernel calls) +- **CPU Multi-thread**: ~1.5-2x faster (better parallelization) +- **GPU**: ~10-50x faster (massive parallelization, no atomics needed) + +# References + +- Golden standard: `docs/src/book/multigpu_nodal_assembly.md` +- PreparedElement: `src/domains/continuum/integration.jl` +- Block kernel: `src/domains/continuum/kernel.jl` + +# Example + +```julia +# Setup +mesh = create_cantilever_mesh(50, 10, 10) +material = LinearElastic(E=210e9, ν=0.3) +kernel = ContinuumKernel( + ContinuumFormulation{FullThreeD}(), + material, + Displacement{3}() +) + +# Create node-based assembler and cache +assembler = NodeBasedCOOAssembler() +cache = create_cache(assembler, mesh, kernel) + +# Assemble (zero allocations after warmup!) +assemble!(cache, assembler, kernel, mesh) + +# Extract system +K, f = extract_system(cache) + +# Solve +apply_dirichlet_bcs!(K, f, kernel, mesh, bc_dirichlet) +u = K \\ f +``` +""" + +using SparseArrays +using Tensors + +""" + NodeBasedCOOCache + +Pre-allocated cache for node-based COO assembly. + +Similar to COOCache but includes inverse connectivity mapping. + +# Fields +- `I::Vector{Int}`: Row indices (COO format) +- `J::Vector{Int}`: Column indices (COO format) +- `V::Vector{Float64}`: Values (COO format) +- `f::Vector{Float64}`: Global force vector +- `counter::Ref{Int}`: Current triplet count +- `capacity::Int`: Maximum triplet capacity +- `node_to_elements::NodeToElementsMap`: Inverse connectivity +- `element_cache::ElementCache`: Cache for element operations +- `ndofs::Int`: Total DOFs in system +""" +struct NodeBasedCOOCache{T<:AbstractTopology,B<:AbstractBasis,IPS} + I::Vector{Int} + J::Vector{Int} + V::Vector{Float64} + f::Vector{Float64} + counter::Ref{Int} + capacity::Int + node_to_elements::NodeToElementsMap + element_cache::ElementCache{T,B,IPS} + ndofs::Int +end + +""" + NodeBasedCOOCache(mesh::AbstractMesh, kernel::ContinuumKernel) + +Create cache for node-based assembly. + +Builds inverse connectivity and allocates buffers. + +# Arguments +- `mesh`: Finite element mesh +- `kernel`: Continuum kernel + +# Returns +- Pre-allocated node-based COO cache +""" +function NodeBasedCOOCache(mesh::AbstractMesh, kernel::ContinuumKernel) + # Build inverse connectivity + node_to_elements = NodeToElementsMap(mesh.connectivity) + + # Estimate triplet count (same as element-based) + ndofs_per_node = dofs_per_node(kernel) + nnodes = length(mesh.nodes) + ndofs = ndofs_per_node * nnodes + + # Estimate: For each node, sum over touching elements + # Each element contributes N blocks (N = nodes per element) + # Each block = 3×3 = 9 triplets + avg_elements_per_node = node_to_elements.nelements / nnodes + N = length(first(mesh.connectivity)) # Nodes per element + estimated_triplets = Int(ceil(1.2 * nnodes * avg_elements_per_node * N * 9)) + + # Allocate triplet arrays + I = Vector{Int}(undef, estimated_triplets) + J = Vector{Int}(undef, estimated_triplets) + V = Vector{Float64}(undef, estimated_triplets) + f = zeros(Float64, ndofs) + counter = Ref(0) + + # Create element cache (for prepare_element! and compute_block!) + element_cache = ElementCache(mesh, kernel) + + return NodeBasedCOOCache(I, J, V, f, counter, estimated_triplets, + node_to_elements, element_cache, ndofs) +end + +""" + reset!(cache::NodeBasedCOOCache) + +Reset cache for new assembly (zero force vector, reset counter). + +Does NOT clear inverse connectivity (that's permanent structure). +""" +function reset!(cache::NodeBasedCOOCache) + fill!(cache.f, 0.0) + cache.counter[] = 0 + return nothing +end + +""" + assemble!( + cache::NodeBasedCOOCache, + assembler::NodeBasedCOOAssembler, + kernel::ContinuumKernel, + mesh::AbstractMesh + ) -> Nothing + +Assemble global system using **node-based traversal**. + +# Algorithm + +```julia +for node_i in 1:nnodes + # Get all elements touching this node + for elem_info in node_to_elements[node_i] + element_id = elem_info.element_id + local_i = elem_info.local_node_idx + + # Prepare element geometry ONCE + prepared = prepare_element!(cache.element_cache, kernel, element_id, mesh) + + # Compute blocks for all nodes in this element + for local_j in 1:N + global_j = connectivity[element_id][local_j] + + # Compute single 3×3 block + K_ij = compute_block!(prepared, kernel.material, local_i, local_j) + + # Scatter to triplets + scatter_block_to_triplets!(cache, K_ij, node_i, global_j) + end + end +end +``` + +# Key Operations + +1. **prepare_element!** - Precompute geometry (Jacobian, gradients) once per element +2. **compute_block!** - Compute single 3×3 stiffness block using prepared geometry +3. **scatter_block_to_triplets!** - Add 9 triplets (i,j,value) for 3×3 block + +# Zero-Allocation (After Warmup) + +All arrays pre-allocated. Element preparation reuses cache buffers. + +# Arguments +- `cache`: Pre-allocated node-based COO cache +- `assembler`: Node-based COO assembler +- `kernel`: Continuum kernel +- `mesh`: Finite element mesh +""" +function assemble!( + cache::NodeBasedCOOCache, + assembler::NodeBasedCOOAssembler, + kernel::ContinuumKernel, + mesh::AbstractMesh +) + # Reset cache + reset!(cache) + + nnodes = length(mesh.nodes) + ndofs_per_node = dofs_per_node(kernel) + + # NODAL LOOP: One iteration per node (GPU: one thread per node!) + for node_i in 1:nnodes + # Get all elements touching this node + touching_elements = cache.node_to_elements.node_to_elements[node_i] + + # Loop over touching elements + for elem_info in touching_elements + element_id = elem_info.element_id + local_i = elem_info.local_node_idx # Position of node_i in element + + # Prepare element geometry ONCE (reuses cache.element_cache) + prepared = prepare_element!(cache.element_cache, kernel, element_id, mesh) + + # Get element connectivity + conn = mesh.connectivity[element_id] + N = length(conn) # Nodes per element + + # Compute blocks for all nodes j in this element + for local_j in 1:N + global_j = conn[local_j] + + # Compute single 3×3 block K[i,j] + # This is THE KEY OPERATION: block-based integration + K_ij = compute_block!( + prepared, + kernel.material, + local_i, + local_j + ) + + # Scatter 3×3 block to triplets (adds 9 entries) + scatter_block_to_triplets!( + cache, + K_ij, + node_i, + global_j, + ndofs_per_node + ) + end + end + end + + return nothing +end + +""" + scatter_block_to_triplets!( + cache::NodeBasedCOOCache, + K_block::Tensor{2,3}, + node_i::Int, + node_j::Int, + ndofs_per_node::Int + ) + +Scatter single 3×3 block to COO triplets **in-place**. + +Maps block[α,β] → triplet at DOF indices: +- Row: 3*(node_i-1) + α +- Col: 3*(node_j-1) + β +- Val: K_block[α,β] + +# Arguments +- `cache`: Node-based COO cache +- `K_block`: 3×3 stiffness block (Tensor{2,3}) +- `node_i`: Global row node index +- `node_j`: Global column node index +- `ndofs_per_node`: DOFs per node (typically 3) + +# Zero-Allocation + +Writes to pre-allocated triplet arrays, updates counter. +""" +function scatter_block_to_triplets!( + cache::NodeBasedCOOCache, + K_block::Tensor{2,3,Float64}, + node_i::Int, + node_j::Int, + ndofs_per_node::Int +) + counter = cache.counter[] + + # Check capacity + new_triplets = ndofs_per_node * ndofs_per_node # 3×3 = 9 + if counter + new_triplets > cache.capacity + error("Node-based COO cache overflow: need $(counter + new_triplets) triplets, " * + "capacity is $(cache.capacity). Increase cache size.") + end + + # DOF offsets for nodes i and j + row_offset = ndofs_per_node * (node_i - 1) + col_offset = ndofs_per_node * (node_j - 1) + + # Scatter 3×3 block to triplets + for β in 1:ndofs_per_node # Column (node j DOF) + j_global = col_offset + β + for α in 1:ndofs_per_node # Row (node i DOF) + i_global = row_offset + α + counter += 1 + cache.I[counter] = i_global + cache.J[counter] = j_global + cache.V[counter] = K_block[α, β] + end + end + + cache.counter[] = counter + return nothing +end + +""" + extract_system(cache::NodeBasedCOOCache) -> (K, f) + +Build sparse matrix from triplets and return system. + +Calls `sparse(I, J, V)` to build CSC matrix. Duplicates are summed automatically. + +# Arguments +- `cache`: Assembled node-based COO cache + +# Returns +- `K::SparseMatrixCSC`: Global stiffness matrix +- `f::Vector`: Global force vector + +# Allocation + +Allocates sparse matrix structure (CSC format). This is the only allocation +outside cache construction. +""" +function extract_system(cache::NodeBasedCOOCache) + ntriplets = cache.counter[] + + # Build sparse matrix (duplicates are summed automatically) + I_used = @view cache.I[1:ntriplets] + J_used = @view cache.J[1:ntriplets] + V_used = @view cache.V[1:ntriplets] + + K = sparse(I_used, J_used, V_used, cache.ndofs, cache.ndofs) + + return K, cache.f +end + +# ============================================================================ +# HELPER FUNCTIONS +# ============================================================================ + +""" + create_cache( + assembler::NodeBasedCOOAssembler, + mesh::AbstractMesh, + kernel::ContinuumKernel + ) -> NodeBasedCOOCache + +Create pre-allocated cache for node-based COO assembly. + +Convenience function that wraps `NodeBasedCOOCache(mesh, kernel)`. + +# Example + +```julia +assembler = NodeBasedCOOAssembler() +cache = create_cache(assembler, mesh, kernel) +assemble!(cache, assembler, kernel, mesh) +K, f = extract_system(cache) +``` +""" +function create_cache( + assembler::NodeBasedCOOAssembler, + mesh::AbstractMesh, + kernel::ContinuumKernel +) + return NodeBasedCOOCache(mesh, kernel) +end + +""" + dofs_per_node(kernel::ContinuumKernel) -> Int + +Return DOFs per node for continuum kernel (always 3 for displacement). + +Dispatches on kernel field dimension. +""" +function dofs_per_node(kernel::ContinuumKernel{Theory,Mat}) where {Theory,Mat} + field = kernel.field + return field.dim # Displacement{3} → 3 +end + +# ============================================================================ +# PERFORMANCE NOTES +# ============================================================================ + +#= +# CPU Performance Comparison (Estimated) + +**Element-Based Assembly:** +- Elements: 1000 Tet4 +- Operations: 1000 elements × 4×4 blocks × 3×3 entries = 48,000 block computations +- Time: ~5ms (baseline) + +**Node-Based Assembly:** +- Nodes: 500 nodes +- Operations: 500 nodes × 8 elements/node × 4 blocks/element = 16,000 block computations +- But: 3× more kernel calls due to overlaps +- Time: ~7-10ms (1.5-2× slower single-threaded) + +**Why slower on CPU?** +- Each block computed once in element assembly +- Each block computed 2× on average in nodal assembly (shared between 2 elements) +- More function call overhead + +**Why faster on GPU?** +- Element assembly: Sequential (can't parallelize over elements efficiently) +- Nodal assembly: Massive parallelism (one thread per node) +- GPU speedup: ~10-50× depending on problem size + +**Multi-threaded CPU (Threads.@threads):** +- Can parallelize outer node loop +- Expected speedup: 1.5-2× over element-based +- No race conditions (each node writes different triplets) + +# Memory Comparison + +**Element-Based:** +- Triplet storage: ~50 KB per 1000 elements +- Element cache: ~2 KB per thread + +**Node-Based:** +- Triplet storage: Same (~50 KB) +- Element cache: ~2 KB per thread +- Inverse connectivity: ~10-20 KB (one-time) + +→ Nearly identical memory usage! + +# When to Use Node-Based Assembly + +**Use when:** +- ✅ GPU acceleration needed +- ✅ Contact mechanics (naturally nodal) +- ✅ Matrix-free methods (K*v without forming K) +- ✅ Adaptive refinement (local node operations) +- ✅ Multi-threading on CPU + +**Don't use when:** +- ❌ Single-threaded CPU only +- ❌ Simple problems (< 1000 nodes) +- ❌ Prototyping/debugging (element-based is clearer) +=#