mirror of
https://github.com/JuliaFEM/JuliaFEM.jl.git
synced 2026-08-06 04:21:33 +00:00
test(assemblers): add partitioned matvec regression
This commit is contained in:
Jukka Aho
@@ -0,0 +1,230 @@
|
||||
# This file is a part of JuliaFEM.
|
||||
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
|
||||
|
||||
"""
|
||||
Orchestration API [`partitioned_owned_matvec!`](@ref): multi-part sum vs [`apply_K!`](@ref).
|
||||
"""
|
||||
|
||||
using Test
|
||||
using JuliaFEM
|
||||
using JuliaFEM: brick_hex_partition_slabs, validate_partition
|
||||
using JuliaFEM: node_partition_owner_min!
|
||||
using JuliaFEM: build_partition_packed_layout_for_matvec, build_matvec_halo_exchanges
|
||||
using JuliaFEM: partitioned_matvec_workspace, partitioned_mpi_owned_matvec_workspace, partitioned_owned_matvec!
|
||||
using JuliaFEM: simulate_halo_recv_from_global!
|
||||
using JuliaFEM: copy_owned_subset_to_packed_owned_prefix!, extract_owned_subset_from_global!, gather_owned_rows_from_workspace!
|
||||
using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache
|
||||
using JuliaFEM: apply_K!, apply_K_owned_rows!, apply_K_owned_rows_from_packed!
|
||||
using JuliaFEM: allocate_halo_recv_buffers, unpack_halo_recv_to_packed!, gather_owned_from_global_to_packed!, expand_packed_to_global!
|
||||
using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex
|
||||
using Random
|
||||
using Tensors
|
||||
|
||||
@testset "partitioned_owned_matvec!: nparts sum ≡ apply_K!" begin
|
||||
Random.seed!(20260518)
|
||||
nx, ny, nz = 3, 4, 2
|
||||
nparts = 3
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(),
|
||||
material, Displacement{3}())
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
asm = DOFBasedCOOAssembler()
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y)
|
||||
validate_partition(layout, length(elements))
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
|
||||
exch_all = build_matvec_halo_exchanges(
|
||||
handler, layout, mesh, node_own, elements, cache.dof_connectivity)
|
||||
|
||||
n = cache.ndofs
|
||||
x = randn(n)
|
||||
y_ref = zeros(n)
|
||||
apply_K!(y_ref, cache, asm, kernel, mesh, x)
|
||||
|
||||
y_sum = zeros(n)
|
||||
for p in 1:nparts
|
||||
L = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, p, cache.dof_connectivity)
|
||||
ex = exch_all[p]
|
||||
ws = partitioned_matvec_workspace(L, ex)
|
||||
y_p = zeros(n)
|
||||
partitioned_owned_matvec!(
|
||||
y_p, x, ws.packed, ws.work, ws.recv_vals,
|
||||
L, ex, cache, asm, kernel, mesh)
|
||||
y_sum .+= y_p
|
||||
end
|
||||
@test y_sum ≈ y_ref rtol = 1e-11 atol = 1e-11
|
||||
end
|
||||
|
||||
@testset "partitioned_owned_matvec! fill_recv_from_global=false" begin
|
||||
nx, ny, nz = 3, 4, 2
|
||||
nparts = 3
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(),
|
||||
material, Displacement{3}())
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
asm = DOFBasedCOOAssembler()
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y)
|
||||
validate_partition(layout, length(elements))
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
exch_all = build_matvec_halo_exchanges(
|
||||
handler, layout, mesh, node_own, elements, cache.dof_connectivity)
|
||||
|
||||
n = cache.ndofs
|
||||
x = randn(n)
|
||||
p = 2
|
||||
L = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, p, cache.dof_connectivity)
|
||||
ex = exch_all[p]
|
||||
ws = partitioned_matvec_workspace(L, ex)
|
||||
simulate_halo_recv_from_global!(ws.recv_vals, x, ex)
|
||||
y_a = zeros(n)
|
||||
partitioned_owned_matvec!(
|
||||
y_a, x, ws.packed, ws.work, ws.recv_vals,
|
||||
L, ex, cache, asm, kernel, mesh; fill_recv_from_global = false)
|
||||
y_b = zeros(n)
|
||||
partitioned_owned_matvec!(
|
||||
y_b, x, ws.packed, ws.work, ws.recv_vals,
|
||||
L, ex, cache, asm, kernel, mesh; fill_recv_from_global = true)
|
||||
@test y_a ≈ y_b
|
||||
end
|
||||
|
||||
@testset "partitioned_mpi_owned_matvec_workspace rejects part mismatch" begin
|
||||
nx, ny, nz = 2, 2, 2
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(),
|
||||
material, Displacement{3}())
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :x)
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
exch = build_matvec_halo_exchanges(
|
||||
handler, layout, mesh, node_own, elements, cache.dof_connectivity)
|
||||
L1 = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, 1, cache.dof_connectivity)
|
||||
@test_throws ArgumentError partitioned_mpi_owned_matvec_workspace(L1, exch[2])
|
||||
end
|
||||
|
||||
@testset "partitioned_matvec_workspace rejects part mismatch" begin
|
||||
nx, ny, nz = 2, 2, 2
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(),
|
||||
material, Displacement{3}())
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :x)
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
exch = build_matvec_halo_exchanges(
|
||||
handler, layout, mesh, node_own, elements, cache.dof_connectivity)
|
||||
L1 = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, 1, cache.dof_connectivity)
|
||||
@test_throws ArgumentError partitioned_matvec_workspace(L1, exch[2])
|
||||
end
|
||||
|
||||
@testset "owned subset ↔ packed owned-prefix helpers" begin
|
||||
nx, ny, nz = 3, 4, 2
|
||||
nparts = 3
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(
|
||||
ContinuumFormulation{FullThreeD}(),
|
||||
material,
|
||||
Displacement{3}(),
|
||||
)
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
asm = DOFBasedCOOAssembler()
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y)
|
||||
validate_partition(layout, length(elements))
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
p = 2
|
||||
L = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, p, cache.dof_connectivity)
|
||||
n = cache.ndofs
|
||||
x = randn(n)
|
||||
packed_a = zeros(L.n_packed)
|
||||
packed_b = zeros(L.n_packed)
|
||||
gather_owned_from_global_to_packed!(packed_a, x, L)
|
||||
v_owned = zeros(L.n_owned)
|
||||
extract_owned_subset_from_global!(v_owned, x, L)
|
||||
copy_owned_subset_to_packed_owned_prefix!(packed_b, v_owned, L)
|
||||
@test view(packed_a, 1:L.n_owned) ≈ view(packed_b, 1:L.n_owned)
|
||||
|
||||
y = zeros(n)
|
||||
apply_K!(y, cache, asm, kernel, mesh, x)
|
||||
Ap_own = zeros(L.n_owned)
|
||||
gather_owned_rows_from_workspace!(Ap_own, y, L)
|
||||
@inbounds for k in 1:L.n_owned
|
||||
g = L.packed_to_global[k]
|
||||
@test Ap_own[k] ≈ y[g]
|
||||
end
|
||||
end
|
||||
|
||||
@testset "apply_K_owned_rows_from_packed! ≡ apply_K_owned_rows! on expanded trial" begin
|
||||
Random.seed!(20260519)
|
||||
nx, ny, nz = 3, 4, 2
|
||||
nparts = 3
|
||||
mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz)
|
||||
material = LinearElastic(E = 210e9, ν = 0.3)
|
||||
kernel = ContinuumKernel(
|
||||
ContinuumFormulation{FullThreeD}(),
|
||||
material,
|
||||
Displacement{3}(),
|
||||
)
|
||||
S = @DOFSet{u::DOF{Displacement{3}, Vertex}}
|
||||
elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S})
|
||||
asm = DOFBasedCOOAssembler()
|
||||
cache = DOFBasedCOOCache(elements, handler, mesh, kernel)
|
||||
layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y)
|
||||
validate_partition(layout, length(elements))
|
||||
nnodes = length(mesh.nodes)
|
||||
node_own = Vector{Int}(undef, nnodes)
|
||||
node_partition_owner_min!(node_own, layout, mesh)
|
||||
exch_all = build_matvec_halo_exchanges(
|
||||
handler, layout, mesh, node_own, elements, cache.dof_connectivity)
|
||||
|
||||
n = cache.ndofs
|
||||
x_global = randn(n)
|
||||
p = 2
|
||||
L = build_partition_packed_layout_for_matvec(
|
||||
handler, layout, mesh, node_own, elements, p, cache.dof_connectivity)
|
||||
ex = exch_all[p]
|
||||
packed = zeros(L.n_packed)
|
||||
recv_vals = allocate_halo_recv_buffers(ex)
|
||||
gather_owned_from_global_to_packed!(packed, x_global, L)
|
||||
simulate_halo_recv_from_global!(recv_vals, x_global, ex)
|
||||
unpack_halo_recv_to_packed!(packed, recv_vals, ex, L)
|
||||
|
||||
work = zeros(n)
|
||||
fill!(work, 0.0)
|
||||
expand_packed_to_global!(work, packed, L)
|
||||
y_ws = zeros(n)
|
||||
apply_K_owned_rows!(y_ws, L.owned_rows, cache, asm, kernel, mesh, work)
|
||||
Ap_ref = zeros(L.n_owned)
|
||||
gather_owned_rows_from_workspace!(Ap_ref, y_ws, L)
|
||||
|
||||
Ap_packed = zeros(L.n_owned)
|
||||
apply_K_owned_rows_from_packed!(Ap_packed, packed, L, cache, asm, kernel, mesh)
|
||||
@test Ap_packed ≈ Ap_ref rtol = 1e-11 atol = 1e-11
|
||||
end
|
||||
Reference in New Issue
Block a user