compiler: Temp fix for memory leaks. (#70)

- Temp fix until VecReplaceArray works properly in conjunction with PetscMemoryAllocator

Author: ZoeLeibowitz

devito/mpi/distributed.py

@@ -91,7 +91,7 @@ def devito_mpi_finalize():
     """
     Finalize MPI, if initialized by Devito.
     """
-    global init_by_devito
+    global init_by_devito  # noqa
     if init_by_devito and MPI.Is_initialized() and not MPI.Is_finalized():
         MPI.Finalize()
 

devito/petsc/iet/passes.py

@@ -11,7 +11,7 @@
 from devito.types.basic import DataSymbol
 from devito.tools import frozendict
 from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData,
-                                PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES,
+                                PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
                                 SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
@@ -40,14 +40,13 @@ def lower_petsc(iet, **kwargs):
             f"{petsc_languages}, but got '{kwargs['language']}'"
         )
 
-    metadata = core_metadata()
     data = FindNodes(PetscMetaData).visit(iet)
 
     if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)):
-        return initialize(iet), metadata
+        return initialize(iet), core_metadata()
 
     if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)):
-        return finalize(iet), metadata
+        return finalize(iet), core_metadata()
 
     unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()}
     # Assumption is that all solves are on the same grid
@@ -79,7 +78,7 @@ def lower_petsc(iet, **kwargs):
     body = core + tuple(setup) + (BlankLine,) + iet.body.body
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
-    metadata.update({'efuncs': tuple(efuncs.values())})
+    metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())}
     return iet, metadata
 
 
@@ -218,13 +217,13 @@ def populate_matrix_context(efuncs, objs):
     'rowidx': PetscInt('rowidx'),
     'colidx': PetscInt('colidx'),
     'J': Mat('J'),
-    'X': GlobalVec('X'),
-    'xloc': LocalVec('xloc'),
-    'Y': GlobalVec('Y'),
-    'yloc': LocalVec('yloc'),
-    'F': GlobalVec('F'),
-    'floc': LocalVec('floc'),
-    'B': GlobalVec('B'),
+    'X': Vec('X'),
+    'xloc': CallbackVec('xloc'),
+    'Y': Vec('Y'),
+    'yloc': CallbackVec('yloc'),
+    'F': Vec('F'),
+    'floc': CallbackVec('floc'),
+    'B': Vec('B'),
     'nfields': PetscInt('nfields'),
     'irow': PointerIS(name='irow'),
     'icol': PointerIS(name='icol'),

devito/petsc/iet/routines.py

@@ -16,7 +16,7 @@
                                     MatShellSetOp, PetscMetaData)
 from devito.petsc.iet.utils import petsc_call, petsc_struct
 from devito.petsc.utils import solver_mapper
-from devito.petsc.types import (DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES,
+from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES,
                                 PetscInt, StartPtr, PointerIS, PointerDM, VecScatter,
                                 DMCast, JacobianStructCast, JacobianStruct,
                                 SubMatrixStruct, CallbackDM)
@@ -448,8 +448,13 @@ def _create_form_rhs_body(self, body, fielddata):
             'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)]
         )
 
+        dm_restore_local_bvec = petsc_call(
+            'DMRestoreLocalVector', [dmda, Byref(sobjs['blocal'])]
+        )
+
         body = body._rebuild(body=body.body + (
-            dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array
+            dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array,
+            dm_restore_local_bvec
         ))
 
         stacks = (
@@ -870,10 +875,10 @@ def _build(self):
         targets = self.fielddata.targets
         base_dict = {
             'Jac': Mat(sreg.make_name(prefix='J')),
-            'xglobal': GlobalVec(sreg.make_name(prefix='xglobal')),
-            'xlocal': LocalVec(sreg.make_name(prefix='xlocal')),
-            'bglobal': GlobalVec(sreg.make_name(prefix='bglobal')),
-            'blocal': LocalVec(sreg.make_name(prefix='blocal')),
+            'xglobal': Vec(sreg.make_name(prefix='xglobal')),
+            'xlocal': Vec(sreg.make_name(prefix='xlocal')),
+            'bglobal': Vec(sreg.make_name(prefix='bglobal')),
+            'blocal': CallbackVec(sreg.make_name(prefix='blocal')),
             'ksp': KSP(sreg.make_name(prefix='ksp')),
             'pc': PC(sreg.make_name(prefix='pc')),
             'snes': SNES(sreg.make_name(prefix='snes')),
@@ -939,9 +944,9 @@ def _extend_build(self, base_dict):
                 name=f'{key}ctx',
                 fields=objs['subctx'].fields,
             )
-            base_dict[f'{key}X'] = LocalVec(f'{key}X')
-            base_dict[f'{key}Y'] = LocalVec(f'{key}Y')
-            base_dict[f'{key}F'] = LocalVec(f'{key}F')
+            base_dict[f'{key}X'] = CallbackVec(f'{key}X')
+            base_dict[f'{key}Y'] = CallbackVec(f'{key}Y')
+            base_dict[f'{key}F'] = CallbackVec(f'{key}F')
 
         return base_dict
 
@@ -953,22 +958,22 @@ 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}'] = LocalVec(
+            base_dict[f'xlocal{name}'] = CallbackVec(
                 sreg.make_name(prefix=f'xlocal{name}'), liveness='eager'
             )
-            base_dict[f'Fglobal{name}'] = LocalVec(
+            base_dict[f'Fglobal{name}'] = CallbackVec(
                 sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager'
             )
-            base_dict[f'Xglobal{name}'] = LocalVec(
+            base_dict[f'Xglobal{name}'] = CallbackVec(
                 sreg.make_name(prefix=f'Xglobal{name}')
             )
-            base_dict[f'xglobal{name}'] = GlobalVec(
+            base_dict[f'xglobal{name}'] = Vec(
                 sreg.make_name(prefix=f'xglobal{name}')
             )
-            base_dict[f'blocal{name}'] = LocalVec(
+            base_dict[f'blocal{name}'] = CallbackVec(
                 sreg.make_name(prefix=f'blocal{name}'), liveness='eager'
             )
-            base_dict[f'bglobal{name}'] = GlobalVec(
+            base_dict[f'bglobal{name}'] = Vec(
                 sreg.make_name(prefix=f'bglobal{name}')
             )
             base_dict[f'da{name}'] = DM(
@@ -1021,6 +1026,12 @@ def _setup(self):
         global_x = petsc_call('DMCreateGlobalVector',
                               [dmda, Byref(sobjs['xglobal'])])
 
+        local_x = petsc_call('DMCreateLocalVector',
+                             [dmda, Byref(sobjs['xlocal'])])
+
+        get_local_size = petsc_call('VecGetSize',
+                                    [sobjs['xlocal'], Byref(sobjs['localsize'])])
+
         global_b = petsc_call('DMCreateGlobalVector',
                               [dmda, Byref(sobjs['bglobal'])])
 
@@ -1078,6 +1089,8 @@ def _setup(self):
             snes_set_jac,
             snes_set_type,
             global_x,
+            local_x,
+            get_local_size,
             global_b,
             snes_get_ksp,
             ksp_set_tols,
@@ -1235,9 +1248,6 @@ def _execute_solve(self):
 
         rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']])
 
-        local_x = petsc_call('DMCreateLocalVector',
-                             [dmda, Byref(sobjs['xlocal'])])
-
         vec_replace_array = self.timedep.replace_array(target)
 
         dm_local_to_global_x = petsc_call(
@@ -1253,13 +1263,15 @@ def _execute_solve(self):
             dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']]
         )
 
+        vec_reset_array = self.timedep.reset_array(target)
+
         run_solver_calls = (struct_assignment,) + (
             rhs_call,
-            local_x
         ) + vec_replace_array + (
             dm_local_to_global_x,
             snes_solve,
             dm_global_to_local_x,
+            vec_reset_array,
             BlankLine,
         )
         return List(body=run_solver_calls)
@@ -1402,7 +1414,16 @@ def replace_array(self, target):
             target.function._C_field_data, target.function._C_symbol
         )
         xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal'])
-        return (petsc_call('VecReplaceArray', [xlocal, field_from_ptr]),)
+        return (petsc_call('VecPlaceArray', [xlocal, field_from_ptr]),)
+
+    def reset_array(self, target):
+        """
+        """
+        sobjs = self.sobjs
+        xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal'])
+        return (
+            petsc_call('VecResetArray', [xlocal])
+        )
 
     def assign_time_iters(self, struct):
         return []
@@ -1530,15 +1551,14 @@ def replace_array(self, target):
 
             caster = cast(target.dtype, '*')
             return (
-                petsc_call('VecGetSize', [xlocal, Byref(sobjs['localsize'])]),
                 DummyExpr(
                     start_ptr,
                     caster(
                         FieldFromPointer(target._C_field_data, target._C_symbol)
                     ) + Mul(target_time, sobjs['localsize']),
                     init=True
                 ),
-                petsc_call('VecReplaceArray', [xlocal, start_ptr])
+                petsc_call('VecPlaceArray', [xlocal, start_ptr])
             )
         return super().replace_array(target)
 

devito/petsc/solve.py

@@ -70,29 +70,33 @@ def build_function_eqns(self, eq, target, arrays):
         formfunc = self.make_formfunc(eq, F_target, arrays, targets)
         formrhs = self.make_rhs(eq, b, arrays)
 
-        return tuple(expr.subs(self.time_mapper) for expr in (formfunc, formrhs))
+        return (formfunc, formrhs)
 
     def build_matvec_eqns(self, eq, target, arrays):
         b, F_target, targets = separate_eqn(eq, target)
         if not F_target:
             return None
         matvec = self.make_matvec(eq, F_target, arrays, targets)
-        return matvec.subs(self.time_mapper)
+        return matvec
 
     def make_matvec(self, eq, F_target, arrays, targets):
-        rhs = arrays['x'] if isinstance(eq, EssentialBC) else F_target.subs(
-            targets_to_arrays(arrays['x'], targets)
-        )
+        if isinstance(eq, EssentialBC):
+            rhs = arrays['x']
+        else:
+            rhs = F_target.subs(targets_to_arrays(arrays['x'], targets))
+            rhs = rhs.subs(self.time_mapper)
         return Eq(arrays['y'], rhs, subdomain=eq.subdomain)
 
     def make_formfunc(self, eq, F_target, arrays, targets):
-        rhs = 0. if isinstance(eq, EssentialBC) else F_target.subs(
-            targets_to_arrays(arrays['x'], targets)
-        )
+        if isinstance(eq, EssentialBC):
+            rhs = 0.
+        else:
+            rhs = F_target.subs(targets_to_arrays(arrays['x'], targets))
+            rhs = rhs.subs(self.time_mapper)
         return Eq(arrays['f'], rhs, subdomain=eq.subdomain)
 
     def make_rhs(self, eq, b, arrays):
-        rhs = 0. if isinstance(eq, EssentialBC) else b
+        rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper)
         return Eq(arrays['b'], rhs, subdomain=eq.subdomain)
 
     def generate_arrays(self, target):

devito/petsc/types/object.py

@@ -66,28 +66,14 @@ def _C_free_priority(self):
         return 2
 
 
-class LocalVec(LocalObject):
+class CallbackVec(LocalObject):
     """
-    PETSc local vector object (Vec).
-    A local vector has ghost locations that contain values that are
-    owned by other MPI ranks.
+    PETSc vector object (Vec).
     """
     dtype = CustomDtype('Vec')
 
 
-class CallbackGlobalVec(LocalVec):
-    """
-    PETSc global vector object (Vec). For example, used for coupled
-    solves inside the `WholeFormFunc` callback.
-    """
-
-
-class GlobalVec(LocalVec):
-    """
-    PETSc global vector object (Vec).
-    A global vector is a parallel vector that has no duplicate values
-    between MPI ranks. A global vector has no ghost locations.
-    """
+class Vec(CallbackVec):
     @property
     def _C_free(self):
         return petsc_call('VecDestroy', [Byref(self.function)])

devito/petsc/utils.py

@@ -76,12 +76,25 @@ def get_petsc_variables():
 
 petsc_variables = get_petsc_variables()
 
-# TODO: Check to see whether Petsc is compiled with
-# 32-bit or 64-bit integers
-# TODO: Check whether PetscScalar is a float or double
-# and only map the right one
-petsc_type_mappings = {ctypes.c_int: 'PetscInt',
-                       ctypes.c_float: 'PetscScalar',
-                       ctypes.c_double: 'PetscScalar'}
+
+def get_petsc_type_mappings():
+    try:
+        petsc_precision = petsc_variables['PETSC_PRECISION']
+    except KeyError:
+        mapper = {}
+    else:
+        petsc_scalar = 'PetscScalar'
+        # TODO: Check to see whether Petsc is compiled with
+        # 32-bit or 64-bit integers
+        mapper = {ctypes.c_int: 'PetscInt'}
+
+        if petsc_precision == 'single':
+            mapper[ctypes.c_float] = petsc_scalar
+        elif petsc_precision == 'double':
+            mapper[ctypes.c_double] = petsc_scalar
+    return mapper
+
+
+petsc_type_mappings = get_petsc_type_mappings()
 
 petsc_languages = ['petsc']

tests/test_petsc.py

@@ -9,7 +9,7 @@
                            FindNodes, retrieve_iteration_tree)
 from devito.types import Constant, LocalCompositeObject
 from devito.passes.iet.languages.C import CDataManager
-from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP,
+from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP,
                                 PC, KSPConvergedReason, PETScArray,
                                 LinearSolveExpr, FieldData, MultipleFieldData)
 from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil
@@ -33,7 +33,7 @@ def test_petsc_local_object():
     """
     lo0 = DM('da', stencil_width=1)
     lo1 = Mat('A')
-    lo2 = LocalVec('x')
+    lo2 = Vec('x')
     lo3 = PetscMPIInt('size')
     lo4 = KSP('ksp')
     lo5 = PC('pc')
@@ -590,7 +590,7 @@ def test_apply():
 
     pn = Function(name='pn', grid=grid, space_order=2)
     rhs = Function(name='rhs', grid=grid, space_order=2)
-    mu = Constant(name='mu', value=2.0)
+    mu = Constant(name='mu', value=2.0, dtype=np.float64)
 
     eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior)
 
@@ -604,7 +604,7 @@ def test_apply():
         op.apply()
 
     # Verify that users can override `mu`
-    mu_new = Constant(name='mu_new', value=4.0)
+    mu_new = Constant(name='mu_new', value=4.0, dtype=np.float64)
     op.apply(mu=mu_new)
 
 
@@ -627,9 +627,10 @@ def test_petsc_frees():
     # Check the frees appear in the following order
     assert str(frees[0]) == 'PetscCall(VecDestroy(&(bglobal0)));'
     assert str(frees[1]) == 'PetscCall(VecDestroy(&(xglobal0)));'
-    assert str(frees[2]) == 'PetscCall(MatDestroy(&(J0)));'
-    assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes0)));'
-    assert str(frees[4]) == 'PetscCall(DMDestroy(&(da0)));'
+    assert str(frees[2]) == 'PetscCall(VecDestroy(&(xlocal0)));'
+    assert str(frees[3]) == 'PetscCall(MatDestroy(&(J0)));'
+    assert str(frees[4]) == 'PetscCall(SNESDestroy(&(snes0)));'
+    assert str(frees[5]) == 'PetscCall(DMDestroy(&(da0)));'
 
 
 @skipif('petsc')