From bfbf7abf09e3ff462754ddd4b2c65a89f21e1fb8 Mon Sep 17 00:00:00 2001 From: filoferra <66474175+filoferra@users.noreply.github.com> Date: Fri, 25 Sep 2026 19:39:44 +0200 Subject: [PATCH 1/2] local_to_local!, with_product_coordinates, set_matrix_preallocate_only! local_to_local!(dst, dm, src) and the in-place local_to_local!(v, dm) refresh ghost points. with_product_coordinates(f, dm) hands out a DMStag's per-axis coordinates read-only. set_matrix_preallocate_only! lets the first assembly define a DM matrix's pattern. The Read restore of the product coordinates called the writable restore. The KSP and SNES docstrings say when constructor options are applied. --- CHANGELOG.md | 5 ++ docs/src/man/dmstag.md | 12 ++- scripts/renames.jl | 3 + src/autowrapped/DM_wrappers.jl | 2 +- src/dm.jl | 48 ++++++++++++ src/dmstag.jl | 39 ++++++++++ src/ksp.jl | 6 ++ src/public_names.jl | 37 ++++----- src/snes.jl | 5 ++ src/ts.jl | 2 + test/dmstag_halos.jl | 77 +++++++++++++++++++ test/mpi_dmstag_halos.jl | 44 +++++++++++ test/runtests.jl | 3 +- ...MStagRestoreProductCoordinateArraysRead.jl | 2 +- 14 files changed, 263 insertions(+), 22 deletions(-) create mode 100644 test/dmstag_halos.jl create mode 100644 test/mpi_dmstag_halos.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index da36ed446..db84d2a72 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ - `LibPETSc.VecNestRestoreSubVecsRead` accepts the vector `VecNestGetSubVecsRead` returns; before, it only took a raw pointer that no wrapper handed out, so the read lock could not be released. - `LibPETSc.PCASMDestroySubdomains` and `PCGASMDestroySubdomains` accept the vectors the subdomain creators return. - On an Int32 PETSc build (the `PetscInt = "Int32"` preference), `LibPETSc.DMStagStencil`, `MatStencil` and the other C structs had `Int64` fields, so PETSc read them wrong. They now take the integer width of the loaded library, and the two stencil constructors convert their indices to it. +- `LibPETSc.DMStagRestoreProductCoordinateArraysRead` restored the arrays through the writable restore, which does not match the read-only get. +- The `KSP` and `SNES` docstrings did not say when the options given to the constructor are applied: once at construction for `KSP`, at every `solve!` for `SNES` (and `TS`). ### Added @@ -26,6 +28,9 @@ - `stencil(dm, loc, I; dof = 0)` builds a `DMStagStencil` from a 1-based element index, allocation free. - Assembly with DMStag stencils: `set_values!(J, dm, rows, cols, vals, mode)`, `set_values!(v, dm, positions, vals, mode)` and `zero_rows_local!(J, dm, rows, diag)`. They take any `AbstractVector`. - `LibPETSc.IS(dm, loc => dof, ...)`, the index set of whole DMStag fields, for `set_fieldsplit_is!`. +- `local_to_local!(dst, dm, src)` and the in-place `local_to_local!(v, dm)`, to refresh ghost points. +- `with_product_coordinates(f, dm)`: read-only access to a DMStag's per-axis coordinates, handed back when `f` returns. +- `set_matrix_preallocate_only!(dm, flag)`, so a matrix from `dm` gets its nonzero pattern from the first assembly. ### Changed diff --git a/docs/src/man/dmstag.md b/docs/src/man/dmstag.md index 2abb5b278..2774d7c98 100644 --- a/docs/src/man/dmstag.md +++ b/docs/src/man/dmstag.md @@ -82,6 +82,9 @@ lvec = PETSc.local_vec(dm) # The written vector comes first, the DM follows it (naming.md ยง8) PETSc.global_to_local!(lvec, dm, gvec, PETSc.INSERT_VALUES) PETSc.local_to_global!(gvec, dm, lvec, PETSc.ADD_VALUES) + +# Refresh the ghost points of a local vector from its neighbours, in place +PETSc.local_to_local!(lvec, dm) ``` ### Getting Location Indices @@ -150,6 +153,11 @@ PETSc.set_uniform_coordinates!(dm, xmin, xmax, ymin, ymax, zmin, zmax) # 3D # Get local coordinate array coords = PETSc.local_coordinate_array(dm) + +# Read the per-axis coordinates: x[i, 1] is the lower face of element i, x[i, 2] its centre +PETSc.with_product_coordinates(dm) do x, y + x[i, 2], y[j, 2] +end ``` ## Stencil Types @@ -176,9 +184,11 @@ PETSc.set_uniform_coordinates!(dm, 0.0, 1.0, 0.0, 1.0) # Create vectors and matrix x = PETSc.global_vec(dm) b = PETSc.global_vec(dm) -A = LibPETSc.DMCreateMatrix(petsclib, dm) +A = PETSc.PetscMat(dm) ``` +`PetscMat(dm)` stores every coupling the stencil allows, as explicit zeros. To let the first assembly define the pattern instead, call `PETSc.set_matrix_preallocate_only!(dm, true)` before creating the matrix. + ## Functions ```@autodocs diff --git a/scripts/renames.jl b/scripts/renames.jl index 0b9961ae9..e288cd92d 100644 --- a/scripts/renames.jl +++ b/scripts/renames.jl @@ -484,6 +484,9 @@ const UNCHANGED_PUBLIC = Symbol[ :face_location, # new in v0.5.2 :element_location, # new in v0.5.2 :stencil, # new in v0.5.2 + :local_to_local!, # new in v0.5.2 + :set_matrix_preallocate_only!, # new in v0.5.2 + :with_product_coordinates, # new in v0.5.2 :PetscDS, :AbstractPetscDS, :add_boundary!, diff --git a/src/autowrapped/DM_wrappers.jl b/src/autowrapped/DM_wrappers.jl index a7870f9ff..c253a4221 100644 --- a/src/autowrapped/DM_wrappers.jl +++ b/src/autowrapped/DM_wrappers.jl @@ -43681,7 +43681,7 @@ function DMStagRestoreProductCoordinateArraysRead(petsclib::PetscLibType, dm::Ab refZ = arrZ === nothing ? C_NULL : Ref(arrZ.ptr) @chk ccall( - (:DMStagRestoreProductCoordinateArrays, $petsc_library), + (:DMStagRestoreProductCoordinateArraysRead, $petsc_library), PetscErrorCode, (CDM, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}), dm, refX, refY, refZ, diff --git a/src/dm.jl b/src/dm.jl index a7f0d5ba9..2d432c30d 100644 --- a/src/dm.jl +++ b/src/dm.jl @@ -408,6 +408,35 @@ function global_to_local!( return lvec end +""" + local_to_local!(dst::AbstractPetscVec, dm::AbstractPetscDM, src::AbstractPetscVec, mode = INSERT_VALUES) + local_to_local!(v::AbstractPetscVec, dm::AbstractPetscDM, mode = INSERT_VALUES) + +Fill the ghost points of the local vector `dst` from the owned values of the local +vector `src` on the neighbouring ranks, and return `dst`. `dst` may be `src`: the +second form refreshes the ghost points of `v` in place. + +# External Links +$(doc_external("DM/DMLocalToLocalBegin")) +$(doc_external("DM/DMLocalToLocalEnd")) +""" +function local_to_local!( + dst::AbstractPetscVec{PetscLib}, + dm::AbstractPetscDM{PetscLib}, + src::AbstractPetscVec{PetscLib}, + mode::InsertMode = INSERT_VALUES, +) where {PetscLib} + LibPETSc.DMLocalToLocalBegin(getlib(PetscLib), dm, src, mode, dst) + LibPETSc.DMLocalToLocalEnd(getlib(PetscLib), dm, src, mode, dst) + return dst +end + +local_to_local!( + v::AbstractPetscVec{PetscLib}, + dm::AbstractPetscDM{PetscLib}, + mode::InsertMode = INSERT_VALUES, +) where {PetscLib} = local_to_local!(v, dm, v, mode) + """ set_uniform_coordinates!( @@ -594,3 +623,22 @@ function LibPETSc.PetscMat(da::AbstractPetscDM{PetscLib}) where {PetscLib} return J end +""" + set_matrix_preallocate_only!(dm::AbstractPetscDM, flag::Bool) + +With `flag` set, a matrix made from `dm` by [`PetscMat`](@ref) is preallocated but +its nonzero pattern is not inserted, so the first assembly defines the pattern +instead of every coupling the stencil allows being stored as an explicit zero. +The setting stays on `dm` for every later matrix. Returns `dm`. + +# External Links +$(doc_external("DM/DMSetMatrixPreallocateOnly")) +""" +function set_matrix_preallocate_only!( + dm::AbstractPetscDM{PetscLib}, + flag::Bool, +) where {PetscLib} + LibPETSc.DMSetMatrixPreallocateOnly(getlib(PetscLib), dm, LibPETSc.PetscBool(flag)) + return dm +end + diff --git a/src/dmstag.jl b/src/dmstag.jl index 360405a71..ef5abf81c 100644 --- a/src/dmstag.jl +++ b/src/dmstag.jl @@ -669,3 +669,42 @@ value_vector(::Type{PetscLib}, v::Vector) where {PetscLib} = eltype(v) === scalartype(PetscLib) ? v : Vector{scalartype(PetscLib)}(v) value_vector(::Type{PetscLib}, v::AbstractVector) where {PetscLib} = Vector{scalartype(PetscLib)}(v) + +# ============================================================================ +# Coordinates +# ============================================================================ + +""" + with_product_coordinates(f, dm::DMStag) + +Call `f(x)`, `f(x, y)` or `f(x, y, z)` with the local coordinate arrays of `dm`, +one per axis, and return what `f` returns. It needs product coordinates, which +[`set_uniform_coordinates!`](@ref) sets. The arrays are PETSc's own and are read +only: they are handed back when `f` returns or throws, and writing to them, or +using them after `f`, is undefined. + +Each array is indexed `[i, slot]`: `i` is the 1-based element index, ghost +elements included, as [`ghost_corners`](@ref) gives it, and `slot` is 1 for the +coordinate of the element's lower face and 2 for its centre. + +```julia +with_product_coordinates(dm) do x, y + x[i, 1], x[i, 2] # x of the lower face and of the centre of element i +end +``` + +# External Links +$(doc_external("DMStag/DMStagGetProductCoordinateArraysRead")) +$(doc_external("DMStag/DMStagRestoreProductCoordinateArraysRead")) +""" +function with_product_coordinates(f, dm::DMStag{PetscLib, N}) where {PetscLib, N} + lib = getlib(PetscLib) + x, y, z = LibPETSc.DMStagGetProductCoordinateArraysRead(lib, dm) + try + N == 1 && return f(x) + N == 2 && return f(x, y) + return f(x, y, z) + finally + LibPETSc.DMStagRestoreProductCoordinateArraysRead(lib, dm, x, y, z) + end +end diff --git a/src/ksp.jl b/src/ksp.jl index d3efa9e2e..51a9f7657 100644 --- a/src/ksp.jl +++ b/src/ksp.jl @@ -61,6 +61,9 @@ end Create a `KSP` using the matrix `A` and preconditioner construction matrix `P` with optional `prefix` and `options`. +The options are applied here, once, through `KSPSetFromOptions`. +A setting made in code afterwards, such as [`set_tolerances!`](@ref), overrides them. + The communicator is obtained from `A` and if it has size `1` then the garbage collector is set, otherwise the user is responsible for calling [`destroy!`](@ref). @@ -110,6 +113,9 @@ end Create a `KSP` associated with the `dm` with optional `prefix` and `options`. +The options are applied here, once, through `KSPSetFromOptions`. +A setting made in code afterwards overrides them. + The communicator is obtained from `dm`. The KSP can be used with geometric multigrid when the DM provides grid hierarchy information. diff --git a/src/public_names.jl b/src/public_names.jl index 212afa57d..523c9a3ad 100644 --- a/src/public_names.jl +++ b/src/public_names.jl @@ -20,24 +20,25 @@ public @bd_fn, @jacobian_fn, @residual_fn, @simple_fn, AbstractPetscDS, inttype, isassembled, isfinalized, isinitialized, issimplex, iteration_number, ksp, ksp_iterations, l2diff, label, library_info, local_arrays, local_coordinate_array, local_coordinates, local_indices, local_interior_linear_index, local_to_global!, - local_vec, mat_nullspace_create, max_steps, max_time, memtype, memtype_backend, - narrow, ndofs, options_prefix, ownership_range, owns, parse_option, parse_options, - pc, prev_time, project_field!, project_function!, release_local_array, reset!, - reshape_local_array, restore_local_arrays!, save_vtk, scalartype, set_adapt_type!, - set_auxiliary_vec!, set_compute_operators!, set_compute_rhs!, set_constants!, - set_convergence_test!, set_dm!, set_dm_active!, set_equation_type!, + local_to_local!, local_vec, mat_nullspace_create, max_steps, max_time, memtype, + memtype_backend, narrow, ndofs, options_prefix, ownership_range, owns, parse_option, + parse_options, pc, prev_time, project_field!, project_function!, release_local_array, + reset!, reshape_local_array, restore_local_arrays!, save_vtk, scalartype, + set_adapt_type!, set_auxiliary_vec!, set_compute_operators!, set_compute_rhs!, + set_constants!, set_convergence_test!, set_dm!, set_dm_active!, set_equation_type!, set_error_if_step_fails!, set_exact_final_time!, set_exact_solution!, set_field!, set_fieldsplit_is!, set_from_options!, set_function!, set_function_domain_error!, set_ifunction!, set_ijacobian!, set_jacobian!, set_jacobian_nullspace!, - set_jacobian_preconditioner!, set_library!, set_max_snes_failures!, - set_max_step_rejections!, set_max_steps!, set_max_time!, set_monitor!, set_name!, - set_nullspace!, set_operators!, set_option!, set_options_prefix!, set_petscint!, - set_post_step!, set_pre_step!, set_problem_type!, set_residual!, set_rhs_function!, - set_rhs_jacobian!, set_shell_apply!, set_shell_setup!, set_snes_jacobian!, - set_snes_local_fem!, set_solution!, set_step_number!, set_time!, set_timestep!, - set_tolerances!, set_type!, set_uniform_coordinates!, set_user_ctx!, set_values!, - setup!, snes, snes_failures, snes_iterations, solution, solve!, solve_time, - star_fd_coloring, stencil, step!, step_number, step_rejections, - tao_usable_after_reinitialize, timestep, tolerances, type_name, unsafe_local_array, - unset_library!, user_ctx, vertex_location, vtk_merge_tensor, with_local_array!, - wrap_local_array, zero_rows!, zero_rows_local! + set_jacobian_preconditioner!, set_library!, set_matrix_preallocate_only!, + set_max_snes_failures!, set_max_step_rejections!, set_max_steps!, set_max_time!, + set_monitor!, set_name!, set_nullspace!, set_operators!, set_option!, + set_options_prefix!, set_petscint!, set_post_step!, set_pre_step!, set_problem_type!, + set_residual!, set_rhs_function!, set_rhs_jacobian!, set_shell_apply!, + set_shell_setup!, set_snes_jacobian!, set_snes_local_fem!, set_solution!, + set_step_number!, set_time!, set_timestep!, set_tolerances!, set_type!, + set_uniform_coordinates!, set_user_ctx!, set_values!, setup!, snes, snes_failures, + snes_iterations, solution, solve!, solve_time, star_fd_coloring, stencil, step!, + step_number, step_rejections, tao_usable_after_reinitialize, timestep, tolerances, + type_name, unsafe_local_array, unset_library!, user_ctx, vertex_location, + vtk_merge_tensor, with_local_array!, with_product_coordinates, wrap_local_array, + zero_rows!, zero_rows_local! diff --git a/src/snes.jl b/src/snes.jl index 06fb767d2..bf08cac89 100644 --- a/src/snes.jl +++ b/src/snes.jl @@ -67,6 +67,11 @@ end Create a PETSc nonlinear solver (SNES) context on the communicator `comm`. +The options are stored and applied at the start of every [`solve!`](@ref), so that a +DM and callbacks attached after construction are visible to `SNESSetFromOptions`. +Because they are applied on every solve, they override a setting made in code for +the same option. + # Arguments - `petsclib`: The PETSc library instance - `comm::MPI.Comm`: MPI communicator diff --git a/src/ts.jl b/src/ts.jl index db39cc786..61ee9bdca 100644 --- a/src/ts.jl +++ b/src/ts.jl @@ -248,6 +248,8 @@ Create a PETSc time stepper on the communicator `comm`. Options are stored and applied in [`solve!`](@ref) rather than here, so that a DM and callbacks attached after construction are visible to `TSSetFromOptions`. +They are applied on every `solve!`, and override a setting made in code for the +same option. [`step!`](@ref) does not apply them. `exact_final_time` defaults to `TS_EXACTFINALTIME_MATCHSTEP`, so the last step lands on the time set by [`set_max_time!`](@ref). diff --git a/test/dmstag_halos.jl b/test/dmstag_halos.jl new file mode 100644 index 000000000..7f2d0a98d --- /dev/null +++ b/test/dmstag_halos.jl @@ -0,0 +1,77 @@ +# test/dmstag_halos.jl +# local_to_local! (in place and into another vector), the product coordinate +# do-block, and matrices preallocated without their nonzero pattern. + +using Test +using PETSc +using MPI + +MPI.Initialized() || MPI.Init() + +@testset "DMStag halos, coordinates and preallocation" begin + petsclib = PETSc.petsclibs[1] + PETSc.initialize(petsclib) + comm = MPI.COMM_SELF + LibPETSc = PETSc.LibPETSc + none(n) = ntuple(_ -> PETSc.DM_BOUNDARY_NONE, n) + + @testset "local_to_local!" begin + # 1D, periodic, one element dof: the local vector is [x4, x1, x2, x3, x4, x1] + dm = PETSc.DMStag(petsclib, comm, (PETSc.DM_BOUNDARY_PERIODIC,), (4,), (0, 1), 1) + g = PETSc.global_vec(dm) + g[1:4] = [1.0, 2.0, 3.0, 4.0] + PETSc.assemble!(g) + l = PETSc.local_vec(dm) + PETSc.global_to_local!(l, dm, g) + @test l[:] == [4.0, 1.0, 2.0, 3.0, 4.0, 1.0] + + # change the first owned element; its periodic ghost follows in place + l[2] = 10.0 + PETSc.assemble!(l) + @test PETSc.local_to_local!(l, dm) === l + @test l[6] == 10.0 + + dst = PETSc.local_vec(dm) + @test PETSc.local_to_local!(dst, dm, l, PETSc.INSERT_VALUES) === dst + @test dst[1] == 4.0 + @test dst[6] == 10.0 + foreach(PETSc.destroy!, (dst, l, g, dm)) + end + + @testset "with_product_coordinates" begin + dm = PETSc.DMStag(petsclib, comm, none(2), (3, 2), (1, 1, 1), 1) + PETSc.set_uniform_coordinates!(dm, (0.0, 0.0), (3.0, 2.0)) + + # slot 1 is the lower face, slot 2 the centre + got = PETSc.with_product_coordinates(dm) do x, y + (x[1, 1], x[1, 2], x[3, 2], y[2, 1], y[2, 2]) + end + @test got == (0.0, 0.5, 2.5, 1.0, 1.5) + + # the arrays go back when the block throws, so the coordinates can be set again + @test_throws ErrorException PETSc.with_product_coordinates((x, y) -> error("inside"), dm) + PETSc.set_uniform_coordinates!(dm, (0.0, 0.0), (6.0, 4.0)) + @test PETSc.with_product_coordinates((x, y) -> x[1, 2], dm) == 1.0 + + dm1 = PETSc.DMStag(petsclib, comm, none(1), (4,), (1, 1), 1) + PETSc.set_uniform_coordinates!(dm1, (0.0,), (1.0,)) + @test PETSc.with_product_coordinates(x -> x[4, 2], dm1) == 0.875 + dm3 = PETSc.DMStag(petsclib, comm, none(3), (2, 2, 2), (0, 0, 0, 1), 1) + PETSc.set_uniform_coordinates!(dm3, (0.0, 0.0, 0.0), (2.0, 2.0, 2.0)) + @test PETSc.with_product_coordinates((x, y, z) -> z[2, 1], dm3) == 1.0 + foreach(PETSc.destroy!, (dm3, dm1, dm)) + end + + @testset "set_matrix_preallocate_only!" begin + dm = PETSc.DMStag(petsclib, comm, none(2), (3, 2), (0, 1, 1), 1) + nonzeros(A) = LibPETSc.MatGetInfo(petsclib, A, LibPETSc.MAT_LOCAL).nz_used + A = PETSc.PetscMat(dm) + @test nonzeros(A) > 0 # the stencil pattern, stored as zeros + @test PETSc.set_matrix_preallocate_only!(dm, true) === dm + B = PETSc.PetscMat(dm) + @test nonzeros(B) == 0 + foreach(PETSc.destroy!, (B, A, dm)) + end + + PETSc.finalize(petsclib) +end diff --git a/test/mpi_dmstag_halos.jl b/test/mpi_dmstag_halos.jl new file mode 100644 index 000000000..af98efd69 --- /dev/null +++ b/test/mpi_dmstag_halos.jl @@ -0,0 +1,44 @@ +# test/mpi_dmstag_halos.jl +# Run on several ranks: local_to_local! in place fills each rank's ghost +# points from the owned values its neighbours changed in their local vectors. + +using Test +using MPI +MPI.Initialized() || MPI.Init() +using PETSc + +@testset "local_to_local! across ranks" begin + petsclib = PETSc.petsclibs[1] + PETSc.initialize(petsclib) + comm = MPI.COMM_WORLD + nranks = MPI.Comm_size(comm) + + # 1D, one element dof, two elements per rank + dm = PETSc.DMStag(petsclib, comm, (PETSc.DM_BOUNDARY_NONE,), (2nranks,), (0, 1), 1) + owned = PETSc.corners(dm) + ghosted = PETSc.ghost_corners(dm) + element(k) = ghosted.lower[1] + k - 1 # global element of local entry k + isowned(k) = owned.lower[1] <= element(k) <= owned.upper[1] + + g = PETSc.global_vec(dm) + PETSc.with_local_array!(g) do a + a .= owned.lower[1]:owned.upper[1] + end + l = PETSc.local_vec(dm) + PETSc.global_to_local!(l, dm, g) + + # negate the owned entries only, then refresh the ghosts in place + PETSc.with_local_array!(l) do a + for k in eachindex(a) + isowned(k) && (a[k] = -a[k]) + end + end + PETSc.local_to_local!(l, dm) + PETSc.with_local_array!(l) do a + # DMStag pads the last rank's ghost region past the grid; those entries stay unset + @test all(a[k] == -element(k) for k in eachindex(a) if element(k) <= 2nranks) + end + + foreach(PETSc.destroy!, (l, g, dm)) + PETSc.finalize(petsclib) +end diff --git a/test/runtests.jl b/test/runtests.jl index 97cda9c53..948547be2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,7 +19,7 @@ else end # Do the MPI tests first so we do not have mpi running inside MPI -mpi_tests = ("mpivec.jl", "mpimat.jl", "ksp.jl", "dmstag.jl") +mpi_tests = ("mpivec.jl", "mpimat.jl", "ksp.jl", "dmstag.jl", "mpi_dmstag_halos.jl") # PETSc_jll >= 3.25.4 has MPI-enabled Windows binaries (MicrosoftMPI), so the MPI tests run everywhere do_mpi = true @@ -47,6 +47,7 @@ include("dmproduct.jl") # test for DMProduct example include("matshell.jl") # autowrapped! include("test_dmstag.jl") include("dmstag_stencils.jl") # DMStag locations by axis, stencils, stencil assembly +include("dmstag_halos.jl") # local_to_local!, with_product_coordinates, preallocate only include("dm_dimension.jl") # dimension-correct DM returns: shape, inference, allocations include("test_snes.jl") include("test_audit.jl") # leak auditor diff --git a/wrapping/generator/overrides/DMStagRestoreProductCoordinateArraysRead.jl b/wrapping/generator/overrides/DMStagRestoreProductCoordinateArraysRead.jl index 317d70134..4e710c4c6 100644 --- a/wrapping/generator/overrides/DMStagRestoreProductCoordinateArraysRead.jl +++ b/wrapping/generator/overrides/DMStagRestoreProductCoordinateArraysRead.jl @@ -29,7 +29,7 @@ function DMStagRestoreProductCoordinateArraysRead(petsclib::PetscLibType, dm::Ab refZ = arrZ === nothing ? C_NULL : Ref(arrZ.ptr) @chk ccall( - (:DMStagRestoreProductCoordinateArrays, $petsc_library), + (:DMStagRestoreProductCoordinateArraysRead, $petsc_library), PetscErrorCode, (CDM, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}), dm, refX, refY, refZ, From f7f6c51a745c539fa486f17cd70ea035302fa2a7 Mon Sep 17 00:00:00 2001 From: filoferra <66474175+filoferra@users.noreply.github.com> Date: Fri, 25 Sep 2026 20:04:17 +0200 Subject: [PATCH 2/2] with_field_views!: typed views of DMStag local vectors by field with_field_views!(f, dm, vecs...; fields, read, write) checks out local vectors and hands f one view per location => dof pair, indexed like stencil, ghosts included, or the whole [I..., slot] array without fields. The views are concretely typed; a checkout costs 9 allocations per vector (independent of the grid size) --- CHANGELOG.md | 1 + docs/src/man/dmstag.md | 17 ++++++++ scripts/renames.jl | 4 ++ src/dmstag.jl | 90 ++++++++++++++++++++++++++++++++++++++ src/public_names.jl | 4 +- test/dmstag_views.jl | 99 ++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 7 files changed, 214 insertions(+), 2 deletions(-) create mode 100644 test/dmstag_views.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index db84d2a72..a54cbeb8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ - `local_to_local!(dst, dm, src)` and the in-place `local_to_local!(v, dm)`, to refresh ghost points. - `with_product_coordinates(f, dm)`: read-only access to a DMStag's per-axis coordinates, handed back when `f` returns. - `set_matrix_preallocate_only!(dm, flag)`, so a matrix from `dm` gets its nonzero pattern from the first assembly. +- `with_field_views!(f, dm, vecs...; fields, read, write)`: concretely typed views of DMStag local vectors by `location => dof`, indexed like `stencil`, handed back when `f` returns. About 9 allocations per vector, whatever the grid size. ### Changed diff --git a/docs/src/man/dmstag.md b/docs/src/man/dmstag.md index 2774d7c98..37e11d35b 100644 --- a/docs/src/man/dmstag.md +++ b/docs/src/man/dmstag.md @@ -99,6 +99,23 @@ indices = PETSc.local_indices(dm) indices = PETSc.global_indices(dm) ``` +### Views by field + +[`with_field_views!`](@ref) checks out local vectors and hands the block one view per field. Views are indexed by the element index, ghosts included, like `stencil`, and their types are concrete, so a loop over them compiles to plain array code: + +```julia +flow = (PETSc.face_location(dm, 1) => 0, PETSc.face_location(dm, 2) => 0, + PETSc.element_location(dm) => 0) +c = PETSc.corners(dm) +PETSc.with_field_views!(dm, x_local, r_local; fields = flow, write = (false, true)) do (Vx, Vy, P), (Rx, Ry, Rp) + for I in c.lower:c.upper + Rp[I] = Vx[I + CartesianIndex(1, 0)] - Vx[I] + Vy[I + CartesianIndex(0, 1)] - Vy[I] + end +end +``` + +Without `fields`, the block gets each whole array, indexed `[I..., slot]` with `slot` from `dof_slot`. `write = false` checks out read-only. + ### Locations and stencils DMStag names a point by where it sits on an element: `DMSTAG_LEFT`, `DMSTAG_DOWN`, `DMSTAG_BACK_DOWN_LEFT`, and so on. Those names shift meaning between dimensions (`DOWN` is the second axis, whatever it's called in your model), so the location functions count axes instead: diff --git a/scripts/renames.jl b/scripts/renames.jl index e288cd92d..ae0c7c39d 100644 --- a/scripts/renames.jl +++ b/scripts/renames.jl @@ -398,6 +398,9 @@ const INTERNAL = Set{Symbol}([ :check_axis, :stencil_vector, :value_vector, + :field_slots, + :call_with_views, + :field_views, :TSStepHookFn, :_unsafe_local_array, :_local_arrays, @@ -487,6 +490,7 @@ const UNCHANGED_PUBLIC = Symbol[ :local_to_local!, # new in v0.5.2 :set_matrix_preallocate_only!, # new in v0.5.2 :with_product_coordinates, # new in v0.5.2 + :with_field_views!, # new in v0.5.2 :PetscDS, :AbstractPetscDS, :add_boundary!, diff --git a/src/dmstag.jl b/src/dmstag.jl index ef5abf81c..b7b6cc014 100644 --- a/src/dmstag.jl +++ b/src/dmstag.jl @@ -708,3 +708,93 @@ function with_product_coordinates(f, dm::DMStag{PetscLib, N}) where {PetscLib, N LibPETSc.DMStagRestoreProductCoordinateArraysRead(lib, dm, x, y, z) end end + +# ============================================================================ +# Views of a local vector by field +# ============================================================================ + +""" + with_field_views!(f, dm::DMStag, vecs::AbstractPetscVec...; fields = nothing, read = true, write = true) + +Check out the local vectors `vecs` of `dm`, call `f` with one argument per vector, +hand the vectors back when `f` returns or throws, and return what `f` returns. + +`fields` is a tuple of `location => dof` pairs, as [`IS`](@ref LibPETSc.IS) takes +them, with the 0-based component of [`stencil`](@ref). For each vector `f` gets a +tuple of views, one per field and in that order. Without `fields`, `f` gets the +whole array per vector, indexed `[I..., slot]` with `slot` from [`dof_slot`](@ref). +Either way the element index `I` is 1-based and covers the ghost elements, as in +[`ghost_corners`](@ref), so a view and [`stencil`](@ref) address the same point. + +`read` and `write` work as in [`with_local_array!`](@ref): a `Bool`, or one per +vector. `write = false` checks out read-only; the views have the same types either +way. The vectors must be local vectors of `dm` ([`local_vec`](@ref)); anything else +throws an `ArgumentError` before any vector is checked out. + +```julia +flow = (face_location(dm, 1) => 0, face_location(dm, 2) => 0, element_location(dm) => 0) +c = corners(dm) +with_field_views!(dm, x, r; fields = flow, write = (false, true)) do (Vx, Vy, P), (Rx, Ry, Rp) + for I in c.lower:c.upper # the owned elements + Rp[I] = Vx[I + CartesianIndex(1, 0)] - Vx[I] + Vy[I + CartesianIndex(0, 1)] - Vy[I] + end +end +``` + +# External Links +$(doc_external("DMStag/DMStagVecGetArray")) +$(doc_external("DMStag/DMStagGetLocationSlot")) +""" +function with_field_views!( + f, + dm::DMStag{PetscLib, N}, + vecs::Vararg{AbstractPetscVec{PetscLib}, M}; + fields = nothing, + read::Union{Bool, NTuple{M, Bool}} = true, + write::Union{Bool, NTuple{M, Bool}} = true, +) where {PetscLib, N, M} + read isa Bool && (read = ntuple(_ -> read, Val(M))) + write isa Bool && (write = ntuple(_ -> write, Val(M))) + lib = getlib(PetscLib) + gc = ghost_corners(dm) + q = Int(LibPETSc.DMStagGetEntriesPerElement(lib, dm)) + dims = (q, gc.size...) + offsets = ntuple(d -> gc.lower[d] - 1, Val(N)) + slots = field_slots(dm, fields) + for v in vecs + LibPETSc.VecGetLocalSize(lib, v) == prod(dims) || throw( + ArgumentError("with_field_views! takes local vectors of the DMStag (local_vec(dm))"), + ) + end + acquired = map(vecs, read, write) do v, r, w + acquire_local_array(v; read = r, write = w) + end + try + return call_with_views(f, map(first, acquired), dims, offsets, slots) + finally + foreach(vecs, acquired, read, write) do v, (_, cpu_arr, backend), r, w + release_local_array(cpu_arr, backend, v; read = r, write = w) + end + end +end + +# The 1-based slot of each field, or `nothing` for the whole array +field_slots(dm, ::Nothing) = nothing +field_slots(dm::DMStag{PetscLib}, fields::Tuple) where {PetscLib} = map(fields) do (loc, dof) + dof_slot(dm, loc, Int(dof)) +end + +# A function barrier: the array types come from a runtime memory type, and +# everything below, `f` included, compiles for the concrete ones +function call_with_views(f, arrays, dims, offsets, slots) + return f(map(a -> field_views(reshape(a, dims), offsets, slots), arrays)...) +end + +# The array is stored with the component first; the views put it last, or drop it +function field_views(A::AbstractArray{T, D}, offsets::NTuple{N, Int}, ::Nothing) where {T, D, N} + whole = PermutedDimsArray(A, (ntuple(d -> d + 1, Val(N))..., 1)) + return OffsetArray(whole, offsets..., 0) +end + +field_views(A::AbstractArray{T, D}, offsets::NTuple{N, Int}, slots::Tuple) where {T, D, N} = + map(slot -> OffsetArray(view(A, slot, ntuple(_ -> Colon(), Val(N))...), offsets...), slots) diff --git a/src/public_names.jl b/src/public_names.jl index 523c9a3ad..0cdda16d4 100644 --- a/src/public_names.jl +++ b/src/public_names.jl @@ -40,5 +40,5 @@ public @bd_fn, @jacobian_fn, @residual_fn, @simple_fn, AbstractPetscDS, snes_iterations, solution, solve!, solve_time, star_fd_coloring, stencil, step!, step_number, step_rejections, tao_usable_after_reinitialize, timestep, tolerances, type_name, unsafe_local_array, unset_library!, user_ctx, vertex_location, - vtk_merge_tensor, with_local_array!, with_product_coordinates, wrap_local_array, - zero_rows!, zero_rows_local! + vtk_merge_tensor, with_field_views!, with_local_array!, with_product_coordinates, + wrap_local_array, zero_rows!, zero_rows_local! diff --git a/test/dmstag_views.jl b/test/dmstag_views.jl new file mode 100644 index 000000000..f8b84ef43 --- /dev/null +++ b/test/dmstag_views.jl @@ -0,0 +1,99 @@ +# test/dmstag_views.jl +# with_field_views!: views of DMStag local vectors by field, their types and +# allocations, and that every vector is handed back. + +using Test +using PETSc +using MPI + +MPI.Initialized() || MPI.Init() + +@testset "with_field_views!" begin + petsclib = PETSc.petsclibs[1] + PETSc.initialize(petsclib) + comm = MPI.COMM_SELF + LibPETSc = PETSc.LibPETSc + none(n) = ntuple(_ -> PETSc.DM_BOUNDARY_NONE, n) + stag(n, dofs) = PETSc.DMStag(petsclib, comm, none(length(n)), n, dofs, 1) + + @testset "views address the points stencil names" begin + dm = stag((3, 2), (1, 1, 1)) + flow = (PETSc.face_location(dm, 1) => 0, PETSc.face_location(dm, 2) => 0, + PETSc.element_location(dm) => 0) + x = PETSc.global_vec(dm) + I = CartesianIndex(2, 1) + PETSc.set_values!(x, dm, [PETSc.stencil(dm, flow[1].first, I), + PETSc.stencil(dm, flow[3].first, I)], [5.0, 7.0]) + PETSc.assemble!(x) + l = PETSc.local_vec(dm) + PETSc.global_to_local!(l, dm, x) + + got = PETSc.with_field_views!(dm, l; fields = flow, write = false) do (Vx, Vy, P) + (Vx[I], P[I], Vy[I], axes(P) == axes(Vx)) + end + @test got == (5.0, 7.0, 0.0, true) + + # the whole array, indexed [I..., slot] + slot = PETSc.dof_slot(dm, PETSc.element_location(dm), 0) + @test PETSc.with_field_views!(A -> A[I, slot], dm, l) == 7.0 + + # writing through a view changes the vector + PETSc.with_field_views!(dm, l; fields = flow[3:3]) do (P,) + P[I] = 9.0 + end + @test PETSc.with_field_views!(A -> A[I, slot], dm, l; write = false) == 9.0 + foreach(PETSc.destroy!, (l, x, dm)) + end + + @testset "types, allocations, several vectors" begin + function checkout(n) + dm = stag((n, n), (1, 1, 1)) + flow = (PETSc.face_location(dm, 1) => 0, PETSc.face_location(dm, 2) => 0, + PETSc.element_location(dm) => 0) + vs = ntuple(_ -> PETSc.local_vec(dm), 4) + body(a, b, c, d) = a[3][1, 1] + d[1][2, 2] + run() = PETSc.with_field_views!(body, dm, vs...; fields = flow, + write = (false, true, true, false)) + run() + allocations = @allocations run() + concrete = PETSc.with_field_views!(dm, vs...; fields = flow) do args... + all(t -> all(v -> isconcretetype(typeof(v)), t), args) + end + foreach(PETSc.destroy!, vs) + PETSc.destroy!(dm) + return allocations, concrete + end + small, concrete = checkout(8) + large, _ = checkout(64) + @test concrete + @test small == large # independent of the grid size + @test small <= 4 * 20 # the budget: about 20 per vector + + for (n, dofs) in (((4,), (1, 1)), ((2, 2, 2), (1, 1, 1, 1))) + dm = stag(n, dofs) + l = PETSc.local_vec(dm) + fields = (PETSc.vertex_location(dm) => 0, PETSc.element_location(dm) => 0) + @test PETSc.with_field_views!(dm, l; fields) do (V, E) + isconcretetype(typeof(V)) && ndims(E) == length(n) + end + foreach(PETSc.destroy!, (l, dm)) + end + end + + @testset "errors hand the vectors back" begin + dm = stag((3, 2), (0, 0, 1)) + l = PETSc.local_vec(dm) + g = PETSc.global_vec(dm) + @test_throws ErrorException PETSc.with_field_views!(A -> error("inside"), dm, l) + # a vector still checked out could not be checked out for writing again + @test PETSc.with_field_views!(A -> (A[1, 1, 1] = 1.0; true), dm, l) + + dm2 = stag((3, 3), (0, 0, 1)) + @test_throws ArgumentError PETSc.with_field_views!(A -> nothing, dm, l, g) + l2 = PETSc.local_vec(dm2) + @test_throws ArgumentError PETSc.with_field_views!(A -> nothing, dm, l2) + foreach(PETSc.destroy!, (l2, dm2, g, l, dm)) + end + + PETSc.finalize(petsclib) +end diff --git a/test/runtests.jl b/test/runtests.jl index 948547be2..fe17ac61f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,6 +48,7 @@ include("matshell.jl") # autowrapped! include("test_dmstag.jl") include("dmstag_stencils.jl") # DMStag locations by axis, stencils, stencil assembly include("dmstag_halos.jl") # local_to_local!, with_product_coordinates, preallocate only +include("dmstag_views.jl") # with_field_views!: types, allocations, hand-back include("dm_dimension.jl") # dimension-correct DM returns: shape, inference, allocations include("test_snes.jl") include("test_audit.jl") # leak auditor