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