diff --git a/CHANGELOG.md b/CHANGELOG.md index d149a3f1e..65aa2124d 100644 --- a/CHANGELOG.md +++ b/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`). diff --git a/scripts/renames.jl b/scripts/renames.jl index 37f16f96e..16b537d2f 100644 --- a/scripts/renames.jl +++ b/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, diff --git a/src/vec.jl b/src/vec.jl index d9bcde0b8..7c98272fa 100644 --- a/src/vec.jl +++ b/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 diff --git a/test/mpi/vec.jl b/test/mpi/vec.jl index cb1121b57..cc90b0eb8 100644 --- a/test/mpi/vec.jl +++ b/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) diff --git a/test/vecmat/vec.jl b/test/vecmat/vec.jl index e6b4ee22d..7133da5d8 100644 --- a/test/vecmat/vec.jl +++ b/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 \ No newline at end of file +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