Checkout api

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## Unreleased
+
+### Deprecated
+
+- `local_arrays` and `restore_local_arrays!` warn: use `with_local_arrays!`. An exception between the two calls left both Vecs checked out.
+- `local_coordinate_array(dm)` warns: use `with_local_coordinate_array(f, dm)`. It now returns a copy. It used to leave the coordinate vector checked out until the garbage collector freed the array, and if the DM was destroyed first, that late hand-back crashed Julia.
+
+### Added
+
+- `with_local_arrays!(f!, g_fx, l_x)` calls `f!(fx, lx)` with the local arrays of the two Vecs, on the device when a GPU extension is loaded, and hands them back afterwards, also when `f!` throws.
+- `with_local_coordinate_array(f, dm)` calls `f(coord)` with the local coordinates of a `DMDA` or `DMStag`, addressed with global indices, and hands the array back when `f` returns or throws.
+
 ## v0.5.4
 
 ### Deprecated

docs/src/man/dmstag.md

@@ -173,8 +173,10 @@ PETSc.set_uniform_coordinates!(dm, xmin, xmax)                           # 1D
 PETSc.set_uniform_coordinates!(dm, xmin, xmax, ymin, ymax)               # 2D
 PETSc.set_uniform_coordinates!(dm, xmin, xmax, ymin, ymax, zmin, zmax)   # 3D
 
-# Get local coordinate array
-coords = PETSc.local_coordinate_array(dm)
+# Local coordinates as an array addressed with global indices, valid inside the block
+PETSc.with_local_coordinate_array(dm) do coords
+    coords[1, i, j]   # x of point (i, j)
+end
 
 # 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

docs/src/man/naming.md

@@ -62,7 +62,7 @@ set_type!(ksp, :gmres)
 ghost_corners(dm)
 ownership_range(A)
 local_indices(dm)
-local_coordinate_array(dm)
+with_local_coordinate_array(f, dm)
 project_function!(vec, dm, f)
 ```
 

examples/dmda/porosity_waves.jl

@@ -99,7 +99,7 @@ bcs = (
     PETSc.DM_BOUNDARY_GHOSTED,
    )[1:dim]
 
-# `corners`, `info`, `local_coordinate_array` and `reshape_local_array` are
+# `corners`, `info`, `with_local_coordinate_array` and `reshape_local_array` are
 # dimension-correct (docs/src/man/naming.md §12): on a `dim`-dimensional DMDA
 # they answer with `dim`-tuples, `CartesianIndex{dim}` and `1 + dim` axes.
 
@@ -131,7 +131,7 @@ PETSc.set_uniform_coordinates!(da, (-20, -10, -10), (150, 10, 10))
 
 # Cache local coordinates as a plain Julia array. This avoids repeatedly
 # mapping/restoring PETSc coordinate vectors during AD/Jacobian evaluations.
-coord_cache = copy(PETSc.local_coordinate_array(da))
+coord_cache = PETSc.with_local_coordinate_array(copy, da)
 
 # Create the PETSC snes object
 snes = PETSc.SNES(petsclib, comm; opts...)
@@ -431,7 +431,7 @@ while (it < max_it && time < max_time)
 
     # Visualisation
     if (mod(it, 40) == 0 && CreatePlots == true)
-        coord = PETSc.local_coordinate_array(da)
+        coord = PETSc.with_local_coordinate_array(copy, da)
         if dim==1
             x = collect(dof_slice(coord, 1))
             time_vec = ones(size(x)) * time

scripts/renames.jl

@@ -700,6 +700,46 @@ const EXTRA_SHIMS = NamedTuple{(:name, :code, :test), Tuple{String, String, Stri
             foreach(PETSc.destroy!, (y, x, A))
         end
         """,
+    ),    (
+        name = "local_arrays and restore_local_arrays!",
+        code = """
+        # The pair warns in src/vec.jl: `with_local_arrays!` replaces it and restores
+        # the arrays also when the kernel throws.
+        """,
+        test = """
+        @testset "local_arrays and restore_local_arrays!" begin
+            petsclib = PETSc.petsclibs[1]
+            PETSc.initialize(petsclib)
+            T = petsclib.PetscScalar
+            x = PETSc.PetscVec(petsclib, T[1, 2])
+            fx = PETSc.PetscVec(petsclib, 2)
+            arrays = warns(() -> PETSc.local_arrays(petsclib, fx, x))
+            arrays[1] .= 3 .* arrays[2]
+            @test warns(() -> PETSc.restore_local_arrays!(petsclib, fx, x, arrays...)) === nothing
+            @test fx[:] == T[3, 6]
+            foreach(PETSc.destroy!, (fx, x))
+        end
+        """,
+    ),    (
+        name = "local_coordinate_array",
+        code = """
+        # Warns in src/dm.jl and returns a copy; `with_local_coordinate_array` hands
+        # the array back when its block returns.
+        """,
+        test = """
+        @testset "local_coordinate_array" begin
+            petsclib = PETSc.petsclibs[1]
+            PETSc.initialize(petsclib)
+            B = PETSc.DM_BOUNDARY_NONE
+            da = PETSc.DMDA(petsclib, PETSc.MPI.COMM_SELF, (B,), (3,), 1)
+            PETSc.set_uniform_coordinates!(da, (0.0,), (1.0,))
+            coord = warns(() -> PETSc.local_coordinate_array(da))
+            @test coord[1, 3] ≈ 1
+            # a copy: nothing is left checked out of the coordinate vector
+            PETSc.destroy!(da)
+            @test coord[1, 3] ≈ 1
+        end
+        """,
     ),
 ]
 
@@ -708,6 +748,8 @@ const EXTRA_SHIMS = NamedTuple{(:name, :code, :test), Tuple{String, String, Stri
 # ---------------------------------------------------------------------------
 
 const INTERNAL = Set{Symbol}([
+    # the checkout behind with_local_arrays! and the deprecated local_arrays
+    :checkout_local_arrays,
     # the per-library body of the MatShell constructor
     :create_matshell,
     # leading underscore dropped (naming.md "Internals")
@@ -990,6 +1032,8 @@ const UNCHANGED_PUBLIC = Symbol[
     :PetscNotInitialized,
     :ConvergenceError,       # new in v0.5.3
     :with_options,           # new in v0.5.3
+    :with_local_arrays!,     # new in v0.5.5
+    :with_local_coordinate_array,  # new in v0.5.5
     :with_destroy,           # new in v0.5.3
     :track!,                 # new in v0.5.3
     :untrack!,               # new in v0.5.3

src/deprecations.jl

@@ -424,6 +424,14 @@ function MatShell(
     throw(ArgumentError("MatShell takes the operator first, then the library"))
 end
 
+# local_arrays and restore_local_arrays!
+# The pair warns in src/vec.jl: `with_local_arrays!` replaces it and restores
+# the arrays also when the kernel throws.
+
+# local_coordinate_array
+# Warns in src/dm.jl and returns a copy; `with_local_coordinate_array` hands
+# the array back when its block returns.
+
 
 # Renamed types: a `const` alias, so `isa` and dispatch keep working
 

src/dm.jl

@@ -546,29 +546,56 @@ $(doc_external("DM/DMGetCoordinatesLocal"))
 local_coordinates(dm::AbstractPetscDM{PetscLib}) where {PetscLib} =
     LibPETSc.DMGetCoordinatesLocal(getlib(PetscLib), dm)
 
+"""
+    with_local_coordinate_array(f, dm::Union{DMDA, DMStag})
+
+Call `f(coord)` with the coordinates of the local part of `dm`, ghost points
+included, and hand the array back to PETSc afterwards, also when `f` throws.
+Returns what `f` returns.
+
+`coord` is an `OffsetArray` addressed with global indices, the coordinate
+direction first: `coord[1, i, j]` is the x coordinate of point `(i, j)`. It
+reads PETSc's memory, so it is valid only inside `f`; return `copy(coord)` to
+keep the values.
+
+```julia
+coord = with_local_coordinate_array(copy, da)
+```
+
+# External Links
+$(doc_external("DM/DMGetCoordinatesLocal"))
+$(doc_external("Vec/VecGetArrayRead"))
+"""
+function with_local_coordinate_array(
+    f,
+    dm::Union{DMDA{PetscLib}, DMStag{PetscLib}},
+) where {PetscLib}
+    dim = LibPETSc.DMGetCoordinateDim(PetscLib, dm)
+    return with_local_array!(local_coordinates(dm); write = false) do a
+        f(reshape_local_array(a, dm, dim))
+    end
+end
+
 """
     local_coordinate_array(da::Union{DMDA, DMStag})
 
-Return coordinate arrays for the local portion of the domain.
+Return coordinate arrays for the local portion of the domain, as `OffsetArray`s
+addressed with global indices, ghost points included.
 
-The returned arrays are `OffsetArray`s that can be addressed using global indices,
-accounting for ghost points.
+Deprecated: use [`with_local_coordinate_array`](@ref). The arrays are a copy,
+so writing to them does not change the coordinates.
 
 # External Links
 $(doc_external("DM/DMGetCoordinatesLocal"))
 """
 function local_coordinate_array(
     da::Union{DMDA{PetscLib}, DMStag{PetscLib}},
 ) where {PetscLib}
-    # retrieve local coordinates
-    coord_vec = local_coordinates(da)
-    # array
-    array1D = unsafe_local_array(coord_vec; read = true, write = false)
-    dim = [PetscLib.PetscInt(0)]
-    dim = LibPETSc.DMGetCoordinateDim(PetscLib, da)
-    dim = dim[1]
-
-    return reshape_local_array(array1D, da, dim)
+    # TODO: remove in v0.6 with the other 0.5 deprecations.
+    @warn "local_coordinate_array is deprecated, use with_local_coordinate_array(f, dm)" maxlog = 1
+    # A copy: an array restored by a finalizer can run after `da` and its
+    # coordinate vector are destroyed, and restoring a freed Vec crashes.
+    return with_local_coordinate_array(copy, da)
 end
 
 

src/public_names.jl

@@ -44,5 +44,6 @@ public @bd_fn, @jacobian_fn, @residual_fn, @simple_fn, AO, AbstractAO, AbstractI
        solve_time, star_fd_coloring, stencil, step!, step_number, step_rejections,
        tao_usable_after_reinitialize, timestep, tolerances, track!, type_name,
        unsafe_local_array, unset_library!, untrack!, user_ctx, vertex_location,
-       vtk_merge_tensor, with_destroy, with_field_views!, with_local_array!, with_options,
+       vtk_merge_tensor, with_destroy, with_field_views!, with_local_array!,
+       with_local_arrays!, with_local_coordinate_array, with_options,
        with_product_coordinates, wrap_local_array, zero_rows!, zero_rows_local!

src/vec.jl

@@ -439,8 +439,11 @@ Return an `Array{PetscScalar}` containing local portion of the PETSc `vec`
 
 Use `read=false` if the array is write-only; `write=false` if read-only.
 
-!!! note
-    `Base.finalize` should be called on the `Array` before the data can be used.
+The caller hands the array back: `vec` stays checked out, and the next
+`VecGetArray` or assembly on it fails, until `Base.finalize` runs on the array
+(or the garbage collector gets to it). Prefer [`with_local_array!`](@ref), which
+hands it back when its block returns or throws; [`acquire_local_array`](@ref) and
+[`release_local_array`](@ref) are the explicit pair with no finalizer.
 
 # External Links
 $(doc_external("Vec/VecGetArray"))
@@ -891,38 +894,69 @@ function LinearAlgebra.norm(v::AbstractPetscVec, p::Real)
     throw(ArgumentError("PETSc computes the vector norm only for p = 1, 2 and Inf, got p = $p"))
 end
 
-# ── GPU-aware array access helpers ────────────────────────────────────────────
-#
-# `local_arrays` calls `VecGetArrayAndMemType` on both Vecs, converts the
-# returned `PetscMemType` values to backend singletons, and dispatches to
-# `_local_arrays`.  The base package handles the pure-CPU case
-# (host × host (both backends nothing)).  GPU extensions add `_local_arrays`
-# methods for their backend combinations and a matching
-# `_restore_local_arrays!` method dispatched by `restore_local_arrays!`.
-#
-# Return tuple:  (fx, lx, fx_arr, lx_arr, fx_bounce)
-#   CPU:  fx, lx are plain Arrays with VecRestore finalizers;
-#         fx_arr = lx_arr = fx_bounce = nothing
-#   GPU:  fx, lx are device arrays; fx_arr, lx_arr are raw PETSc arrays
-#         (needed for restore); fx_bounce is a scratch device array or nothing.
+# ── GPU-aware array access for a pair of Vecs ─────────────────────────────────
 
 """
-    local_arrays(petsclib, g_fx, l_x) -> (fx, lx, fx_arr, lx_arr, fx_bounce)
+    with_local_arrays!(f!, g_fx::AbstractPetscVec, l_x::AbstractPetscVec)
+
+Call `f!(fx, lx)` with the local array of `g_fx` (written) and of `l_x` (read),
+as arrays a compute kernel can use, and hand both back to PETSc afterwards,
+also when `f!` throws. Returns what `f!` returns.
 
-Return arrays for `g_fx` (read-write) and `l_x` (read-only) suitable for
-passing to a compute kernel.  Dispatches on the memory location of each Vec
-via `memtype_backend`.
+Where each array lives is decided for the pair: with a GPU backend extension
+loaded and both Vecs on the device, `fx` and `lx` wrap the device memory. When
+`g_fx` is on the host, `fx` is a device scratch array whose contents are copied
+back into `g_fx` on return, so the kernel can run on the device either way.
+Without a GPU, both are plain host arrays.
 
-On the pure-CPU path (`host × host (both backends nothing)`) `fx`/`lx` are plain
-`Array`s and `fx_arr = lx_arr = fx_bounce = nothing`.  When a GPU backend
-extension is loaded and a Vec lives on the device the returned `fx`/`lx` are
-device arrays.  An optional bounce buffer `fx_bounce` is allocated when `g_fx`
-is host-resident; its contents must be written back by `restore_local_arrays!`
-after the kernel completes.
+```julia
+with_local_arrays!(g_fx, l_x) do fx, lx
+    residual_kernel!(fx, lx)
+end
+```
 
-See also: [`restore_local_arrays!`](@ref)
+# External Links
+$(doc_external("Vec/VecGetArrayAndMemType"))
+"""
+function with_local_arrays!(
+    f!,
+    g_fx::AbstractPetscVec{PetscLib},
+    l_x::AbstractPetscVec,
+) where {PetscLib}
+    petsclib = getlib(PetscLib)
+    fx, lx, fx_arr, lx_arr, fx_bounce = checkout_local_arrays(petsclib, g_fx, l_x)
+    try
+        return f!(fx, lx)
+    finally
+        _restore_local_arrays!(petsclib, g_fx, l_x, fx, lx, fx_arr, lx_arr, fx_bounce)
+    end
+end
+
+"""
+    local_arrays(petsclib, g_fx, l_x) -> (fx, lx, fx_arr, lx_arr, fx_bounce)
+
+Check out the arrays of `g_fx` (read-write) and `l_x` (read-only) for a compute
+kernel; [`restore_local_arrays!`](@ref) must hand them back.
+
+Deprecated: use [`with_local_arrays!`](@ref), which restores the arrays also
+when the kernel throws.
 """
 function local_arrays(petsclib, g_fx, l_x)
+    # TODO: remove in v0.6 with the other 0.5 deprecations.
+    @warn "local_arrays is deprecated, use with_local_arrays!(f!, g_fx, l_x)" maxlog = 1
+    return checkout_local_arrays(petsclib, g_fx, l_x)
+end
+
+# Checks out both arrays and picks where they live: `VecGetArrayAndMemType` gives
+# each Vec's memory type, and `_local_arrays` dispatches on the pair of backends.
+# The base package handles host × host; GPU extensions add the other cases and a
+# matching `_restore_local_arrays!`.
+#
+# Returns (fx, lx, fx_arr, lx_arr, fx_bounce):
+#   CPU:  fx, lx are the PETSc arrays; fx_arr = lx_arr = fx_bounce = nothing
+#   GPU:  fx, lx are device arrays; fx_arr, lx_arr are the raw PETSc arrays
+#         (needed for restore); fx_bounce is a scratch device array or nothing.
+function checkout_local_arrays(petsclib, g_fx, l_x)
     T = petsclib.PetscScalar
     fx_arr, fx_mtype = LibPETSc.VecGetArrayAndMemType(petsclib, as_petsc_vec(g_fx))
     lx_arr, lx_mtype = LibPETSc.VecGetArrayReadAndMemType(petsclib, as_petsc_vec(l_x))
@@ -932,8 +966,8 @@ function local_arrays(petsclib, g_fx, l_x)
     )
 end
 
-# CPU base case: return arrays directly. restore_local_arrays! calls VecRestore
-# explicitly — no finalizers to avoid the double-finalization crash.
+# CPU base case: the PETSc arrays themselves, restored explicitly (no finalizers,
+# which could restore a second time and crash PETSc).
 function _local_arrays(
     petsclib, g_fx, l_x, ::Type, fx_arr, lx_arr, ::Nothing, ::Nothing,
 )
@@ -943,15 +977,15 @@ end
 """
     restore_local_arrays!(petsclib, g_fx, l_x, fx, lx, fx_arr, lx_arr, fx_bounce)
 
-Restore PETSc Vecs after a kernel launched via [`local_arrays`](@ref).
+Hand back the arrays [`local_arrays`](@ref) checked out.
 
-Dispatches to `_restore_local_arrays!`.  On the CPU path (`fx_arr`,
-`lx_arr`, `fx_bounce` all `nothing`) this simply finalizes `fx` and `lx`,
-triggering the registered `VecRestoreArray*AndMemType` finalizers.  GPU backend
-extensions add a `_restore_local_arrays!` method for their array types.
+Deprecated: use [`with_local_arrays!`](@ref).
 """
 function restore_local_arrays!(petsclib, g_fx, l_x, fx, lx, fx_arr, lx_arr, fx_bounce)
+    # TODO: remove in v0.6 with the other 0.5 deprecations.
+    @warn "restore_local_arrays! is deprecated, use with_local_arrays!(f!, g_fx, l_x)" maxlog = 1
     _restore_local_arrays!(petsclib, g_fx, l_x, fx, lx, fx_arr, lx_arr, fx_bounce)
+    return nothing
 end
 
 # CPU base case: call VecRestore directly (no finalizers).

test/core/deprecations.jl

@@ -354,4 +354,30 @@ const SHIMS = [
         foreach(PETSc.destroy!, (y, x, A))
     end
 
+    @testset "local_arrays and restore_local_arrays!" begin
+        petsclib = PETSc.petsclibs[1]
+        PETSc.initialize(petsclib)
+        T = petsclib.PetscScalar
+        x = PETSc.PetscVec(petsclib, T[1, 2])
+        fx = PETSc.PetscVec(petsclib, 2)
+        arrays = warns(() -> PETSc.local_arrays(petsclib, fx, x))
+        arrays[1] .= 3 .* arrays[2]
+        @test warns(() -> PETSc.restore_local_arrays!(petsclib, fx, x, arrays...)) === nothing
+        @test fx[:] == T[3, 6]
+        foreach(PETSc.destroy!, (fx, x))
+    end
+
+    @testset "local_coordinate_array" begin
+        petsclib = PETSc.petsclibs[1]
+        PETSc.initialize(petsclib)
+        B = PETSc.DM_BOUNDARY_NONE
+        da = PETSc.DMDA(petsclib, PETSc.MPI.COMM_SELF, (B,), (3,), 1)
+        PETSc.set_uniform_coordinates!(da, (0.0,), (1.0,))
+        coord = warns(() -> PETSc.local_coordinate_array(da))
+        @test coord[1, 3] ≈ 1
+        # a copy: nothing is left checked out of the coordinate vector
+        PETSc.destroy!(da)
+        @test coord[1, 3] ≈ 1
+    end
+
 end