From e4ce3d12eb103432a976d1ecdb8b5527e297ecc1 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 18:37:22 +0300 Subject: [PATCH] test(assemblers): add KernelAbstractions apply_K regression --- test/assemblers/test_dof_based_apply_K_ka.jl | 217 +++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 test/assemblers/test_dof_based_apply_K_ka.jl diff --git a/test/assemblers/test_dof_based_apply_K_ka.jl b/test/assemblers/test_dof_based_apply_K_ka.jl new file mode 100644 index 0000000..347035c --- /dev/null +++ b/test/assemblers/test_dof_based_apply_K_ka.jl @@ -0,0 +1,217 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Backend-agnostic `apply_K!` (KernelAbstractions.jl) regression. + +What we *can* test locally on darwin / aarch64: + +1. The KA cache mirror builds without errors from a CPU `DOFBasedCOOCache`. +2. `apply_K!(y, cache_ka, kernel, x)` on `KernelAbstractions.CPU()` matches + the direct CPU `apply_K!(y, cache, asm, kernel, mesh, x)` to round-off, + for both Tet4 and Hex8 meshes of various sizes (max rel error == 0.0 + in practice). +3. The result also matches `K * x` from the assembled matrix (transitivity + check: same path the original `apply_K!` test guarantees). +4. Re-running Pass 1 + re-syncing the KA cache produces identical output + on subsequent calls (Newton-Krylov rehearsal). + +What we cannot test here: + +* Actual CUDA / Metal / AMDGPU execution. Those require the + corresponding hardware. The KA `@kernel` is defined once and + dispatches on backend; once a GPU host runs this same test with + `Adapt.adapt(CUDABackend(), cache_ka)` it should pass identically. + +# Allocation note + +The KA path on `CPU()` backend allocates ≈700 bytes per `apply_K!` launch +for KA's host-side scheduler metadata (workgroup, ndrange, etc.). This +overhead is *not* present on real GPU backends and is not asserted here; +the canonical zero-allocation CPU path is the direct +`apply_K!(y, cache, asm, kernel, mesh, x)` (see `test_dof_based_apply_K.jl`). +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache, apply_K! +using JuliaFEM: DOFBasedCOOCacheKA, sync_from_cpu!, _prepare_caches! +using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex +using LinearAlgebra +using SparseArrays +using Tensors +using Random +using KernelAbstractions + +# ---------------------------------------------------------------------------- +# Mesh helpers (kept local so this file is independent). +# ---------------------------------------------------------------------------- + +function _ka_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 _ka_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 _ka_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}} + ElemType = Element{Topo, Lagrange{1}, S} + elements, dof_mgr = create_elements!(mesh, ElemType) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + return cache, asm, kernel, mesh +end + +# ---------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------- + +@testset "apply_K! (KernelAbstractions, CPU backend) correctness" begin + println("\n" * "=" ^ 70) + println("DOF-BASED APPLY_K! VIA KERNELABSTRACTIONS — CPU BACKEND") + println("=" ^ 70) + + Random.seed!(20260508) + + # ------------------------------------------------------------------ + # Single Tet4 + # ------------------------------------------------------------------ + @testset "Single Tet4" begin + mesh = _ka_build_single_tet4() + cache, asm, kernel, m = _ka_setup(mesh, Tetrahedron{4}) + + # Pass 1 + assemble K (CPU); also build the KA mirror. + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + n = size(K, 1) + cache_ka = DOFBasedCOOCacheKA(cache) + sync_from_cpu!(cache_ka, cache) + + max_rel_vs_ref = 0.0 + max_rel_vs_cpu = 0.0 + for trial in 1:8 + x = randn(n) + y_ref = K * x + + y_cpu = zeros(n) + apply_K!(y_cpu, cache, asm, kernel, m, x) + + y_ka = zeros(n) + apply_K!(y_ka, cache_ka, kernel, x) + + err_ref = norm(y_ka .- y_ref) / max(norm(y_ref), 1.0) + err_cpu = norm(y_ka .- y_cpu) / max(norm(y_cpu), 1.0) + max_rel_vs_ref = max(max_rel_vs_ref, err_ref) + max_rel_vs_cpu = max(max_rel_vs_cpu, err_cpu) + end + + @test max_rel_vs_ref < 1e-12 + @test max_rel_vs_cpu < 1e-14 + println(" Single Tet4 n=$n max rel vs K*x=$(round(max_rel_vs_ref, sigdigits=3)) max rel vs CPU apply_K!=$(round(max_rel_vs_cpu, sigdigits=3))") + end + + # ------------------------------------------------------------------ + # Hex8 cubes — sweep size to exercise multi-element row coverage + # ------------------------------------------------------------------ + @testset "Hex8 cubes" begin + for (nx, ny, nz) in ((1, 1, 1), (2, 1, 1), (4, 2, 2), (6, 3, 3)) + @testset "Hex8 $(nx)×$(ny)×$(nz)" begin + mesh = _ka_build_hex8_box(nx, ny, nz) + cache, asm, kernel, m = _ka_setup(mesh, Hexahedron{8}) + + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + n = size(K, 1) + cache_ka = DOFBasedCOOCacheKA(cache) + sync_from_cpu!(cache_ka, cache) + + max_rel_vs_ref = 0.0 + max_rel_vs_cpu = 0.0 + for trial in 1:5 + x = randn(n) + y_ref = K * x + + y_cpu = zeros(n) + apply_K!(y_cpu, cache, asm, kernel, m, x) + + y_ka = zeros(n) + apply_K!(y_ka, cache_ka, kernel, x) + + err_ref = norm(y_ka .- y_ref) / max(norm(y_ref), 1.0) + err_cpu = norm(y_ka .- y_cpu) / max(norm(y_cpu), 1.0) + max_rel_vs_ref = max(max_rel_vs_ref, err_ref) + max_rel_vs_cpu = max(max_rel_vs_cpu, err_cpu) + end + + @test max_rel_vs_ref < 1e-12 + @test max_rel_vs_cpu < 1e-14 + println(" Hex8 $(nx)×$(ny)×$(nz) $(nx*ny*nz) elem $(n) dof max rel vs K*x=$(round(max_rel_vs_ref, sigdigits=3)) max rel vs CPU apply_K!=$(round(max_rel_vs_cpu, sigdigits=3))") + end + end + end +end + + +@testset "apply_K! (KernelAbstractions, CPU backend) — re-prepare flow" begin + # Verify that re-running Pass 1 on the CPU cache and re-syncing to KA + # mirror produces the same result on a second call. This exercises + # the path a Newton-Krylov solver would take across iterations. + mesh = _ka_build_hex8_box(2, 1, 1) + cache, asm, kernel, m = _ka_setup(mesh, Hexahedron{8}) + + assemble!(cache, asm, kernel, m) + K, _ = extract_system(cache) + n = size(K, 1) + cache_ka = DOFBasedCOOCacheKA(cache) + + Random.seed!(11) + x = randn(n) + + # First call + sync_from_cpu!(cache_ka, cache) + y1 = zeros(n) + apply_K!(y1, cache_ka, kernel, x) + + # Re-run Pass 1 (no parameter change), re-sync, check identical y + JuliaFEM.reset!(cache) + _prepare_caches!(cache, kernel, m) + sync_from_cpu!(cache_ka, cache) + y2 = zeros(n) + apply_K!(y2, cache_ka, kernel, x) + + @test y1 ≈ y2 + @test maximum(abs, y1 .- y2) ≤ eps(Float64) * max(norm(y1), 1.0) +end