From 4edf8301b17f11aad91a3560b5b217e46d71c23e Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 16:30:36 +0300 Subject: [PATCH] feat(src): add partitioned_matvec.jl src/assemblers/partitioned_matvec.jl | 276 +++++++++++++++++++++++++++++++++++ 1 file changed, 276 insertions(+) --- src/assemblers/partitioned_matvec.jl | 276 +++++++++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 src/assemblers/partitioned_matvec.jl diff --git a/src/assemblers/partitioned_matvec.jl b/src/assemblers/partitioned_matvec.jl new file mode 100644 index 0000000..52bdb3c --- /dev/null +++ b/src/assemblers/partitioned_matvec.jl @@ -0,0 +1,276 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +# Single-partition orchestration for matrix-free `K*x` on owned rows using a +# packed patch + matvec halo exchange (MPI-ready). + +""" + allocate_halo_recv_buffers(exchange::RankHaloExchange) -> Vector{Vector{Float64}} + +Preallocate receive buffers aligned with `exchange.recv_dof` (setup only). + +After MPI non-blocking receives complete, or [`simulate_halo_recv_from_global!`](@ref) +in serial, pass the result +to [`unpack_halo_recv_to_packed!`](@ref). +""" +function allocate_halo_recv_buffers(exchange::RankHaloExchange)::Vector{Vector{Float64}} + return [zeros(length(rd)) for rd in exchange.recv_dof] +end + +""" + allocate_halo_send_buffers(exchange::RankHaloExchange) -> Vector{Vector{Float64}} + +Preallocate send buffers aligned with `exchange.send_dof` (setup only). + +Used by [`pack_halo_send_from_packed!`](@ref) before [`exchange_matvec_halos_mpi!`](@ref) +(with both `JuliaFEM` and `MPI` loaded). +""" +function allocate_halo_send_buffers(exchange::RankHaloExchange)::Vector{Vector{Float64}} + return [zeros(length(sd)) for sd in exchange.send_dof] +end + +""" + partitioned_matvec_workspace(layout::PartitionPackedLayout, exchange::RankHaloExchange) + +Named tuple `(packed, work, recv_vals)` with correct lengths for +[`partitioned_owned_matvec!`](@ref): + +- `packed` — length `layout.n_packed` +- `work` — length `layout.ndofs_global` (expanded stencil workspace) +- `recv_vals` — from [`allocate_halo_recv_buffers`](@ref)`(exchange)` + +`exchange.part` must match `layout.part`. +""" +function partitioned_matvec_workspace( + layout::PartitionPackedLayout, + exchange::RankHaloExchange, +) + exchange.part == layout.part || + throw(ArgumentError("exchange.part $(exchange.part) != layout.part $(layout.part)")) + packed = zeros(layout.n_packed) + work = zeros(layout.ndofs_global) + recv_vals = allocate_halo_recv_buffers(exchange) + return (; packed, work, recv_vals) +end + +""" + partitioned_mpi_owned_matvec_workspace(layout::PartitionPackedLayout, exchange::RankHaloExchange) + +Like [`partitioned_matvec_workspace`](@ref), but omits the `ndofs_global`-length `work` buffer used +by [`expand_packed_to_global!`](@ref). Intended for the lean MPI path +[`mpi_partitioned_operator_matvec_owned!`](@ref), which applies the stiffness rows directly from +[`packed`](@ref). + +Returns `(; packed, recv_vals, send_vals)` with `exchange.part == layout.part`. +""" +function partitioned_mpi_owned_matvec_workspace( + layout::PartitionPackedLayout, + exchange::RankHaloExchange, +) + exchange.part == layout.part || + throw(ArgumentError("exchange.part $(exchange.part) != layout.part $(layout.part)")) + packed = zeros(layout.n_packed) + recv_vals = allocate_halo_recv_buffers(exchange) + send_vals = allocate_halo_send_buffers(exchange) + return (; packed, recv_vals, send_vals) +end + +""" + simulate_halo_recv_from_global!(recv_vals, x_global, exchange) -> recv_vals + +Serial stand-in for completed MPI receives: + +`recv_vals[k][j] = x_global[exchange.recv_dof[k][j]]`. + +Buffers must be preallocated (`recv_vals[k]` length matches `recv_dof[k]`). +Allocation-free in the inner loops. +""" +function simulate_halo_recv_from_global!( + recv_vals::Vector{Vector{Float64}}, + x_global::AbstractVector{Float64}, + exchange::RankHaloExchange, +) + length(recv_vals) == length(exchange.recv_dof) == + length(exchange.recv_neighbor) || + throw(DimensionMismatch("recv_vals length mismatch")) + @inbounds for k in eachindex(exchange.recv_dof) + rd = exchange.recv_dof[k] + rv = recv_vals[k] + length(rv) == length(rd) || + throw(DimensionMismatch("recv_vals[$k] length $(length(rv)) != recv_dof length $(length(rd))")) + for j in eachindex(rd) + g = rd[j] + rv[j] = x_global[g] + end + end + return recv_vals +end + +""" + partitioned_owned_matvec!( + y_contrib, x_global, packed, work, recv_vals, + layout, exchange, cache, assembler, kernel, mesh; + fill_recv_from_global = true, + ) -> y_contrib + +One partition’s additive contribution to `K*x` on **owned rows** (zeros elsewhere +in `y_contrib`), using the matvec packed layout and [`build_matvec_halo_exchanges`](@ref) +metadata: + +1. [`gather_owned_from_global_to_packed!`](@ref)`(packed, x_global, layout)` +2. If `fill_recv_from_global` — [`simulate_halo_recv_from_global!`](@ref)`(recv_vals, x_global, exchange)` (skip when MPI fills `recv_vals`) +3. [`unpack_halo_recv_to_packed!`](@ref)`(packed, recv_vals, exchange, layout)` +4. [`fill!`](@ref)`(work, 0)` then [`expand_packed_to_global!`](@ref)`(work, packed, layout)` +5. [`apply_K_owned_rows!`](@ref)`(y_contrib, layout.owned_rows, ...)` + +Requires `exchange.part == layout.part`, `length(y_contrib) == cache.ndofs`, and +matching `work` / `packed` sizes from [`partitioned_matvec_workspace`](@ref). + +Summing `y_contrib` over disjoint owned partitions recovers global [`apply_K!`](@ref) +when `x_global` is the full vector (serial replica). + +Allocation-free after workspace setup when `fill_recv_from_global` follows the same pattern. +""" +function partitioned_owned_matvec!( + y_contrib::AbstractVector{Float64}, + x_global::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + work::AbstractVector{Float64}, + recv_vals::Vector{Vector{Float64}}, + layout::PartitionPackedLayout, + exchange::RankHaloExchange, + cache::DOFBasedCOOCache, + assembler::DOFBasedCOOAssembler, + kernel::AbstractKernel, + mesh::AbstractMesh; + fill_recv_from_global::Bool = true, +) + exchange.part == layout.part || + throw(ArgumentError("exchange.part $(exchange.part) != layout.part $(layout.part)")) + nd = cache.ndofs + length(y_contrib) == nd || + throw(DimensionMismatch("y_contrib length $(length(y_contrib)), ndofs $nd")) + length(x_global) == nd || + throw(DimensionMismatch("x_global length $(length(x_global)), ndofs $nd")) + length(work) == nd || + throw(DimensionMismatch("work length $(length(work)), ndofs $nd")) + layout.ndofs_global == nd || + throw(DimensionMismatch("layout.ndofs_global $(layout.ndofs_global), cache.ndofs $nd")) + + gather_owned_from_global_to_packed!(packed, x_global, layout) + if fill_recv_from_global + simulate_halo_recv_from_global!(recv_vals, x_global, exchange) + end + unpack_halo_recv_to_packed!(packed, recv_vals, exchange, layout) + fill!(work, 0.0) + expand_packed_to_global!(work, packed, layout) + apply_K_owned_rows!(y_contrib, layout.owned_rows, cache, assembler, kernel, mesh, work) + return y_contrib +end + +""" + apply_K_owned_rows_from_packed!( + Ap_owned, packed, layout, cache, assembler, kernel, mesh, + ) -> Ap_owned + +Like [`apply_K_owned_rows!`](@ref)`(y, layout.owned_rows, …, x)` with trial vector `x` implied by +`packed`: for each global column index `j`, use `packed[layout.global_to_packed[j]]`. Only owned +rows `packed_to_global[1:n_owned]` are computed; results go to `Ap_owned[k]` for row +`packed_to_global[k]`. + +Requires every stencil neighbor `j` of those rows to satisfy `layout.global_to_packed[j] ≠ 0` +(as ensured by [`build_partition_packed_layout_for_matvec`](@ref)). + +`length(Ap_owned) == layout.n_owned`, `length(packed) ≥ layout.n_packed`, `layout.ndofs_global == cache.ndofs`. +Allocation-free after warmup. +""" +function apply_K_owned_rows_from_packed!( + Ap_owned::AbstractVector{Float64}, + packed::AbstractVector{Float64}, + layout::PartitionPackedLayout, + cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType}, + assembler::DOFBasedCOOAssembler, + kernel::AbstractKernel, + mesh::AbstractMesh, +) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType} + ndofs = cache.ndofs + layout.ndofs_global == ndofs || + throw(DimensionMismatch( + "layout.ndofs_global $(layout.ndofs_global) != cache.ndofs $ndofs", + )) + no = layout.n_owned + n_packed = layout.n_packed + length(Ap_owned) == no || + throw(DimensionMismatch("Ap_owned length $(length(Ap_owned)), n_owned $no")) + length(packed) ≥ n_packed || + throw(DimensionMismatch("packed length $(length(packed)) < n_packed $n_packed")) + + owned_rows = layout.owned_rows + length(owned_rows) == ndofs || + throw(DimensionMismatch("owned_rows length $(length(owned_rows)); expected $ndofs")) + + g2p = layout.global_to_packed + p2g = layout.packed_to_global + + _prepare_caches!(cache, kernel, mesh) + + elements = cache.elements + element_caches = cache.element_caches + geometry_caches = cache.geometry_caches + qp_buffers = cache.qp_buffers + + dof_connectivity = cache.dof_connectivity + dof_to_elements = dof_connectivity.dof_to_elements + + loc_layout = local_dof_layout(E) + ndofs_elem = length(loc_layout) + + @inbounds for k in 1:no + dof_i = p2g[k] + owned_rows[dof_i] || + throw(ArgumentError("packed owned slot $k → global $dof_i not marked owned in layout")) + + yi = 0.0 + touching_elements = dof_to_elements[dof_i] + n_conns = length(touching_elements) + + @inbounds for conn_idx in 1:n_conns + conn = touching_elements[conn_idx] + elem_id_val = elem_id(conn) + local_i = local_dof_idx(conn) + + element = elements[elem_id_val]::E + element_cache = element_caches[elem_id_val] + geometry_cache = geometry_caches[elem_id_val] + qp_buffer = view(qp_buffers, :, elem_id_val) + + entry_i = loc_layout[local_i] + dofs_elem = element_cache.dofs + + @inbounds for local_j in 1:ndofs_elem + dof_j_global = Int(dofs_elem[local_j]) + entry_j = loc_layout[local_j] + + K_ij = evaluate_entry( + kernel, + geometry_cache, + qp_buffer, + entry_i, + entry_j, + Int(elem_id_val), + ) + + pk_j = g2p[dof_j_global] + pk_j == 0 && + throw(ArgumentError( + "global trial dof $dof_j_global not in packed patch (row $dof_i)", + )) + yi += K_ij * packed[pk_j] + end + end + + Ap_owned[k] = yi + end + + return Ap_owned +end