Add a parallel matrix type and test

Author: Jeremy E Kozdon

src/mat.jl

@@ -100,6 +100,109 @@ function MatSeqDense(
     return mat
 end
 
+"""
+    MatAIJ(
+        petsclib::PetscLib,
+        comm::MPI.Comm,
+        loc_num_rows::Integer,
+        loc_num_cols::Integer,
+        diag_nonzeros::Union{Integer, Vector},
+        off_diag_nonzeros::Union{Integer, Vector};
+        glo_num_rows = PETSC_DETERMINE,
+        glo_num_cols = PETSC_DETERMINE,
+        setup = true
+    ) where {PetscLib <: PetscLibType}
+
+Create an MPI PETSc sparse array on the `comm` using AIJ format (also known as a
+compressed sparse row or CSR format) of size `glo_num_rows X glo_num_cols` with
+local size `loc_num_rows X loc_num_cols`.
+
+The diagonal block and off-diagonal block non-zeros are `diag_nonzeros` and
+`off_diag_nonzeros` which can be either an integer (same for all rows) or a
+Vector of `PetscInt`s with on entry per row.
+
+Memory allocation is handled by PETSc and garbage collection can be used.
+
+If `glo_num_rows isa Integer` or `glo_num_cols isa Integer` then the
+corresponding local variable can be `PETSC_DECIDE`.
+
+If `setup == true` then [`setup!`](@ref) is called
+
+# External Links
+$(_doc_external("Mat/MatCreateAIJ"))
+$(_doc_external("Mat/MatSetUp"))
+
+!!! note
+
+    The user is responsible for calling `destroy(mat)` on the `MatAIJ` since
+    this cannot be handled by the garbage collector do to the MPI nature of the
+    object.
+"""
+mutable struct MatAIJ{PetscLib, PetscScalar} <:
+               AbstractMat{PetscLib, PetscScalar}
+    ptr::CMat
+end
+
+function MatAIJ(
+    petsclib::PetscLib,
+    comm::MPI.Comm,
+    loc_num_rows::Integer,
+    loc_num_cols::Integer,
+    diag_nonzeros::Union{Integer, Vector},
+    off_diag_nonzeros::Union{Integer, Vector};
+    glo_num_rows = PETSC_DETERMINE,
+    glo_num_cols = PETSC_DETERMINE,
+    setup = true,
+) where {PetscLib <: PetscLibType}
+    @assert initialized(petsclib)
+    PetscScalar = petsclib.PetscScalar
+    mat = MatAIJ{PetscLib, PetscScalar}(C_NULL)
+    if diag_nonzeros isa Integer
+        diag_nonzero = diag_nonzeros
+        diag_nonzeros = C_NULL
+    else
+        diag_nonzero = -1
+    end
+    if off_diag_nonzeros isa Integer
+        off_diag_nonzero = off_diag_nonzeros
+        off_diag_nonzeros = C_NULL
+    else
+        off_diag_nonzero = -1
+    end
+
+    LibPETSc.MatCreateAIJ(
+        petsclib,
+        comm,
+        loc_num_rows,
+        loc_num_cols,
+        glo_num_rows,
+        glo_num_cols,
+        diag_nonzero,
+        diag_nonzeros,
+        off_diag_nonzero,
+        off_diag_nonzeros,
+        mat,
+    )
+
+    setup && setup!(mat)
+
+    return mat
+end
+
+"""
+    setup!(mat::AbstractMat)
+
+Set up the interal data for `mat`
+
+# External Links
+$(_doc_external("Mat/MatSetUp"))
+"""
+function setup!(mat::AbstractMat{PetscLib}) where {PetscLib}
+    @assert initialized(PetscLib)
+    LibPETSc.MatSetUp(PetscLib, mat)
+    return mat
+end
+
 """
     assemble!(M::AbstractMat[, t::MatAssemblyType = MAT_FINAL_ASSEMBLY)
 
@@ -151,6 +254,35 @@ function assemblyend!(
     return M
 end
 
+"""
+    ownershiprange(mat::AbstractMat, [base_one = true])
+
+The range of row indices owned by this processor, assuming that the `mat` is
+laid out with the first `n1` rows on the first processor, next `n2` rows on the
+second, etc. For certain parallel layouts this range may not be well defined.
+
+If the optional argument `base_one == true` then base-1 indexing is used,
+otherwise base-0 index is used.
+
+!!! note
+
+    unlike the C function, the range returned is inclusive (`idx_first:idx_last`)
+
+# External Links
+$(_doc_external("Mat/MatGetOwnershipRange"))
+"""
+function ownershiprange(
+    mat::AbstractMat{PetscLib},
+    base_one::Bool = true,
+) where {PetscLib}
+    PetscInt = PetscLib.PetscInt
+    r_lo = Ref{PetscInt}()
+    r_hi = Ref{PetscInt}()
+    LibPETSc.MatGetOwnershipRange(PetscLib, mat, r_lo, r_hi)
+    return base_one ? ((r_lo[] + PetscInt(1)):(r_hi[])) :
+           ((r_lo[]):(r_hi[] - PetscInt(1)))
+end
+
 function Base.size(A::AbstractMat{PetscLib}) where {PetscLib}
     m = Ref{PetscLib.PetscInt}()
     n = Ref{PetscLib.PetscInt}()
@@ -342,3 +474,27 @@ function Base.copyto!(M::PETSc.MatSeqAIJ{T}, S::SparseMatrixCSC{T}) where {T}
     assemble(M);  
 end
 =#
+
+"""
+    createvecs(
+        M::AbstractMat{PetscLib},
+    )
+
+Returns vectors `V` which are compatible with `M`. A right compatible vectors is
+`V.right` and a left compatible vector is `V.left`; positionally these are
+returned as `(right, left)`
+
+The created vectors are not garbage collected and should be destroyed with
+[`destroy`](@ref).
+
+# External Links
+$(_doc_external("Mat/MatCreateVecs"))
+"""
+function createvecs(M::AbstractMat{PetscLib}) where {PetscLib}
+    r_right = Ref{CVec}()
+    r_left = Ref{CVec}()
+    LibPETSc.MatCreateVecs(PetscLib, M, r_right, r_left)
+    right = VecPtr(PetscLib, r_right[])
+    left = VecPtr(PetscLib, r_left[])
+    return (right = right, left = left)
+end

src/vec.jl

@@ -239,7 +239,7 @@ function setvalues!(
     idxs0::Vector{PetscInt},
     vals::Array{PetscScalar},
     insertmode::InsertMode;
-    num_idxs = length(idxs0)
+    num_idxs = length(idxs0),
 ) where {PetscLib, PetscInt, PetscScalar}
     @assert length(vals) >= num_idxs
     @assert PetscInt == PetscLib.PetscInt
@@ -267,7 +267,7 @@ function getvalues!(
     vals::Array{PetscScalar},
     v::AbstractVec{PetscLib},
     idxs0::Vector{PetscInt};
-    num_idxs = length(idxs0)
+    num_idxs = length(idxs0),
 ) where {PetscLib, PetscInt, PetscScalar}
     @assert length(vals) >= num_idxs
     @assert PetscInt == PetscLib.PetscInt
@@ -425,7 +425,7 @@ end
 !!! note
     `Base.finalize` is automatically called on the array.
 """
-function with_localarray!(f!, vecs::NTuple{N, AbstractVec}; kwargs...) where N
+function with_localarray!(f!, vecs::NTuple{N, AbstractVec}; kwargs...) where {N}
     arrays = map(vecs) do v
         unsafe_localarray(v; kwargs...)
     end
@@ -462,6 +462,23 @@ function assemblyend!(vec::AbstractVec{PetscLib}) where {PetscLib}
     return nothing
 end
 
+"""
+    assemble!(v::AbstractVec)
+
+Assemble the vector `v`.
+
+For overlapping assembly see [`assemblybegin!`](@ref) and  [`assemblyend!`](@ref)
+
+# External Links
+$(_doc_external("Vec/VecAssemblyBegin"))
+$(_doc_external("Vec/VecAssemblyEnd"))
+"""
+function assemble!(v::AbstractVec)
+    assemblybegin!(v)
+    assemblyend!(v)
+    return v
+end
+
 mutable struct LocalVec{PetscLib, PetscScalar, GVec} <:
                AbstractVec{PetscLib, PetscScalar}
     ptr::CVec
@@ -485,7 +502,7 @@ Obtains the local ghosted representation of a [`Vec`](@ref).
 # External Links
 $(_doc_external("Vec/VecGhostGetLocalForm"))
 """
-function getlocalform(gvec::AbstractVec{PetscLib}) where PetscLib
+function getlocalform(gvec::AbstractVec{PetscLib}) where {PetscLib}
     lvec = LocalVec(gvec)
     LibPETSc.VecGhostGetLocalForm(PetscLib, gvec, lvec)
     if lvec.ptr == C_NULL
@@ -504,7 +521,7 @@ Restore the `local_vec` to the associated global vector after a call to
 # External Links
 $(_doc_external("Vec/VecGhostRestoreLocalForm"))
 """
-function restorelocalform!(lvec::LocalVec{PetscLib}) where PetscLib
+function restorelocalform!(lvec::LocalVec{PetscLib}) where {PetscLib}
     LibPETSc.VecGhostRestoreLocalForm(PetscLib, lvec.gvec, lvec)
     lvec.ptr = C_NULL
     return lvec.gvec
@@ -548,7 +565,7 @@ function ghostupdatebegin!(
     vec::AbstractVec{PetscLib},
     insertmode = INSERT_VALUES,
     scattermode = SCATTER_FORWARD,
-) where PetscLib
+) where {PetscLib}
     LibPETSc.VecGhostUpdateBegin(PetscLib, vec, insertmode, scattermode)
     return nothing
 end
@@ -569,13 +586,12 @@ function ghostupdateend!(
     vec::AbstractVec{PetscLib},
     insertmode = INSERT_VALUES,
     scattermode = SCATTER_FORWARD,
-) where PetscLib
+) where {PetscLib}
     LibPETSc.VecGhostUpdateEnd(PetscLib, vec, insertmode, scattermode)
     return nothing
 end
 
-
-function getpetsctype(vec::AbstractVec{PetscLib}) where PetscLib
+function getpetsctype(vec::AbstractVec{PetscLib}) where {PetscLib}
     name_r = Ref{LibPETSc.VecType}()
     LibPETSc.VecGetType(PetscLib, vec, name_r)
     return unsafe_string(name_r[])

test/mpimat.jl

@@ -0,0 +1,86 @@
+using Test
+using MPI
+MPI.Initialized() || MPI.Init()
+using PETSc
+using LinearAlgebra: mul!, norm
+
+@testset "MatAIJ" begin
+    comm = MPI.COMM_WORLD
+    mpisize = MPI.Comm_size(comm)
+    mpirank = MPI.Comm_rank(comm)
+
+    for petsclib in PETSc.petsclibs
+        PETSc.initialize(petsclib)
+        PetscScalar = petsclib.PetscScalar
+        PetscInt = petsclib.PetscInt
+
+        loc_num_rows = 5
+        loc_num_cols = 5
+        diag_nonzeros = 3
+        off_diag_non_zeros = 3
+
+        mat = PETSc.MatAIJ(
+            petsclib,
+            comm,
+            loc_num_rows,
+            loc_num_cols,
+            diag_nonzeros,
+            off_diag_non_zeros,
+        )
+
+        # Get compatible vectors
+        (right, left) = PETSc.createvecs(mat)
+
+        # Fill the matrix and right vector
+        row_rng = PETSc.ownershiprange(mat, false)
+        for i in row_rng
+            if i == 0
+                vals = [-2, 1]
+                row0idxs = [i]
+                col0idxs = [i, i + 1]
+            elseif i == mpisize * loc_num_rows - 1
+                vals = [-2, 1]
+                row0idxs = [i]
+                col0idxs = [i, i - 1]
+            else
+                vals = [1, -2, 1]
+                row0idxs = [i]
+                col0idxs = [i - 1, i, i + 1]
+            end
+            PETSc.setvalues!(
+                mat,
+                PetscInt.(row0idxs),
+                PetscInt.(col0idxs),
+                PetscScalar.(vals),
+            )
+            right[i + 1] = i^3
+        end
+        PETSc.assemble!(mat)
+        PETSc.assemble!(right)
+
+        # Do matrix multiply and check result
+        mul!(left, mat, right)
+        for i in row_rng
+            if i == 0
+                v = -2 * i^3 + (i + 1)^3
+            elseif i == mpisize * loc_num_rows - 1
+                v = (i - 1)^3 - 2 * i^3
+            else
+                v = (i - 1)^3 - 2 * i^3 + (i + 1)^3
+            end
+            @test v == left[i + 1]
+        end
+
+        # Check the norm in parallel
+        sz = loc_num_rows * mpisize
+        exact_norm = sqrt(sz * 2^2 + 2 * (sz - 1))
+        @test norm(mat) ≈ exact_norm
+
+        PETSc.destroy(mat)
+        PETSc.destroy(right)
+        PETSc.destroy(left)
+
+        PETSc.finalize(petsclib)
+    end
+end
+nothing

test/runtests.jl

@@ -3,9 +3,12 @@ using MPI: mpiexec
 
 # Do the MPI tests first so we do not have mpi running inside MPI
 @testset "mpi tests" begin
-    @test mpiexec() do mpi_cmd
-        cmd = `$mpi_cmd -n 4 $(Base.julia_cmd()) --startup-file=no --project mpivec.jl`
-        success(pipeline(cmd, stderr = stderr))
+    for file in ("mpivec.jl", "mpimat.jl")
+        @test mpiexec() do mpi_cmd
+            cmd =
+                `$mpi_cmd -n 4 $(Base.julia_cmd()) --startup-file=no --project $file`
+            success(pipeline(cmd, stderr = stderr))
+        end
     end
 end