diff --git a/ext/JuliaFEMMPIExt.jl b/ext/JuliaFEMMPIExt.jl index f68aaf2..70e0ccb 100644 --- a/ext/JuliaFEMMPIExt.jl +++ b/ext/JuliaFEMMPIExt.jl @@ -1,5 +1,5 @@ -# This file is a part of JuliaFEM. -# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md +# SPDX-FileCopyrightText: 2015-2026 Jukka Aho +# SPDX-License-Identifier: MIT module JuliaFEMMPIExt @@ -99,17 +99,24 @@ function JuliaFEM.mpi_partitioned_operator_matvec_owned!( exchange::RankHaloExchange, cache::DOFBasedCOOCache, assembler::DOFBasedCOOAssembler, - kernel::AbstractKernel, mesh::AbstractMesh, comm::MPI.Comm; dirichlet = nothing, mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, ) copy_owned_subset_to_packed_owned_prefix!(packed, p_owned, layout) exchange_matvec_halos_mpi!( recv_vals, send_vals, packed, layout, exchange, comm; mpi_requests = mpi_requests) unpack_halo_recv_to_packed!(packed, recv_vals, exchange, layout) - apply_K_owned_rows_from_packed!(Ap_owned, packed, layout, cache, assembler, kernel, mesh) + apply_K_owned_rows_from_packed!( + Ap_owned, packed, layout, cache, assembler, mesh; + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) if dirichlet !== nothing dirichlet isa PenaltyDirichlet || throw(ArgumentError( @@ -120,6 +127,98 @@ function JuliaFEM.mpi_partitioned_operator_matvec_owned!( return Ap_owned end +function JuliaFEM.mpi_partitioned_operator_matvec_owned!( + Ap_owned::AbstractVector{Float64}, + p_owned::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + recv_vals::Vector{Vector{Float64}}, + send_vals::Vector{Vector{Float64}}, + layout::PartitionPackedLayout, + exchange::RankHaloExchange, + cache::DOFBasedCOOCache, + assembler::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + comm::MPI.Comm; + dirichlet = nothing, + mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, +) + JuliaFEM._depwarn_redundant_kernel_arg!(:mpi_partitioned_operator_matvec_owned!) + return mpi_partitioned_operator_matvec_owned!( + Ap_owned, p_owned, packed, recv_vals, send_vals, layout, exchange, cache, assembler, mesh, comm; + dirichlet = dirichlet, + mpi_requests = mpi_requests, + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + +# Owned-row internal force: pack `u`, halo exchange, then [`apply_f_int_owned_rows_from_packed!`](@ref). +# Requires `work` of length `cache.ndofs` for Pass~1 (expand + optional explicit `configuration`). +function JuliaFEM.mpi_partitioned_internal_force_owned!( + f_int_owned::AbstractVector{Float64}, + u_owned::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + work::AbstractVector{Float64}, + recv_vals::Vector{Vector{Float64}}, + send_vals::Vector{Vector{Float64}}, + layout::PartitionPackedLayout, + exchange::RankHaloExchange, + cache::DOFBasedCOOCache, + assembler::DOFBasedCOOAssembler, + mesh::AbstractMesh, + comm::MPI.Comm; + mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, +) + copy_owned_subset_to_packed_owned_prefix!(packed, u_owned, layout) + exchange_matvec_halos_mpi!( + recv_vals, send_vals, packed, layout, exchange, comm; mpi_requests = mpi_requests) + unpack_halo_recv_to_packed!(packed, recv_vals, exchange, layout) + apply_f_int_owned_rows_from_packed!( + f_int_owned, packed, work, layout, cache, assembler, mesh; + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) + return f_int_owned +end + +function JuliaFEM.mpi_partitioned_internal_force_owned!( + f_int_owned::AbstractVector{Float64}, + u_owned::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + work::AbstractVector{Float64}, + recv_vals::Vector{Vector{Float64}}, + send_vals::Vector{Vector{Float64}}, + layout::PartitionPackedLayout, + exchange::RankHaloExchange, + cache::DOFBasedCOOCache, + assembler::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + comm::MPI.Comm; + mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, +) + JuliaFEM._depwarn_redundant_kernel_arg!(:mpi_partitioned_internal_force_owned!) + return mpi_partitioned_internal_force_owned!( + f_int_owned, u_owned, packed, work, recv_vals, send_vals, layout, exchange, cache, assembler, mesh, comm; + mpi_requests = mpi_requests, + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + # Global replicated matvec: owned-row stiffness + `MPI.Allreduce!`, then optional penalty BC post-hook. function JuliaFEM.mpi_partitioned_operator_matvec!( Ap::AbstractVector{Float64}, @@ -132,11 +231,13 @@ function JuliaFEM.mpi_partitioned_operator_matvec!( exchange::RankHaloExchange, cache::DOFBasedCOOCache, assembler::DOFBasedCOOAssembler, - kernel::AbstractKernel, mesh::AbstractMesh, comm::MPI.Comm; dirichlet = nothing, mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, ) gather_owned_from_global_to_packed!(packed, p, layout) exchange_matvec_halos_mpi!( @@ -145,7 +246,12 @@ function JuliaFEM.mpi_partitioned_operator_matvec!( fill!(work, 0.0) expand_packed_to_global!(work, packed, layout) fill!(Ap, 0.0) - apply_K_owned_rows!(Ap, layout.owned_rows, cache, assembler, kernel, mesh, work) + apply_K_owned_rows!( + Ap, layout.owned_rows, cache, assembler, mesh, work; + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) MPI.Allreduce!(Ap, MPI.SUM, comm) if dirichlet !== nothing apply_constraint_post!(Ap, p, dirichlet) @@ -153,4 +259,35 @@ function JuliaFEM.mpi_partitioned_operator_matvec!( return Ap end +function JuliaFEM.mpi_partitioned_operator_matvec!( + Ap::AbstractVector{Float64}, + p::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + work::AbstractVector{Float64}, + recv_vals::Vector{Vector{Float64}}, + send_vals::Vector{Vector{Float64}}, + layout::PartitionPackedLayout, + exchange::RankHaloExchange, + cache::DOFBasedCOOCache, + assembler::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + comm::MPI.Comm; + dirichlet = nothing, + mpi_requests = nothing, + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Float64 = 0.0, +) + JuliaFEM._depwarn_redundant_kernel_arg!(:mpi_partitioned_operator_matvec!) + return mpi_partitioned_operator_matvec!( + Ap, p, packed, work, recv_vals, send_vals, layout, exchange, cache, assembler, mesh, comm; + dirichlet = dirichlet, + mpi_requests = mpi_requests, + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + end # module JuliaFEMMPIExt diff --git a/test/mpi/partitioned_internal_force_smoke.jl b/test/mpi/partitioned_internal_force_smoke.jl new file mode 100644 index 0000000..dfaf998 --- /dev/null +++ b/test/mpi/partitioned_internal_force_smoke.jl @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: 2015-2026 Jukka Aho +# SPDX-License-Identifier: MIT +# +# Multi-rank smoke: `mpi_partitioned_internal_force_owned!` on disjoint owned rows, +# scattered + `MPI.Allreduce(SUM)` vs serial [`assemble_internal_force!`](@ref). +# +# Run like `partitioned_matvec_smoke.jl` (see header there); CI runs this via mpiexec. + +using Test +using LinearAlgebra +using JuliaFEM +using MPI +using Random + +function main() + MPI.Init() + try + comm = MPI.COMM_WORLD + rank = MPI.Comm_rank(comm) + nprocs = MPI.Comm_size(comm) + @test nprocs >= 2 + + nparts = nprocs + part = rank + 1 + + nx, ny, nz = 3, 4, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel( + ContinuumFormulation{ThreeDimensional}(), + 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)) + @test maximum(layout.element_part_id) == nparts + + 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 + u = zeros(n) + if rank == 0 + Random.seed!(20260522) + randn!(u) + u .*= 1.0e-4 + end + MPI.Bcast!(u, 0, comm) + + L = build_partition_packed_layout_for_matvec( + handler, + layout, + mesh, + node_own, + elements, + part, + cache.dof_connectivity, + ) + ex = exch_all[part] + ws = partitioned_mpi_owned_matvec_workspace(L, ex) + work = zeros(n) + mpi_reqs = allocate_exchange_matvec_halo_mpi_requests(ex) + + no = L.n_owned + u_owned = zeros(no) + extract_owned_subset_from_global!(u_owned, u, L) + + f_owned = zeros(no) + mpi_partitioned_internal_force_owned!( + f_owned, + u_owned, + ws.packed, + work, + ws.recv_vals, + ws.send_vals, + L, + ex, + cache, + asm, + mesh, + comm; + mpi_requests = mpi_reqs, + ) + + f_scat = zeros(n) + @inbounds for k in 1:no + g = L.packed_to_global[k] + f_scat[g] = f_owned[k] + end + f_sum = similar(f_scat) + MPI.Allreduce!(f_scat, f_sum, MPI.SUM, comm) + + f_ref = zeros(n) + assemble_internal_force!(f_ref, cache, asm, mesh; configuration = u) + @test f_sum ≈ f_ref rtol = 1e-10 atol = 1e-10 + + rank == 0 && println( + "MPI partitioned internal force smoke: OK ($(nprocs) ranks)", + ) + finally + MPI.Finalize() + end + return nothing +end + +main()