Repeated PyMat.createSubMatrix leaks submatrix

Consider:

import petsc4py
import sys
petsc4py.init(sys.argv)
from petsc4py import PETSc

class pymat(object):

    def createSubMatrix(self, mat, row_is, col_is, target=None):
        if target is not None:
            target.assemble()
            return target

        rsize = row_is.getSize()
        csize = row_is.getSize()

        m = PETSc.Mat().create(comm=mat.comm)
        m.setType("python")
        m.setSizes(((rsize, None), (csize, None)))

        ctx = pymat()

        m.setPythonContext(ctx)

        m.setUp()
        return m


mat = PETSc.Mat().create()
mat.setSizes((10, 10))

mat.setType("python")
mat.setPythonContext(pymat())
mat.setUp()


iset = PETSc.IS().createStride(0,4, 1)

sub = mat.createSubMatrix(iset, iset)

sub = mat.createSubMatrix(iset, iset, submat=sub)

sub.destroy()
mat.destroy()
iset.destroy()

$ python foo.py -log_view
...
--- Event Stage 0: Main Stage

              Viewer     1              0            0     0.
           Index Set     1              1          800     0.
              Matrix     2              1         2792     0.

So we leak a Mat. It looks like the problem is that MatCreateSubMatrix_Python always calls addRef on the matrix returned from the user-provided createSubMatrix call. But I think this is the wrong thing to do in the MAT_REUSE_MATRIX case. We already hold a reference to the submatrix, and just need to update the contents.

How about this change:

diff --git a/src/libpetsc4py/libpetsc4py.pyx b/src/libpetsc4py/libpetsc4py.pyx
index f254191..93a801a 100644
--- a/src/libpetsc4py/libpetsc4py.pyx
+++ b/src/libpetsc4py/libpetsc4py.pyx
@@ -703,7 +703,7 @@ cdef PetscErrorCode MatCreateSubMatrix_Python(
     PetscMat *out,
     ) \
     except IERR with gil:
-    FunctionBegin(b"MatCopy_Python")
+    FunctionBegin(b"MatCreateSubMatrix_Python")
     cdef createSubMatrix = PyMat(mat).createSubMatrix
     if createSubMatrix is None: return UNSUPPORTED(b"createSubMatrix")
     cdef Mat sub = None
@@ -711,10 +711,11 @@ cdef PetscErrorCode MatCreateSubMatrix_Python(
         sub = None
     elif op == MAT_INITIAL_MATRIX:
         sub = createSubMatrix(Mat_(mat), IS_(row), IS_(col), None)
+        if sub is not None:
+            addRef(sub.mat)
     elif op == MAT_REUSE_MATRIX:
         sub = createSubMatrix(Mat_(mat), IS_(row), IS_(col), Mat_(out[0]))
     if sub is not None:
-        addRef(sub.mat)
         out[0] = sub.mat
     return FunctionEnd()

Comments (3)