Broadcasting into a PetscVec runs on the local arrays

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- Broadcasting into a `PetscVec` (`y .= 2 .* x`, a typical `MatShell` body) read and wrote the vectors one entry at a time, through `VecGetValues` and `VecSetValues`: at 10 000 entries a `MatShell` product took 139 000 allocations and was 500× slower than the same loop on the local arrays. The vectors now take part as their local arrays, at a fixed cost of about 10 allocations. On several ranks the broadcast runs over the entries each rank owns.
 - `LibPETSc.ISColoringGetIS` with `PETSC_OWN_POINTER` returned the index sets as borrowed, so they and the C array holding them leaked. The index sets are now owned by the caller and the array is freed. `ISColoringRestoreIS` accepts the vector `ISColoringGetIS` returns.
 - `LibPETSc.DMCreateFieldIS` leaked the field names and both C arrays.
 - The TS and SNES manual pages still described the 0.5.0 callback rules (return an error code, exceptions become a `PetscError`).

scripts/renames.jl

@@ -458,6 +458,7 @@ const INTERNAL = Set{Symbol}([
     :stencil_type_enum,
     :petsclib_of,
     :make_local_array,
+    :with_broadcast_arrays,
     :to_petscint_tuple,
     :audit_walk,
     :audit_targets,

src/vec.jl

@@ -157,6 +157,9 @@ Base.ndims(::Type{<:AbstractPetscVec}) = 1
 Base.IndexStyle(::Type{<:AbstractPetscVec}) = IndexLinear()
 Base.axes(v::AbstractPetscVec) = (Base.OneTo(length(v)),)
 Base.BroadcastStyle(::Type{<:AbstractPetscVec}) = Broadcast.DefaultArrayStyle{1}()
+# A Vec takes part in a broadcast as itself; the default would `collect` it
+# entry by entry, since it iterates but is not an AbstractArray
+Base.Broadcast.broadcastable(v::AbstractPetscVec) = v
 
 # The library a wrapper carries in its type parameter, as a value. Used where a
 # call has to recover `petsclib` from an object rather than take it as an
@@ -254,23 +257,40 @@ function Base.copyto!(
     return dst
 end
 
-# Broadcasting assignment support (dest is not an AbstractArray)
-function Base.copyto!(dest::AbstractPetscVec{PetscLib}, bc::Base.Broadcast.Broadcasted) where {PetscLib}
-    # Evaluate the broadcasted RHS
-    rhs = Base.materialize(bc)
-    if rhs isa Number
-        # Fast path for scalar RHS
-        Base.fill!(dest, rhs)
-        return dest
-    end
-    # Array-like RHS: shape must match
-    axes(dest) == axes(rhs) || throw(DimensionMismatch("broadcast axes $(axes(rhs)) do not match destination axes $(axes(dest))"))
-    @inbounds for i in eachindex(rhs)
-        dest[i] = rhs[i]
+# In-place broadcasting, `dest .= f.(args...)`, runs on the local arrays:
+# `dest` and every `PetscVec` among the arguments are checked out, the broadcast runs
+# over the arrays, and all are handed back. So it is elementwise over the entries
+# this rank owns, which on one rank is every entry.
+# Other array arguments must have the local length.
+function Base.copyto!(dest::AbstractPetscVec, bc::Base.Broadcast.Broadcasted)
+    flat = Base.Broadcast.flatten(bc)
+    arr, raw, backend = acquire_local_array(dest; read = true, write = true)
+    try
+        with_broadcast_arrays(dest, arr, flat.args) do args
+            copyto!(arr, Base.Broadcast.instantiate(Base.Broadcast.Broadcasted(flat.f, args)))
+        end
+    finally
+        release_local_array(raw, backend, dest; read = true, write = true)
     end
     return dest
 end
 
+# Calls `f` with `args`, each `PetscVec` replaced by its local array:
+# `dest_arr` for `dest` itself, a read-only checkout for any other
+with_broadcast_arrays(f, dest, dest_arr, ::Tuple{}) = f(())
+function with_broadcast_arrays(f, dest, dest_arr, args::Tuple)
+    a = first(args)
+    rest = Base.tail(args)
+    a isa AbstractPetscVec || return with_broadcast_arrays(r -> f((a, r...)), dest, dest_arr, rest)
+    a.ptr == dest.ptr && return with_broadcast_arrays(r -> f((dest_arr, r...)), dest, dest_arr, rest)
+    arr, raw, backend = acquire_local_array(a; read = true, write = false)
+    try
+        return with_broadcast_arrays(r -> f((arr, r...)), dest, dest_arr, rest)
+    finally
+        release_local_array(raw, backend, a; read = true, write = false)
+    end
+end
+
 Base.materialize!(dest::AbstractPetscVec, bc::Base.Broadcast.Broadcasted) = (Base.copyto!(dest, bc); dest)
 
 # Iterator interface

test/mpi/vec.jl

@@ -72,6 +72,15 @@ using LinearAlgebra: norm
                 @test x == julia_x
             end
 
+            # broadcasting works on the entries each rank owns
+            petsc_y = similar(petsc_x)
+            petsc_y .= 2 .* petsc_x .+ 1
+            PETSc.with_local_array!(petsc_y) do y
+                @test y == 2 .* julia_x .+ 1
+            end
+            @test norm(petsc_y) ≈ norm(2 .* (0:ne) .+ 1)
+            PETSc.destroy!(petsc_y)
+
             @test :mpi === PETSc.type_name(petsc_x)
             PETSc.destroy!(petsc_x)
 

test/vecmat/vec.jl

@@ -1,6 +1,6 @@
 using Test
 using PETSc, MPI
-using LinearAlgebra: norm
+using LinearAlgebra: norm, mul!
 
 MPI.Initialized() || MPI.Init()
 comm = LibPETSc.PETSC_COMM_SELF
@@ -361,4 +361,51 @@ end
         PETSc.destroy!(petsc_y)
         PETSc.finalize(petsclib)
     end
-end
+end
+@testset "broadcasting into a PetscVec" begin
+    petsclib = PETSc.petsclibs[1]
+    PETSc.initialize(petsclib)
+    PetscScalar = petsclib.PetscScalar
+    RealT = real(PetscScalar)
+
+    function vecs(n)
+        xv, yv = rand(RealT, n), rand(RealT, n)
+        x = PETSc.PetscVec(petsclib, comm, PetscScalar.(xv))
+        y = PETSc.PetscVec(petsclib, comm, PetscScalar.(yv))
+        return x, y, similar(x), xv, yv
+    end
+
+    n = 50
+    x, y, z, xv, yv = vecs(n)
+    z .= 2 .* x .+ sin.(y)
+    @test z[:] ≈ 2 .* xv .+ sin.(yv)
+    z .= 3
+    @test all(==(3), z[:])
+    z .= y .* (1:n)                   # a Julia array of the same length
+    @test z[:] ≈ yv .* (1:n)
+    x .= x .+ y                       # the destination read on the right
+    @test x[:] ≈ xv .+ yv
+    @test_throws DimensionMismatch z .= y .* ones(n + 1)
+    foreach(PETSc.destroy!, (x, y, z))
+
+    # the Vecs take part as arrays, so the cost does not grow with the length
+    function allocations(n)
+        x, y, z = vecs(n)
+        f!(z, x, y) = (z .= 2 .* x .+ y; nothing)
+        f!(z, x, y)
+        a = @allocations f!(z, x, y)
+        foreach(PETSc.destroy!, (x, y, z))
+        return a
+    end
+    small, large = allocations(100), allocations(10_000)
+    @test small == large
+    @test small <= 16
+
+    # a MatShell body written as a broadcast
+    x, y, z, _, yv = vecs(n)
+    shell = PETSc.MatShell(petsclib, (b, a) -> (b .= 2 .* a), comm, n, n)
+    mul!(z, shell, y)
+    @test z[:] ≈ 2 .* yv
+    foreach(PETSc.destroy!, (shell, x, y, z))
+    PETSc.finalize(petsclib)
+end