From f2d5f44acc791ba95a4d7f575f14a0bc03ba1d71 Mon Sep 17 00:00:00 2001 From: filoferra <66474175+filoferra@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:59:50 +0100 Subject: [PATCH] Add copy, dot, axpy! and arithmetic for Vec and Mat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PetscVec gets copy, dot, axpy!, axpby!, rmul!, lmul! and the operators x + y, x - y, -x, α * x and x / α; an assembled PetscMat gets copy, A + B, A - B, -A and α * A. They call the PETSc routines (VecDot, VecAXPY, VecWAXPY, MatAXPY, MatScale, ...), so they are collective and work on distributed objects, and dot conjugates its first argument as it does for Julia arrays. The docs index now lists Base and LinearAlgebra methods under their own names. Purely additive --- CHANGELOG.md | 4 + docs/api_index.jl | 5 +- docs/src/man/mat.md | 2 + docs/src/man/naming.md | 6 +- docs/src/man/vec.md | 17 ++++ scripts/renames.jl | 5 ++ src/mat.jl | 59 +++++++++++++ src/vec.jl | 156 ++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 +- test/vecmat/linear_algebra.jl | 145 +++++++++++++++++++++++++++++++ 10 files changed, 397 insertions(+), 5 deletions(-) create mode 100644 test/vecmat/linear_algebra.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 71207cd9b..ad24891bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## Unreleased +### Added + +- `PetscVec` has `copy`, `dot`, `axpy!`, `axpby!`, `rmul!`, `lmul!` and the operators `x + y`, `x - y`, `-x`, `α * x` and `x / α`, through `VecCopy`, `VecDot`, `VecAXPY`, `VecAXPBY`, `VecScale` and `VecWAXPY`. An assembled `PetscMat` has `copy`, `A + B`, `A - B`, `-A` and `α * A`, through `MatDuplicate`, `MatAXPY` and `MatScale`. All are collective and work on distributed objects; `copy` and the operators return a new Vec or Mat, which on more than one rank needs `destroy!` like any other. `dot(x, y)` conjugates `x`, as for Julia arrays. + ### Deprecated - `local_arrays` and `restore_local_arrays!` warn: use `with_local_arrays!`. An exception between the two calls left both Vecs checked out. diff --git a/docs/api_index.jl b/docs/api_index.jl index e3eb47f84..b92ac1e68 100644 --- a/docs/api_index.jl +++ b/docs/api_index.jl @@ -29,8 +29,9 @@ const HIGH_LEVEL_SOURCES = [ const DOC_EXTERNAL_RE = r"_?doc_external\(\"([^\"]+)\"\)" # The name a definition line defines, whatever wraps it: `function f(`, -# `LibPETSc.@for_petsc function f(`, `macro f(`, `mutable struct F{`, `f(x) = …`. -const DEFINITION_RE = r"^\s*(?:LibPETSc\.@for_petsc\s+)?(?:@for_petsc\s+)?(?:function|macro|(?:mutable\s+)?struct|const|abstract\s+type)?\s*((?:Base\.|LibPETSc\.)?@?[A-Za-z_][A-Za-z0-9_!]*)" +# `LibPETSc.@for_petsc function f(`, `macro f(`, `mutable struct F{`, `f(x) = …`, +# and methods of Base and LinearAlgebra, operators included: `function Base.:+(`. +const DEFINITION_RE = r"^\s*(?:LibPETSc\.@for_petsc\s+)?(?:@for_petsc\s+)?(?:function|macro|(?:mutable\s+)?struct|const|abstract\s+type)?\s*((?:Base\.|LibPETSc\.|LinearAlgebra\.)?(?:@?[A-Za-z_][A-Za-z0-9_!]*|:\(?[-+*/\\=<>]+\)?))" """ collect_doc_external(srcdir) -> Dict{String, Vector{String}} diff --git a/docs/src/man/mat.md b/docs/src/man/mat.md index 5cf986422..ba4a77219 100644 --- a/docs/src/man/mat.md +++ b/docs/src/man/mat.md @@ -92,6 +92,8 @@ PETSc.setup!(A) # Complete matrix setup PETSc.destroy!(A) # Release it ``` +An assembled `PetscMat` supports `copy(A)`, `A + B`, `A - B`, `-A` and `α * A`. Each returns a new matrix and is collective; `A ± B` has the union of the two nonzero patterns. On more than one rank the result has no finalizer, so call `destroy!` on it on every rank. + `ownership_range(A)` is 1-based and takes no second argument. v0.4's `ownershiprange(A, false)` still works and warns; it is a `MethodError` in v0.6 ([naming conventions](naming.md), §12.1). diff --git a/docs/src/man/naming.md b/docs/src/man/naming.md index e47ab64ca..6c185c957 100644 --- a/docs/src/man/naming.md +++ b/docs/src/man/naming.md @@ -497,9 +497,11 @@ Add methods to `Base` and `LinearAlgebra` wherever the PETSc operation satisfies ```julia Base.size, Base.length, Base.ndims, Base.eltype, Base.axes -Base.similar, Base.fill!, Base.iterate, Base.copyto! +Base.similar, Base.copy, Base.fill!, Base.iterate, Base.copyto! Base.getindex, Base.setindex!, Base.show -LinearAlgebra.norm, LinearAlgebra.mul!, +Base.:+, Base.:-, Base.:*, Base.:/ +LinearAlgebra.norm, LinearAlgebra.mul!, LinearAlgebra.dot, +LinearAlgebra.axpy!, LinearAlgebra.axpby!, LinearAlgebra.rmul!, LinearAlgebra.lmul!, LinearAlgebra.issymmetric, LinearAlgebra.ishermitian ``` diff --git a/docs/src/man/vec.md b/docs/src/man/vec.md index 5c624bb71..75684abc2 100644 --- a/docs/src/man/vec.md +++ b/docs/src/man/vec.md @@ -50,6 +50,23 @@ length(v) # Get length size(v) # Get size tuple ``` +## Linear Algebra + +The vector operations of `LinearAlgebra` and the arithmetic operators call PETSc directly. All of them are collective, so every rank must make the call: + +```julia +using LinearAlgebra + +w = copy(v) # VecDuplicate + VecCopy +dot(x, y) # xᴴy, the global result on every rank +axpy!(α, x, y) # y = αx + y +axpby!(α, x, β, y) # y = αx + βy +rmul!(v, α) # v = αv +z = x + y # a new vector; so are x - y, -x, α * x and x / α +``` + +`copy` and the operators return a new `PetscVec` with the layout of their first vector argument. On more than one rank it has no finalizer, so call `destroy!` on it on every rank. Broadcasting is different for now: `y .= α .* x .+ y` updates `y` in place and works on any number of ranks, while an out-of-place broadcast such as `2 .* x` returns a Julia `Vector` and needs an undistributed vector. + ## Assembly After setting values, vectors must be assembled: diff --git a/scripts/renames.jl b/scripts/renames.jl index a7cba9363..52266aea5 100644 --- a/scripts/renames.jl +++ b/scripts/renames.jl @@ -923,6 +923,11 @@ const INTERNAL = Set{Symbol}([ :AUDIT_TYPE_CREATORS, :AUDIT_NAMED_CREATORS, :AUDIT_DESTROYER_NAMES, + :finalize_if_serial!, + :check_same_length, + :waxpy, + :scaled_copy, + :add_scaled, ]) # --------------------------------------------------------------------------- diff --git a/src/mat.jl b/src/mat.jl index deef0855b..ceff1fce4 100644 --- a/src/mat.jl +++ b/src/mat.jl @@ -677,6 +677,65 @@ function LinearAlgebra.ishermitian(A::PetscMat{PetscLib}; tol = 0.0) where {Pets return LibPETSc.MatIsHermitian(PetscLib, A, PetscReal(tol)) end +""" + copy(A::PetscMat) + +A new assembled `PetscMat` with the type, layout, nonzero pattern and entries of +`A`, which must be assembled. Collective. As from a constructor, the copy has a +finalizer only on a one-rank communicator; on more ranks call +[`destroy!`](@ref) on it on every rank. + +# External Links +$(doc_external("Mat/MatDuplicate")) +""" +function Base.copy(A::PetscMat{PetscLib}) where {PetscLib} + B = LibPETSc.MatDuplicate(getlib(PetscLib), A, LibPETSc.MAT_COPY_VALUES) + return finalize_if_serial!(B) +end + +""" + A + B + A - B + -A + α * A, A * α + +Matrix arithmetic on assembled `PetscMat`s, each giving a new `PetscMat` of the +type and layout of `A`. Collective. The nonzero pattern of `A ± B` is the union +of the two patterns. The result has a finalizer only on a one-rank communicator, +as [`copy`](@ref) does. + +# External Links +$(doc_external("Mat/MatAXPY")) +$(doc_external("Mat/MatScale")) +""" +function Base.:+(A::PetscMat{PetscLib}, B::PetscMat{PetscLib}) where {PetscLib} + return add_scaled(A, one(scalartype(PetscLib)), B) +end + +function Base.:-(A::PetscMat{PetscLib}, B::PetscMat{PetscLib}) where {PetscLib} + return add_scaled(A, -one(scalartype(PetscLib)), B) +end + +Base.:-(A::PetscMat) = scaled_copy(A, -1) +Base.:*(α::Number, A::PetscMat) = scaled_copy(A, α) +Base.:*(A::PetscMat, α::Number) = scaled_copy(A, α) + +function scaled_copy(A::PetscMat{PetscLib}, α::Number) where {PetscLib} + C = copy(A) + LibPETSc.MatScale(PetscLib, C, scalartype(PetscLib)(α)) + return C +end + +# A copy of A plus αB. UNKNOWN_NONZERO_PATTERN lets PETSc compare the two +# patterns and merge them when they differ. +function add_scaled(A::PetscMat{PetscLib}, α, B::PetscMat{PetscLib}) where {PetscLib} + size(A) == size(B) || + throw(DimensionMismatch("matrices have sizes $(size(A)) and $(size(B))")) + C = copy(A) + LibPETSc.MatAXPY(PetscLib, C, α, B, LibPETSc.UNKNOWN_NONZERO_PATTERN) + return C +end + """ setup!(mat::AbstractMat) diff --git a/src/vec.jl b/src/vec.jl index 66e92f5f9..e84632aa2 100644 --- a/src/vec.jl +++ b/src/vec.jl @@ -894,6 +894,162 @@ 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 +# ============================================================================= +# Linear algebra on Vecs +# ============================================================================= + +# A new Vec or Mat gets a finalizer only on one rank, as from its constructor: +# its Destroy is collective, and a GC finalizer runs at an arbitrary point +function finalize_if_serial!( + obj::Union{AbstractPetscVec{PetscLib}, LibPETSc.AbstractPetscMat{PetscLib}}, +) where {PetscLib} + if MPI.Comm_size(LibPETSc.PetscObjectGetComm(getlib(PetscLib), obj)) == 1 + finalizer(destroy!, obj) + end + return obj +end + +function check_same_length(x::AbstractPetscVec, y::AbstractPetscVec) + length(x) == length(y) || + throw(DimensionMismatch("vectors have lengths $(length(x)) and $(length(y))")) + return nothing +end + +""" + copy(v::AbstractPetscVec) + +A new `PetscVec` with the layout and entries of `v`. Collective. As from a +constructor, the copy has a finalizer only on a one-rank communicator; on more +ranks call [`destroy!`](@ref) on it on every rank. + +# External Links +$(doc_external("Vec/VecDuplicate")) +$(doc_external("Vec/VecCopy")) +""" +function Base.copy(v::AbstractPetscVec{PetscLib}) where {PetscLib} + w = LibPETSc.VecDuplicate(getlib(PetscLib), v) + LibPETSc.VecCopy(PetscLib, v, w) + return finalize_if_serial!(w) +end + +""" + dot(x::AbstractPetscVec, y::AbstractPetscVec) + +`xᴴy`, conjugating `x` as `dot` does for Julia arrays. Collective: every rank +gets the global result. + +# External Links +$(doc_external("Vec/VecDot")) +""" +function LinearAlgebra.dot( + x::AbstractPetscVec{PetscLib}, + y::AbstractPetscVec{PetscLib}, +) where {PetscLib} + check_same_length(x, y) + # VecDot(x, y) is yᴴx + return LibPETSc.VecDot(PetscLib, y, x) +end + +""" + axpy!(α, x::AbstractPetscVec, y::AbstractPetscVec) + +`y = αx + y`, returning `y`. Collective. + +# External Links +$(doc_external("Vec/VecAXPY")) +""" +function LinearAlgebra.axpy!( + α::Number, + x::AbstractPetscVec{PetscLib}, + y::AbstractPetscVec{PetscLib}, +) where {PetscLib} + check_same_length(x, y) + LibPETSc.VecAXPY(PetscLib, y, scalartype(PetscLib)(α), x) + return y +end + +""" + axpby!(α, x::AbstractPetscVec, β, y::AbstractPetscVec) + +`y = αx + βy`, returning `y`. Collective. + +# External Links +$(doc_external("Vec/VecAXPBY")) +""" +function LinearAlgebra.axpby!( + α::Number, + x::AbstractPetscVec{PetscLib}, + β::Number, + y::AbstractPetscVec{PetscLib}, +) where {PetscLib} + check_same_length(x, y) + PetscScalar = scalartype(PetscLib) + LibPETSc.VecAXPBY(PetscLib, y, PetscScalar(α), PetscScalar(β), x) + return y +end + +""" + rmul!(v::AbstractPetscVec, α) + lmul!(α, v::AbstractPetscVec) + +Scale `v` by `α` in place and return it. Collective. + +# External Links +$(doc_external("Vec/VecScale")) +""" +function LinearAlgebra.rmul!(v::AbstractPetscVec{PetscLib}, α::Number) where {PetscLib} + LibPETSc.VecScale(PetscLib, v, scalartype(PetscLib)(α)) + return v +end +LinearAlgebra.lmul!(α::Number, v::AbstractPetscVec) = rmul!(v, α) + +""" + x + y + x - y + -x + α * x, x * α + x / α + +Vector arithmetic on `PetscVec`s, each giving a new `PetscVec` with the layout of +`x`. Collective. The result has a finalizer only on a one-rank communicator, as +[`copy`](@ref) does. For an update in place use [`axpy!`](@ref), [`axpby!`](@ref) +and [`rmul!`](@ref), or broadcast into an existing vector, `y .= α .* x .+ y`. + +# External Links +$(doc_external("Vec/VecWAXPY")) +$(doc_external("Vec/VecScale")) +""" +function Base.:+(x::AbstractPetscVec{PetscLib}, y::AbstractPetscVec{PetscLib}) where {PetscLib} + return waxpy(one(scalartype(PetscLib)), x, y) +end + +function Base.:-(x::AbstractPetscVec{PetscLib}, y::AbstractPetscVec{PetscLib}) where {PetscLib} + return waxpy(-one(scalartype(PetscLib)), y, x) +end + +Base.:-(x::AbstractPetscVec) = rmul!(copy(x), -1) +Base.:*(α::Number, x::AbstractPetscVec) = rmul!(copy(x), α) +Base.:*(x::AbstractPetscVec, α::Number) = rmul!(copy(x), α) +# Through the broadcast rather than VecScale with 1/α, so each entry is divided +# exactly as `x ./ α` would +function Base.:/(x::AbstractPetscVec{PetscLib}, α::Number) where {PetscLib} + w = finalize_if_serial!(similar(x)) + w .= x ./ scalartype(PetscLib)(α) + return w +end + +# A new vector αx + y with the layout of x +function waxpy( + α, + x::AbstractPetscVec{PetscLib}, + y::AbstractPetscVec{PetscLib}, +) where {PetscLib} + check_same_length(x, y) + w = finalize_if_serial!(similar(x)) + LibPETSc.VecWAXPY(PetscLib, w, α, x, y) + return w +end + # ── GPU-aware array access for a pair of Vecs ───────────────────────────────── """ diff --git a/test/runtests.jl b/test/runtests.jl index 1c79503e7..86049210e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -66,6 +66,7 @@ if selected("vecmat") include("vecmat/mat.jl") include("vecmat/mat_vec_methods.jl") # copyto!, fill!, zero_rows!, set_option!, diagonal!, isassembled include("vecmat/matshell.jl") + include("vecmat/linear_algebra.jl") # copy, dot, axpy!, +, -, * on Vecs and Mats end if selected("solvers") @@ -130,7 +131,7 @@ end # Last, so that no MPI run happens inside another if selected("mpi") @testset "MPI Tests" begin - for file in ("mpi/vec.jl", "mpi/mat.jl", "solvers/ksp.jl", "dm/dmstag.jl", + for file in ("mpi/vec.jl", "mpi/mat.jl", "vecmat/linear_algebra.jl", "solvers/ksp.jl", "dm/dmstag.jl", "mpi/dmstag_halos.jl", "mpi/blas_threads.jl", "core/with_destroy.jl") @test run_separately(file; nranks = 4) diff --git a/test/vecmat/linear_algebra.jl b/test/vecmat/linear_algebra.jl new file mode 100644 index 000000000..8a0885eb9 --- /dev/null +++ b/test/vecmat/linear_algebra.jl @@ -0,0 +1,145 @@ +using Test +using MPI +MPI.Initialized() || MPI.Init() +using PETSc +using LinearAlgebra: dot, axpy!, axpby!, rmul!, lmul! + +# Vec and Mat arithmetic on COMM_WORLD, so on several ranks it checks that the +# operations are collective and give the global result. Each rank compares the +# entries it owns against the same computation on Julia arrays. + +# The entries of `v` this rank owns, and the matching slice of a Julia vector +owned(v::PETSc.AbstractPetscVec) = v[PETSc.ownership_range(v)] +owned(v::PETSc.AbstractPetscVec, ref::AbstractVector) = ref[PETSc.ownership_range(v)] + +function petsc_vec(petsclib, comm, vals) + v = PETSc.PetscVec(petsclib, comm, length(vals)) + r = PETSc.ownership_range(v) + v[r] = vals[r] + return PETSc.assemble!(v) +end + +# Rows this rank owns of a distributed AIJ matrix filled from `entries`, a list of +# (i, j, value) with 1-based indices +function petsc_mat(petsclib, comm, n, entries) + PetscInt = petsclib.PetscInt + A = LibPETSc.MatCreateAIJ(petsclib, comm, PetscInt(LibPETSc.PETSC_DECIDE), + PetscInt(LibPETSc.PETSC_DECIDE), PetscInt(n), PetscInt(n), + PetscInt(3), C_NULL, PetscInt(3), C_NULL) + rows = PETSc.ownership_range(A) + for (i, j, a) in entries + i in rows && (A[i, j] = a) + end + return PETSc.assemble!(A) +end + +@testset "Vec and Mat linear algebra" begin + comm = MPI.COMM_WORLD + for petsclib in PETSc.petsclibs + PETSc.initialize(petsclib) + PetscScalar = petsclib.PetscScalar + n = 4 * MPI.Comm_size(comm) + 3 + # Complex entries in a complex build, so `dot` has something to conjugate + z = PetscScalar <: Complex ? im : 0 + xs = PetscScalar[i + z * (n - i) for i in 1:n] + ys = PetscScalar[2i - 1 - z * i for i in 1:n] + α = PetscScalar(3 + 2z) + β = PetscScalar(-0.5) + + @testset "Vec ($PetscScalar)" begin + x = petsc_vec(petsclib, comm, xs) + y = petsc_vec(petsclib, comm, ys) + + w = copy(x) + @test w isa PETSc.PetscVec + @test w.ptr != x.ptr + @test owned(w) == owned(x, xs) + fill!(w, 0) + @test owned(x) == owned(x, xs) + + @test dot(x, y) ≈ dot(xs, ys) + @test dot(x, y) ≈ conj(dot(y, x)) + @test dot(x, x) ≈ sum(abs2, xs) + + results = [ + x + y => xs + ys, + x - y => xs - ys, + -x => -xs, + α * x => α * xs, + x * α => xs * α, + x / α => xs / α, + 2 * x => 2 * xs, + ] + for (v, ref) in results + @test v isa PETSc.PetscVec + @test owned(v) ≈ owned(v, ref) + end + # Division is entry by entry, as for a Julia vector, not a scale by 1/α + @test owned(results[6].first) == owned(x, xs ./ α) + + u = copy(y) + @test axpy!(α, x, u) === u + @test owned(u) ≈ owned(u, α * xs + ys) + @test axpby!(α, x, β, u) === u + @test owned(u) ≈ owned(u, α * xs + β * (α * xs + ys)) + @test rmul!(u, β) === u + @test lmul!(α, u) === u + @test owned(u) ≈ owned(u, α * β * (α * xs + β * (α * xs + ys))) + + # The same vector on both sides, as for Julia arrays + u = copy(x) + axpy!(2, u, u) + @test owned(u) ≈ owned(u, 3 * xs) + axpby!(1, u, 1, u) + @test owned(u) ≈ owned(u, 6 * xs) + + short = PETSc.PetscVec(petsclib, comm, n - 1) + @test_throws DimensionMismatch x + short + @test_throws DimensionMismatch dot(x, short) + @test_throws DimensionMismatch axpy!(1, x, short) + if PetscScalar <: Real + @test_throws InexactError 2im * x + end + + foreach(PETSc.destroy!, (x, y, w, u, short, first.(results)...)) + end + + @testset "Mat ($PetscScalar)" begin + # Upper and lower bidiagonal, so A ± B has the union of the patterns + a = [(i, j, PetscScalar(i + 2j + z)) for i in 1:n for j in (i, i + 1) if j <= n] + b = [(i, j, PetscScalar(3i - j)) for i in 1:n for j in (i - 1, i) if j >= 1] + dense(entries) = (M = zeros(PetscScalar, n, n); foreach(((i, j, v),) -> M[i, j] = v, entries); M) + As, Bs = dense(a), dense(b) + A = petsc_mat(petsclib, comm, n, a) + B = petsc_mat(petsclib, comm, n, b) + rows = PETSc.ownership_range(A) + + C = copy(A) + @test C isa PETSc.PetscMat + @test C.ptr != A.ptr + @test C == A + @test C[rows, 1:n] == As[rows, :] + fill!(C, 0) + @test A[rows, 1:n] == As[rows, :] + + results = [ + A + B => As + Bs, + A - B => As - Bs, + -A => -As, + α * A => α * As, + A * α => As * α, + ] + for (M, ref) in results + @test M isa PETSc.PetscMat + @test M[rows, 1:n] ≈ ref[rows, :] + end + @test A[rows, 1:n] == As[rows, :] + @test B[rows, 1:n] == Bs[rows, :] + + short = petsc_mat(petsclib, comm, n - 1, [(1, 1, PetscScalar(1))]) + @test_throws DimensionMismatch A + short + + foreach(PETSc.destroy!, (A, B, C, short, first.(results)...)) + end + end +end