compiler: Add destroy matop ctx for submats but need to fix when things are destroyed or not..

Author: ZoeLeibowitz

devito/petsc/iet/callbacks.py

@@ -839,6 +839,7 @@ def _uxreplace_efuncs(self):
 class CoupledCallbackBuilder(BaseCallbackBuilder):
     def __init__(self, **kwargs):
         self._submatrices_callback = None
+        self._destroy_submat_callback = None
         super().__init__(**kwargs)
 
     @property
@@ -866,6 +867,7 @@ def _make_core(self):
         self._make_whole_matvec()
         self._make_whole_formfunc()
         self._make_user_struct_efunc()
+        self._create_destroy_submatrix()
         self._create_submatrices()
         self._efuncs['PopulateMatContext'] = self.objs['dummyefunc']
 
@@ -1064,6 +1066,28 @@ def _whole_formfunc_body(self, body):
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields}
 
         return Uxreplace(subs).visit(formfunc_body)
+    
+    def _create_destroy_submatrix(self):
+        # Need a special destroy because each submatrix has a manually
+        # PetscMalloc'ed context attached via MatShellSetContext
+
+        objs = self.objs
+
+        get_ctx = petsc_call(
+            'MatShellGetContext', [objs['J'], Byref(objs['subctx'])]
+        )
+
+        free_ctx = petsc_call(
+            'PetscFree', [objs['subctx']]
+        )
+
+        body = self._make_callable_body((get_ctx, free_ctx))
+
+        cb = self._make_petsc_callable(
+            'DestroySubMatrixCtx', body, parameters=(objs['J']))
+
+        self._destroy_submat_callback = cb
+        self._efuncs[cb.name] = cb
 
     def _create_submatrices(self):
         body = self._submat_callback_body()
@@ -1153,6 +1177,15 @@ def _submat_callback_body(self):
 
         set_ctx = petsc_call('MatShellSetContext', [objs['block'], objs['subctx']])
 
+        set_destroy_mat_op = petsc_call(
+            'MatShellSetOperation',
+            [
+                objs['block'],
+                'MATOP_DESTROY',
+                MatShellSetOp(self._destroy_submat_callback.name, VOID._dtype, VOID._dtype),
+            ],
+        )
+
         mat_setup = petsc_call('MatSetUp', [objs['block']])
 
         assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block'])
@@ -1169,6 +1202,7 @@ def _submat_callback_body(self):
             dm_set_ctx,
             matset_dm,
             set_ctx,
+            set_destroy_mat_op,
             mat_setup,
             assign_block
         )

devito/petsc/iet/logging.py

@@ -22,6 +22,7 @@ def __init__(self, level, **kwargs):
         self.section_mapper = kwargs.get('section_mapper', {})
         self.inject_solve = kwargs.get('inject_solve', None)
 
+        # TODO: fix the segfault with kspgettype
         if level <= PERF:
             funcs = [
                 # KSP specific

devito/petsc/iet/type_builder.py

@@ -119,7 +119,7 @@ def _extend_build(self, base_dict):
 
         base_dict['jacctx'] = JacobianStruct(
             name=sreg.make_name(prefix=objs['ljacctx'].name),
-            fields=objs['ljacctx'].fields,
+            fields=objs['ljacctx'].fields, no_of_submats=2
         )
 
         for sm in submatrices:
@@ -175,7 +175,7 @@ def _target_dependent(self, base_dict):
             base_dict[f'{name}_ptr'] = StartPtr(
                 sreg.make_name(prefix=f'{name}_ptr'), t.dtype
             )
-            base_dict[f'xlocal{name}'] = CallbackVec(
+            base_dict[f'xlocal{name}'] = Vec(
                 sreg.make_name(prefix=f'xlocal{name}'), liveness='eager'
             )
             base_dict[f'Fglobal{name}'] = CallbackVec(

devito/petsc/types/object.py

@@ -201,6 +201,10 @@ class MatReuse(PetscObject):
 class VecScatter(PetscObject):
     dtype = CustomDtype('VecScatter')
 
+    @property
+    def _C_free(self):
+        return petsc_call('VecScatterDestroy', [Byref(self.function)])
+
 
 class StartPtr(PetscObject):
     def __init__(self, name, dtype):
@@ -228,7 +232,7 @@ class PetscSF(PetscObject):
     dtype = CustomDtype('PetscSF')
 
 
-class PETScStruct(LocalCompositeObject):
+class PETScStruct(LocalCompositeObject, PetscMixin):
 
     @property
     def time_dim_fields(self):
@@ -268,10 +272,35 @@ def parent(self):
 
 class JacobianStruct(PETScStruct):
     def __init__(self, name='jctx', pname='JacobianCtx', fields=None,
-                 modifier='', liveness='lazy'):
+                 modifier='', liveness='lazy', no_of_submats=0):
         super().__init__(name, pname, fields, modifier, liveness)
+        self._no_of_submats = no_of_submats
     _C_modifier = None
 
+    @property
+    def _C_free(self):
+        # from IPython import embed; embed()
+        submats = [i for i in self.fields if isinstance(i, PointerMat)]
+        submats = submats[0]
+        # from IPython import embed; embed()
+        from devito.symbolics import FieldFromComposite
+        destroy_call = [petsc_call('MatDestroy', [Byref(FieldFromComposite(submats.indexed[i], self.function))]) for i in range(self._no_of_submats)]
+        destroy_call.append(petsc_call('PetscFree', [Byref(FieldFromComposite(submats.base, self.function))]))
+        return destroy_call
+        # return petsc_call('PetscFree', [Byref(self.function)])
+    
+
+    # @property
+    # def _C_free(self):
+
+    #     submats = [i for i in self.fields if isinstance(i, PointerMat)]
+    #     destroy_calls = [
+    #         petsc_call('MatDestroy', [Byref(self.indexify().subs({self.dim: i}))])
+    #         for i in range(self._no_of_submats)
+    #     ]
+    #     destroy_calls.append(petsc_call('PetscFree', [self.function]))
+    #     return destroy_calls
+
 
 class SubMatrixStruct(PETScStruct):
     def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None,
@@ -385,9 +414,11 @@ class NofSubMats(Scalar, LocalType):
 
 FREE_PRIORITY = {
     PETScArrayObject: 0,
-    Vec: 1,
-    Mat: 2,
-    SNES: 3,
-    PetscSectionGlobal: 4,
-    DM: 5,
+    JacobianStruct: 1,
+    VecScatter: 2,
+    Vec: 3,
+    Mat: 4,
+    SNES: 5,
+    PetscSectionGlobal: 6,
+    DM: 7,
 }