diff --git a/test/assemblers/test_dof_based_apply_K.jl b/test/assemblers/test_dof_based_apply_K.jl new file mode 100644 index 0000000..a63e5ce --- /dev/null +++ b/test/assemblers/test_dof_based_apply_K.jl @@ -0,0 +1,308 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Regression + correctness tests for `apply_K!` — the matrix-free `y = K x` +sibling of `assemble!` on the DOF-based assembler. + +For every problem size: + +1. Build assembled `K` via `assemble!` + `extract_system`. +2. Build the same `y_ref = K * x_random` for several random `x`. +3. Compare against `y_mf` from `apply_K!(y_mf, …, x)`. Must agree to + round-off (no accumulation difference: same DOF traversal, same + evaluate_entry inside). +4. Assert `apply_K!` is zero-allocation after warmup. +5. Assert the optimized LLVM IR for `apply_K!` has 0 GC allocation + sites — guarantees the inner accumulate loop never heap-allocates, + which is the whole point of the matrix-free path for Krylov. +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache, apply_K! +using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex +using LinearAlgebra +using SparseArrays +using Tensors +using Random +using InteractiveUtils # code_llvm +using LinearOperators # LinearOperator wrapper for apply_K! +using IterativeSolvers # cg + +# ---------------------------------------------------------------------------- +# Mesh helpers (shared shape with test_dof_based_zero_alloc.jl, kept local +# so the two regression tests stay independent). +# ---------------------------------------------------------------------------- + +function _build_hex8_box(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 + +function _build_single_tet4() + 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 + +"Set up DOF-based cache + kernel for elasticity." +function _setup(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 = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + return cache, asm, kernel, mesh +end + +# ---------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------- + +@testset "apply_K!: correctness vs assembled K" begin + println("\n" * "=" ^ 70) + println("DOF-BASED APPLY_K! CORRECTNESS") + println("=" ^ 70) + + Random.seed!(20260508) + + # ------------------------------------------------------------------ + # 1. Single Tet4 + # ------------------------------------------------------------------ + @testset "Single Tet4" begin + mesh = _build_single_tet4() + cache, asm, kernel, m = _setup(mesh, Tetrahedron{4}) + + # Assemble K once; compare K*x against apply_K!(y, …, x) for + # several random x. + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + n = size(K, 1) + + max_rel = 0.0 + for trial in 1:8 + x = randn(n) + y_ref = K * x + y_mf = zeros(n) + apply_K!(y_mf, cache, asm, kernel, m, x) + rel = norm(y_mf - y_ref) / max(norm(y_ref), 1.0) + @test rel < 1e-12 + max_rel = max(max_rel, rel) + end + println(" Single Tet4 ...... n=$n max rel=$(round(max_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(nx, ny, nz) + cache, asm, kernel, m = _setup(mesh, Hexahedron{8}) + + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + n = size(K, 1) + + max_rel = 0.0 + for trial in 1:5 + x = randn(n) + y_ref = K * x + y_mf = zeros(n) + apply_K!(y_mf, cache, asm, kernel, m, x) + rel = norm(y_mf - y_ref) / max(norm(y_ref), 1.0) + @test rel < 1e-12 + max_rel = max(max_rel, rel) + end + + nelems = length(m.connectivity) + println(" Hex8 $(nx)×$(ny)×$(nz) $(lpad(nelems, 5)) elem " * + "$(lpad(n, 5)) dof max rel=$(round(max_rel; sigdigits=3))") + end +end + +# ---------------------------------------------------------------------------- +# Zero-allocation + LLVM IR check for apply_K! +# ---------------------------------------------------------------------------- +@testset "apply_K!: zero allocation + 0 LLVM gc-alloc sites" begin + println("\n" * "=" ^ 70) + println("DOF-BASED APPLY_K! ZERO-ALLOC") + println("=" ^ 70) + + @testset "Single Tet4" begin + mesh = _build_single_tet4() + cache, asm, kernel, m = _setup(mesh, Tetrahedron{4}) + n = cache.ndofs + + x = ones(n) + y = zeros(n) + + # Warmup: compile + populate caches + apply_K!(y, cache, asm, kernel, m, x) + + GC.gc() + a = @allocated apply_K!(y, cache, asm, kernel, m, x) + @test a == 0 + println(" Single Tet4 ........................ allocs=$a") + end + + @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(nx, ny, nz) + cache, asm, kernel, m = _setup(mesh, Hexahedron{8}) + n = cache.ndofs + + x = ones(n) + y = zeros(n) + apply_K!(y, cache, asm, kernel, m, x) # warmup + + GC.gc() + a = @allocated apply_K!(y, cache, asm, kernel, m, x) + @test a == 0 + + nelems = length(m.connectivity) + println(" Hex8 $(nx)×$(ny)×$(nz) $(lpad(nelems,5)) elem " * + "$(lpad(n,5)) dof allocs=$a") + end + + @testset "Optimized LLVM IR has 0 gc-alloc sites" begin + mesh = _build_hex8_box(2, 1, 1) + cache, asm, kernel, m = _setup(mesh, Hexahedron{8}) + n = cache.ndofs + x = ones(n) + y = zeros(n) + apply_K!(y, cache, asm, kernel, m, x) # warmup + + iob = IOBuffer() + code_llvm(iob, apply_K!, + Tuple{typeof(y), typeof(cache), typeof(asm), + typeof(kernel), typeof(m), typeof(x)}; + 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(" apply_K! optimized LLVM IR: $n_alloc gc-alloc sites " * + "($(length(ir)) IR chars)") + end +end + +# ---------------------------------------------------------------------------- +# Krylov demo: matrix-free CG using `apply_K!` as a `LinearOperator`. +# This is the actual use case for the matrix-free path — once it solves, +# the API is proven end-to-end against an off-the-shelf solver. +# ---------------------------------------------------------------------------- +@testset "apply_K!: matrix-free CG via LinearOperators + IterativeSolvers" begin + println("\n" * "=" ^ 70) + println("DOF-BASED APPLY_K! — KRYLOV VALIDATION") + println("=" ^ 70) + + # Small unit cube, 3×3×3 = 27 elements, 192 DOFs. + nx = ny = nz = 3 + mesh = _build_hex8_box(nx, ny, nz) + cache, asm, kernel, m = _setup(mesh, Hexahedron{8}) + n = cache.ndofs + + # Bottom-face homogeneous Dirichlet via the shared `PenaltyDirichlet` + # abstraction — same struct that the heat domain uses. + fixed_dofs = Int[] + for (nid, X) in enumerate(m.nodes) + if X[3] == 0.0 + push!(fixed_dofs, 3*(nid - 1) + 1) + push!(fixed_dofs, 3*(nid - 1) + 2) + push!(fixed_dofs, 3*(nid - 1) + 3) + end + end + @test !isempty(fixed_dofs) + bc = PenaltyDirichlet(fixed_dofs; penalty = 1e16) + + # Loading: unit downward force on the corner (nx+1, ny+1, nz+1). + top_corner_node = (nx + 1) * (ny + 1) * (nz + 1) + fz_dof = 3 * (top_corner_node - 1) + 3 + b = zeros(n) + b[fz_dof] = -1e6 # 1 MN (problem is in SI; just a number) + + # 1. Direct solution via assembled K with the same penalty BC. + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + Kbc = Matrix(K) + apply_constraint!(Kbc, bc) + u_direct = Kbc \ b + + # 2. Matrix-free CG: `matrix_free_op` builds the closure with the + # Dirichlet contribution baked in. No K is materialized. + op = matrix_free_op(cache, asm, kernel, m; dirichlet = bc) + linop = LinearOperator(Float64, n, n, true, true, op) + + u_mf = zeros(n) + cg!(u_mf, linop, b; abstol = 1e-8, reltol = 1e-10, maxiter = 4 * n) + + # 3. Compare matrix-free to direct. + rel_u = norm(u_mf - u_direct) / max(norm(u_direct), 1.0) + @test rel_u < 1e-6 + + # Also verify residual of matrix-free solution under the same operator. + r = zeros(n) + op(r, u_mf) + rel_r = norm(b - r) / max(norm(b), 1.0) + @test rel_r < 1e-6 + + println(" Hex8 $(nx)×$(ny)×$(nz) ndof=$n fixed=$(length(fixed_dofs)) " * + "rel_u=$(round(rel_u; sigdigits=3)) rel_r=$(round(rel_r; sigdigits=3))") +end + +# ---------------------------------------------------------------------------- +# Type stability: apply_K! must infer `Vector{Float64}` (==typeof(y)) end-to-end +# ---------------------------------------------------------------------------- +@testset "apply_K!: type stability" begin + mesh = _build_hex8_box(2, 1, 1) + cache, asm, kernel, m = _setup(mesh, Hexahedron{8}) + n = cache.ndofs + x = ones(n) + y = zeros(n) + apply_K!(y, cache, asm, kernel, m, x) # warmup + + rt = Base.promote_op(apply_K!, typeof(y), typeof(cache), typeof(asm), + typeof(kernel), typeof(m), typeof(x)) + @test rt === Vector{Float64} + + code = code_typed(apply_K!, + (typeof(y), typeof(cache), typeof(asm), + typeof(kernel), typeof(m), typeof(x)); + optimize = true) + @test !isempty(code) + info = code[1] + @test isconcretetype(info.second) + println(" apply_K! inferred return type: $(info.second)") +end