Rely on PETSc duplicated comms

JuliaParallel · Jul 9, 2021 · 827e7e1 · 827e7e1
1 parent d3f5d21
commit 827e7e1
Show file tree

Hide file tree

Showing 8 changed files with 73 additions and 21 deletions.
diff --git a/src/PETSc.jl b/src/PETSc.jl
@@ -31,5 +31,6 @@ include("matshell.jl")
 include("ksp.jl")
 include("pc.jl")
 include("snes.jl")
+include("sys.jl")
 
 end
diff --git a/src/ksp.jl b/src/ksp.jl
@@ -6,7 +6,6 @@ abstract type AbstractKSP{T, PetscLib} <: Factorization{T} end
 
 mutable struct KSP{T, PetscLib} <: AbstractKSP{T, PetscLib}
     ptr::CKSP
-    comm::MPI.Comm
     # keep around so that they don't get gc'ed
     A  # Operator
     P  # preconditioning operator
@@ -30,7 +29,7 @@ LinearAlgebra.adjoint(ksp) = LinearAlgebra.Adjoint(ksp)
     function KSP{$PetscScalar}(comm::MPI.Comm; kwargs...)
         @assert initialized($petsclib)
         opts = Options($petsclib, kwargs...)
-        ksp = KSP{$PetscScalar, $PetscLib}(C_NULL, comm, nothing, nothing, opts)
+        ksp = KSP{$PetscScalar, $PetscLib}(C_NULL, nothing, nothing, opts)
         @chk ccall((:KSPCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CKSP}), comm, ksp)
         if comm == MPI.COMM_SELF
             finalizer(destroy, ksp)
@@ -75,7 +74,7 @@ LinearAlgebra.adjoint(ksp) = LinearAlgebra.Adjoint(ksp)
         return r_its[]
     end
 
-    function view(ksp::KSP{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, ksp.comm))
+    function view(ksp::KSP{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, getcomm(ksp)))
         @chk ccall((:KSPView, $libpetsc), PetscErrorCode, 
                     (CKSP, CPetscViewer),
                 ksp, viewer);
@@ -129,7 +128,7 @@ Construct a PETSc Krylov subspace solver.
 Any PETSc options prefixed with `ksp_` and `pc_` can be passed as keywords.
 """
 function KSP(A::AbstractMat{T}, P::AbstractMat{T}=A; kwargs...) where {T}
-    ksp = KSP{T}(A.comm; kwargs...)
+    ksp = KSP{T}(getcomm(A); kwargs...)
     setoperators!(ksp, A, P)
     with(ksp.opts) do
         setfromoptions!(ksp)

diff --git a/src/mat.jl b/src/mat.jl
@@ -22,7 +22,6 @@ Memory allocation is handled by PETSc.
 """
 mutable struct MatSeqAIJ{T} <: AbstractMat{T}
     ptr::CMat
-    comm::MPI.Comm
 end
 
 """
@@ -32,7 +31,6 @@ PETSc dense array. This wraps a Julia `Matrix{T}` object.
 """
 mutable struct MatSeqDense{T} <: AbstractMat{T}
     ptr::CMat
-    comm::MPI.Comm
     array::Matrix{T}
 end
 
@@ -42,7 +40,7 @@ end
     function MatSeqAIJ{$PetscScalar}(m::Integer, n::Integer, nnz::Vector{$PetscInt})
         @assert initialized($petsclib)
         comm = MPI.COMM_SELF
-        mat = MatSeqAIJ{$PetscScalar}(C_NULL, comm)
+        mat = MatSeqAIJ{$PetscScalar}(C_NULL)
         @chk ccall((:MatCreateSeqAIJ, $libpetsc), PetscErrorCode,
                 (MPI.MPI_Comm, $PetscInt, $PetscInt, $PetscInt, Ptr{$PetscInt}, Ptr{CMat}),
                 comm, m, n, 0, nnz, mat)
@@ -52,7 +50,7 @@ end
     function MatSeqDense(A::Matrix{$PetscScalar})
         @assert initialized($petsclib)
         comm = MPI.COMM_SELF
-        mat = MatSeqDense(C_NULL, comm, A)
+        mat = MatSeqDense(C_NULL, A)
         @chk ccall((:MatCreateSeqDense, $libpetsc), PetscErrorCode, 
             (MPI.MPI_Comm, $PetscInt, $PetscInt, Ptr{$PetscScalar}, Ptr{CMat}),
             comm, size(A,1), size(A,2), A, mat)
@@ -92,7 +90,7 @@ end
         @chk ccall((:MatAssemblyEnd, $libpetsc), PetscErrorCode, (CMat, MatAssemblyType), M, t)
         return nothing
     end
-    function view(mat::AbstractMat{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, mat.comm))
+    function view(mat::AbstractMat{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, getcomm(mat)))
         @chk ccall((:MatView, $libpetsc), PetscErrorCode, 
                     (CMat, CPetscViewer),
                 mat, viewer);

diff --git a/src/matshell.jl b/src/matshell.jl
@@ -9,7 +9,6 @@ This can be changed by defining `PETSc._mul!`.
 """
 mutable struct MatShell{T,A} <: AbstractMat{T}
     ptr::CMat
-    comm::MPI.Comm
     obj::A
 end
 
@@ -30,7 +29,7 @@ MatShell{T}(obj, m, n) where {T} = MatShell{T}(obj, MPI.COMM_SELF, m, n, m, n)
 
 @for_libpetsc begin
     function MatShell{$PetscScalar}(obj::A, comm::MPI.Comm, m, n, M, N) where {A}
-        mat = MatShell{$PetscScalar,A}(C_NULL, comm, obj)
+        mat = MatShell{$PetscScalar,A}(C_NULL, obj)
         # we use the MatShell object itsel
         ctx = pointer_from_objref(mat)
         @chk ccall((:MatCreateShell, $libpetsc), PetscErrorCode,

diff --git a/src/pc.jl b/src/pc.jl
@@ -5,7 +5,6 @@ const CPCType = Cstring
 
 mutable struct PC{T}
     ptr::Ptr{Cvoid}
-    comm::MPI.Comm
 end
 
 Base.cconvert(::Type{CPC}, obj::PC) = obj.ptr
@@ -17,14 +16,14 @@ scalartype(::PC{T}) where {T} = T
 @for_libpetsc begin
 
     function PC{$PetscScalar}(comm::MPI.Comm)
-        pc = PC{$PetscScalar}(C_NULL, comm)
+        pc = PC{$PetscScalar}(C_NULL)
         @chk ccall((:PCCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CPC}), comm, pc)
         finalizer(destroy, pc)
         return pc
     end
 
     function PC(ksp::KSP{$PetscScalar})
-        pc = PC{$PetscScalar}(C_NULL, ksp.comm)
+        pc = PC{$PetscScalar}(C_NULL)
         @chk ccall((:KSPGetPC, $libpetsc), PetscErrorCode, (CKSP, Ptr{CPC}), ksp, pc)
         incref(pc) # need to manually increment the reference counter
         finalizer(destroy, pc)
@@ -53,7 +52,7 @@ scalartype(::PC{T}) where {T} = T
         return unsafe_string(t_r[])
     end
 
-    function view(pc::PC{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, pc.comm))
+    function view(pc::PC{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, getcomm(pc)))
         @chk ccall((:PCView, $libpetsc), PetscErrorCode,
                     (CPC, CPetscViewer),
                 pc, viewer);

diff --git a/src/snes.jl b/src/snes.jl
@@ -5,7 +5,6 @@ const CSNESType = Cstring
 
 mutable struct SNES{T, PetscLib}
     ptr::CSNES
-    comm::MPI.Comm
     opts::Options{PetscLib}
     fn!
     fn_vec
@@ -51,7 +50,7 @@ end
     function SNES{$PetscScalar}(comm::MPI.Comm; kwargs...)
         @assert initialized($petsclib)
         opts = Options($petsclib, kwargs...)
-        snes = SNES{$PetscScalar, $PetscLib}(C_NULL, comm, opts, nothing, nothing, nothing, nothing, nothing)
+        snes = SNES{$PetscScalar, $PetscLib}(C_NULL, opts, nothing, nothing, nothing, nothing, nothing)
         @chk ccall((:SNESCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CSNES}), comm, snes)
 
         with(snes.opts) do
@@ -104,7 +103,7 @@ end
         return unsafe_string(t_r[])
     end
 
-    function view(snes::SNES{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, snes.comm))
+    function view(snes::SNES{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, getcomm(snes)))
         @chk ccall((:SNESView, $libpetsc), PetscErrorCode,
                     (CSNES, CPetscViewer),
                 snes, viewer);

diff --git a/src/sys.jl b/src/sys.jl
@@ -0,0 +1,58 @@
+const CPetscObject = Ptr{Cvoid}
+
+@for_libpetsc begin
+    function getcomm(
+        obj::Union{
+            AbstractKSP{$PetscScalar},
+            AbstractMat{$PetscScalar},
+            AbstractVec{$PetscScalar},
+        },
+    )
+        comm = MPI.Comm()
+        @chk ccall(
+            (:PetscObjectGetComm, $libpetsc),
+            PetscErrorCode,
+            (CPetscObject, Ptr{MPI.MPI_Comm}),
+            obj,
+            comm,
+        )
+
+        #XXX We should really increase the petsc reference counter.
+        #    But for for some reason the PETSc says that this communicator is
+        #    unknown
+        #=
+        # Call the PetscCommDuplicate to increase reference count
+        @chk ccall(
+            (:PetscCommDuplicate, $libpetsc),
+            PetscErrorCode,
+            (MPI.MPI_Comm, Ptr{MPI.MPI_Comm}, Ptr{Cvoid}),
+            comm,
+            comm,
+            C_NULL,
+        )
+
+        # Register PetscCommDestroy to decriment the reference count
+        finalizer(PetscCommDestroy, comm)
+        =#
+
+        return comm
+    end
+end
+
+#=
+#XXX Not sure why this doesn't work
+@for_libpetsc begin
+    function PetscCommDestroy(
+        comm::MPI.Comm
+    )
+        @show comm.val
+        @chk ccall(
+            (:PetscCommDestroy, $libpetsc),
+            PetscErrorCode,
+            (Ptr{MPI.MPI_Comm},),
+            comm,
+        )
+        return nothing
+    end
+end
+=#
diff --git a/src/vec.jl b/src/vec.jl
@@ -29,7 +29,6 @@ $(_doc_external("Vec/VecCreateSeqWithArray"))
 """
 mutable struct VecSeq{T} <: AbstractVec{T}
     ptr::CVec
-    comm::MPI.Comm
     array::Vector{T}
 end
 
@@ -41,7 +40,7 @@ Base.parent(v::AbstractVec) = v.array
 @for_libpetsc begin
     function VecSeq(comm::MPI.Comm, X::Vector{$PetscScalar}; blocksize=1)
         @assert initialized($petsclib)
-        v = VecSeq(C_NULL, comm, X)
+        v = VecSeq(C_NULL, X)
         @chk ccall((:VecCreateSeqWithArray, $libpetsc), PetscErrorCode,
                 (MPI.MPI_Comm, $PetscInt, $PetscInt, Ptr{$PetscScalar}, Ptr{CVec}),
                 comm, blocksize, length(X), X, v)
@@ -84,7 +83,7 @@ Base.parent(v::AbstractVec) = v.array
         r_lo[]:(r_hi[]-$PetscInt(1))
     end
 
-    function view(vec::AbstractVec{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, vec.comm))
+    function view(vec::AbstractVec{$PetscScalar}, viewer::AbstractViewer{$PetscLib}=ViewerStdout($petsclib, getcomm(vec)))
         @chk ccall((:VecView, $libpetsc), PetscErrorCode,
                     (CVec, CPetscViewer),
                 vec, viewer);