Add copy, dot, axpy! and arithmetic for Vec and Mat

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.

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}}

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).

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
 ```
 

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:

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,
 ])
 
 # ---------------------------------------------------------------------------

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)

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 ─────────────────────────────────
 
 """

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)