From 6786111fd579d920e97c12a1ab38accff10b37c7 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 18:37:30 +0300 Subject: [PATCH] test(assemblers): add DOF assembler zero-allocation regression --- test/assemblers/test_dof_based_zero_alloc.jl | 238 +++++++++++++++++++ 1 file changed, 238 insertions(+) create mode 100644 test/assemblers/test_dof_based_zero_alloc.jl diff --git a/test/assemblers/test_dof_based_zero_alloc.jl b/test/assemblers/test_dof_based_zero_alloc.jl new file mode 100644 index 0000000..84297cd --- /dev/null +++ b/test/assemblers/test_dof_based_zero_alloc.jl @@ -0,0 +1,238 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Regression test: DOF-based assembler must stay zero-allocation +and bit-equivalent to the element-based assembler. + +This test exists to lock in the December 2025 zero-allocation work +and to catch any regression as the assembler evolves. + +Each problem size: +1. Build a structured Hex8 / Tet4 mesh +2. Create elements + DOFManager +3. Build both assemblers +4. Warm up +5. Assert: 0 bytes allocated by the DOF-based `assemble!` +6. Assert: matrix matches the element-based assembler to round-off +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache +using JuliaFEM: COOAssembler, create_cache +using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex +using JuliaFEM: local_dof_layout, DOFLayoutEntry +using LinearAlgebra +using SparseArrays +using Tensors +using InteractiveUtils # @code_typed, code_llvm + +# ---------------------------------------------------------------------------- +# Helpers (kept local to this test, no leak into package) +# ---------------------------------------------------------------------------- + +""" +Build a structured Hex8 box mesh with `nx × ny × nz` elements over [0,1]^3. +""" +function _build_hex8_box_mesh(nx::Int, ny::Int, nz::Int) + nodes = Vec{3,Float64}[] + nidx(i, j, k) = (i - 1) + (j - 1) * (nx + 1) + (k - 1) * (nx + 1) * (ny + 1) + 1 + for k in 1:(nz + 1), j in 1:(ny + 1), i in 1:(nx + 1) + push!(nodes, Vec{3}((Float64(i - 1) / nx, + Float64(j - 1) / ny, + Float64(k - 1) / nz))) + end + conns = NTuple{8,UInt32}[] + for k in 1:nz, j in 1:ny, i in 1:nx + n1 = nidx(i, j, k) + n2 = nidx(i + 1, j, k) + n3 = nidx(i + 1, j + 1, k) + n4 = nidx(i, j + 1, k) + n5 = nidx(i, j, k + 1) + n6 = nidx(i + 1, j, k + 1) + n7 = nidx(i + 1, j + 1, k + 1) + n8 = nidx(i, j + 1, k + 1) + push!(conns, (UInt32(n1), UInt32(n2), UInt32(n3), UInt32(n4), + UInt32(n5), UInt32(n6), UInt32(n7), UInt32(n8))) + end + return Mesh{8,Hexahedron{8}}(nodes, conns) +end + +""" +Build a single Tet4 mesh. +""" +function _build_single_tet4_mesh() + nodes = Vec{3,Float64}[ + Vec{3}((0.0, 0.0, 0.0)), + Vec{3}((1.0, 0.0, 0.0)), + Vec{3}((0.5, 1.0, 0.0)), + Vec{3}((0.5, 0.5, 1.0)), + ] + conns = [(UInt32(1), UInt32(2), UInt32(3), UInt32(4))] + return Mesh{Tetrahedron{4}}(nodes, conns) +end + +""" +Setup elasticity assembly fixture for a given mesh + topology. +Returns (cache_dof, asm_dof, cache_elem, asm_elem, kernel, mesh). +""" +function _setup_assembly(mesh, ::Type{Topo}) where {Topo} + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, dof_mgr = create_elements!(mesh, Element{Topo, Lagrange{1}, S}) + + asm_dof = DOFBasedCOOAssembler() + cache_dof = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + + asm_elem = COOAssembler() + cache_elem = create_cache(asm_elem, mesh, kernel) + + return cache_dof, asm_dof, cache_elem, asm_elem, kernel, mesh +end + +# ---------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------- + +@testset "DOF-based assembler: zero allocation regression" begin + println("\n" * "=" ^ 70) + println("DOF-BASED ASSEMBLER ZERO-ALLOCATION REGRESSION") + println("=" ^ 70) + + # ------------------------------------------------------------------------ + # 1. Single Tet4 + # ------------------------------------------------------------------------ + @testset "Single Tet4" begin + mesh = _build_single_tet4_mesh() + cache_dof, asm_dof, cache_elem, asm_elem, kernel, m = + _setup_assembly(mesh, Tetrahedron{4}) + + # warmup + assemble!(cache_dof, asm_dof, kernel, m) + assemble!(cache_elem, asm_elem, kernel, m) + + K_dof, _ = extract_system(cache_dof) + K_elem, _ = extract_system(cache_elem) + + # Equivalence + rel = maximum(abs, Matrix(K_dof) - Matrix(K_elem)) / + max(maximum(abs, Matrix(K_elem)), 1.0) + @test rel < 1e-12 + + # Zero-alloc + GC.gc() + a = @allocated assemble!(cache_dof, asm_dof, kernel, m) + @test a == 0 + println(" Single Tet4 ........................ allocs=$a, rel=$(round(rel; sigdigits=3))") + end + + # ------------------------------------------------------------------------ + # 2. Growing Hex8 cube meshes + # ------------------------------------------------------------------------ + @testset "Hex8 cube $(nx)×$(ny)×$(nz)" for (nx, ny, nz) in + [(1, 1, 1), (2, 1, 1), (4, 2, 2), (6, 3, 3), (8, 4, 4)] + + mesh = _build_hex8_box_mesh(nx, ny, nz) + cache_dof, asm_dof, cache_elem, asm_elem, kernel, m = + _setup_assembly(mesh, Hexahedron{8}) + + # warmup + assemble!(cache_dof, asm_dof, kernel, m) + assemble!(cache_elem, asm_elem, kernel, m) + + K_dof, _ = extract_system(cache_dof) + K_elem, _ = extract_system(cache_elem) + + rel = maximum(abs, Matrix(K_dof) - Matrix(K_elem)) / + max(maximum(abs, Matrix(K_elem)), 1.0) + @test rel < 1e-12 + + GC.gc() + a = @allocated assemble!(cache_dof, asm_dof, kernel, m) + @test a == 0 + + nelems = length(m.connectivity) + ndofs = 3 * length(m.nodes) + println(" Hex8 $(nx)×$(ny)×$(nz) $(lpad(nelems,5)) elem " * + "$(lpad(ndofs,5)) dof allocs=$a rel=$(round(rel; sigdigits=3))") + end +end + +# ---------------------------------------------------------------------------- +# Type stability check (no untyped ::Any in inferred output) +# ---------------------------------------------------------------------------- +@testset "DOF-based assembler: type stability" begin + mesh = _build_hex8_box_mesh(2, 1, 1) + cache_dof, asm_dof, _, _, kernel, m = _setup_assembly(mesh, Hexahedron{8}) + assemble!(cache_dof, asm_dof, kernel, m) # warmup + + # Inferred return type must be concrete (Nothing). + rt = Base.promote_op(assemble!, typeof(cache_dof), typeof(asm_dof), + typeof(kernel), typeof(m)) + @test rt === Nothing + + # No method ambiguities or `Any` in the top-level signature + code = code_typed(assemble!, (typeof(cache_dof), typeof(asm_dof), + typeof(kernel), typeof(m)); optimize = true) + @test !isempty(code) + info = code[1] + @test isconcretetype(info.second) + println(" assemble! inferred return type: $(info.second)") +end + +# ---------------------------------------------------------------------------- +# Element-as-template: local_dof_layout(E) is compile-time constant and +# `assemble!` has zero GC allocation sites in the optimized LLVM IR. +# ---------------------------------------------------------------------------- +@testset "DOF-based assembler: element template + LLVM allocs" begin + mesh = _build_hex8_box_mesh(2, 1, 1) + cache_dof, asm_dof, _, _, kernel, m = _setup_assembly(mesh, Hexahedron{8}) + assemble!(cache_dof, asm_dof, kernel, m) # warmup + + # ------------------------------------------------------------------ + # 1. local_dof_layout returns the expected NTuple of DOFLayoutEntry + # for an Hex8 displacement element (24 DOFs, 8 vertices × 3 comp.) + # ------------------------------------------------------------------ + ET = eltype(cache_dof.elements) + layout = local_dof_layout(ET) + @test layout isa NTuple{24, DOFLayoutEntry} + @test all(e -> e.field_idx == 1, layout) + @test [Int(e.entity_local) for e in layout] == + repeat(1:8; inner = 3) + @test [Int(e.component) for e in layout] == + repeat(1:3, 8) + + # ------------------------------------------------------------------ + # 2. local_dof_layout(ET) call is constant-folded by the compiler: + # @allocated must be 0 and the optimized typed-IR must report a + # `Core.Const` for the layout in `assemble!` — exercised already + # by the zero-allocation tests above. We assert separately here. + # ------------------------------------------------------------------ + GC.gc() + a_layout = @allocated local_dof_layout(ET) + @test a_layout == 0 + + # ------------------------------------------------------------------ + # 3. Optimized LLVM IR for `assemble!` must have ZERO GC allocation + # sites. This is the strongest guarantee: even if @allocated == 0 + # today, this catches any future change that would re-introduce + # a heap-allocating call inside the inner loop. + # ------------------------------------------------------------------ + iob = IOBuffer() + code_llvm(iob, assemble!, + Tuple{typeof(cache_dof), typeof(asm_dof), + typeof(kernel), typeof(m)}; optimize = true) + ir = String(take!(iob)) + n_alloc = + length(collect(eachmatch(r"call.*julia\.gc_alloc", ir))) + + length(collect(eachmatch(r"call.*jl_gc_pool_alloc", ir))) + + length(collect(eachmatch(r"call.*jl_gc_big_alloc", ir))) + + length(collect(eachmatch(r"call.*jl_gc_alloc_typed", ir))) + @test n_alloc == 0 + println(" assemble! optimized LLVM IR: $n_alloc gc-alloc sites " * + "($(length(ir)) IR chars)") +end