test(assemblers): add KernelAbstractions apply_K regression

2026-09-22 18:52:16 +00:00 · 2026-05-09 18:37:22 +03:00
parent e6437ef902
commit e4ce3d12eb
1 changed files with 217 additions and 0 deletions
@@ -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