misc: Clean up

Author: ZoeLeibowitz

devito/ir/support/basic.py

@@ -18,7 +18,6 @@
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
                           CriticalRegion, Function, Symbol, Temp, TempArray,
                           TBArray)
-from devito.types.misc import PostIncrementIndex
 
 __all__ = ['IterationInstance', 'TimedAccess', 'Scope', 'ExprGeometry']
 
@@ -410,11 +409,7 @@ def distance(self, other):
                 #       `self.itintervals=(time, x, y)`, `n=0`
                 continue
             elif not sai and not oai:
-                # TODO: temp fix
-                if any(isinstance(i, PostIncrementIndex)
-                       for i in (self[n], other[n])):
-                    ret.append(S.Zero)
-                elif self[n] - other[n] == 0:
+                if self[n] - other[n] == 0:
                     # E.g., `self=R<a,[4]>` and `other=W<a,[4]>`
                     ret.append(S.Zero)
                 else:

devito/petsc/config.py

@@ -53,8 +53,7 @@ def core_metadata():
     }
 
 
-# Maximum number of bytes (including the null terminator) reserved for a
-# KSPType string in the profiler struct.
+# Maximum number of bytes for a KSPType string in the profiler struct.
 KSPTYPE_MAX_LEN = 64
 
 

devito/petsc/iet/builder.py

@@ -70,10 +70,7 @@ def _setup(self):
         local_size = math.prod(
             v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space
         )
-        # TODO: Check - VecCreateSeqWithArray
-        # local_x = petsc_call('VecCreateMPIWithArray',
-        #                      [sobjs['comm'], 1, local_size, 'PETSC_DECIDE',
-        #                       field_from_ptr, Byref(sobjs['xlocal'])])
+        # TODO: Check - VecCreateSeqWithArray vs VecCreateMPIWithArray
         local_x = petsc_call('VecCreateSeqWithArray',
                              ['PETSC_COMM_SELF', 1, local_size,
                               field_from_ptr, Byref(sobjs['xlocal'])])
@@ -139,7 +136,7 @@ def _extend_setup(self):
 
     def _create_dmda_calls(self, dmda):
         dmda_create = self._create_dmda(dmda)
-        # TODO: probs need to set the dm options prefix the same as snes?
+        # TODO: probably need to set the dm options prefix the same as snes?
         dm_set_from_opts = petsc_call('DMSetFromOptions', [dmda])
         dm_setup = petsc_call('DMSetUp', [dmda])
         dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL'])
@@ -355,7 +352,7 @@ def _create_dmda_calls(self, dmda):
         dmda_create = self._create_dmda(dmda)
 
         # TODO: likely need to set the dm options prefix the same as snes?
-        # likely shouldn't hardcode this option like this.. (should be set in the options
+        # Probably shouldn't hardcode this option.. (should be set in the options
         # callback)
         da_create_section = petsc_call(
             'PetscOptionsSetValue', [Null, String("-da_use_section"), Null]

devito/petsc/iet/callbacks.py

@@ -122,7 +122,7 @@ def _make_core(self):
         # Make the initial guess callback
         if self.field_data.initial_guess.exprs:
             self._make_initial_guess()
-        # Make the callback used to constrain boundary nodes
+        # Make the callback to constrain boundary nodes
         if self.field_data.constrain_bc:
             self._make_constrain_bc()
         self._make_user_struct_efunc()
@@ -650,14 +650,9 @@ def _make_constrain_bc(self):
         Constructs the `CountBCs` and `SetPointBCs` efuncs. Works for both
         single- and multi-field.
         """
-        constrain_bc = self.field_data.constrain_bc
+        constrain_bc_dict = self.field_data.constrain_bc
         sobjs = self.solver_objs
 
-        # Normalize to dict {target: ConstrainBC}
-        if isinstance(constrain_bc, dict):
-            constrain_bc_dict = constrain_bc
-        else:
-            constrain_bc_dict = {self.field_data.target: constrain_bc}
         targets = list(constrain_bc_dict.keys())
 
         all_increment_exprs = [
@@ -741,136 +736,14 @@ def _create_count_bc_body(self, body, pairs):
         return Uxreplace(subs).visit(body)
 
     def _create_set_point_bc_body(self, body, constrain_bc_dict):
-        """Single-field SetPointBCs body. `constrain_bc_dict` has one entry."""
-        (target, constrain_bc), = constrain_bc_dict.items()
-        tname = target.name
-        linsolve_expr = self.inject_solve.expr.rhs
-        objs = self.objs
-        sobjs = self.solver_objs
-
-        dmda = sobjs['callbackdm']
-        ctx = objs['dummyctx']
-
-        dm_get_local_info = petsc_call(
-            'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)]
-        )
-
-        body = self.time_dependence.uxreplace_time(body)
-
-        fields = get_user_struct_fields(body)
-        self._struct_params.extend(fields)
-
-        dm_get_app_context = petsc_call(
-            'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
-        )
-        petsc_obj_comm = Call('PetscObjectComm', arguments=[PetscObjectCast(dmda)])
-        is_create_general = petsc_call(
-            'ISCreateGeneral',
-            [petsc_obj_comm, sobjs[f'numBC_{tname}'], sobjs[f'bcPointsArr_{tname}'],
-             'PETSC_OWN_POINTER', Byref(sobjs['bcPointsIS'])]
-        )
-        malloc_bc_points_arr = petsc_call(
-            'PetscMalloc1',
-            [sobjs[f'numBC_{tname}'], Byref(sobjs[f'bcPointsArr_{tname}']._C_symbol)]
-        )
-        malloc_bc_points = petsc_call(
-            'PetscMalloc1', [1, Byref(sobjs['bcPoints']._C_symbol)]
-        )
-        dummy_expr = DummyExpr(sobjs['bcPoints'].indexed[0], sobjs['bcPointsIS'])
-        set_point_bc = petsc_call(
-            'DMDASetPointBC', [dmda, 1, sobjs['bcPoints'], Null]
-        )
-        body = body._rebuild(
-            body=(
-                (malloc_bc_points_arr,)
-                + body.body
-                + (is_create_general, malloc_bc_points, dummy_expr, set_point_bc,)
-            )
-        )
-
-        derefs = dereference_funcs(ctx, fields)
-        standalones = [
-            Definition(ctx),
-            dm_get_app_context,
-            Definition(sobjs[f'k_iter_{tname}'])
-        ]
-        body = self._make_callable_body(
-            body, standalones=standalones, stacks=(dm_get_local_info,) + derefs
-        )
-
-        subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
-                i in fields if not isinstance(i.function, AbstractFunction)}
-        subs[constrain_bc.counter._C_symbol] = \
-            sobjs[f'bcPointsArr_{tname}'].indexed[sobjs[f'k_iter_{tname}']]
-
-        return Uxreplace(subs).visit(body)
-
-    def _make_user_struct_efunc(self):
-        """
-        This is the struct initialised inside the main kernel and
-        attached to the DM via DMSetApplicationContext.
-        """
-        mainctx = self.solver_objs['userctx'] = MainUserStruct(
-            name=self.sregistry.make_name(prefix='ctx'),
-            pname=self.sregistry.make_name(prefix='UserCtx'),
-            fields=self.filtered_struct_params,
-            liveness='lazy',
-            modifier=None
-        )
-        body = [
-            DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol)
-            for i in mainctx.callback_fields
-        ]
-        struct_callback_body = self._make_callable_body(body)
-        cb = Callable(
-            self.sregistry.make_name(prefix='PopulateUserContext'),
-            struct_callback_body, self.objs['err'],
-            parameters=[mainctx]
-        )
-        self._efuncs[cb.name] = cb
-        self._user_struct_efunc = cb
-
-    def _uxreplace_efuncs(self):
-        sobjs = self.solver_objs
-        callback_user_struct = CallbackUserStruct(
-            name=sobjs['userctx'].name,
-            pname=sobjs['userctx'].pname,
-            fields=self.filtered_struct_params,
-            liveness='lazy',
-            modifier=' *',
-            parent=sobjs['userctx']
-        )
-        mapper = {}
-        visitor = Uxreplace({self.objs['dummyctx']: callback_user_struct})
-        for k, v in self._efuncs.items():
-            mapper.update({k: visitor.visit(v)})
-        return mapper
-
-
-class CoupledCallbackBuilder(BaseCallbackBuilder):
-    def __init__(self, **kwargs):
-        self._submatrices_callback = None
-        self._destroy_submat_callback = None
-        super().__init__(**kwargs)
-
-    @property
-    def submatrices_callback(self):
-        return self._submatrices_callback
-
-    def _create_set_point_bc_body(self, body, _constrain_bc_dict):
-        return self._create_set_point_bc_body_coupled(body)
-
-    def _create_set_point_bc_body_coupled(self, body):
         """
-        # TODO : ADD DOCS - MAKE IT CLEARER
-        Combined SetPointBCs body for all target fields.
+        Generic SetPointBCs body, handles single- and multi-field.
         """
+        targets = list(constrain_bc_dict.keys())
+        nfields = len(targets)
         linsolve_expr = self.inject_solve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
-        constrain_bc = self.field_data.constrain_bc
-        targets = self.field_data.targets
-        nfields = len(targets)
         dmda = sobjs['callbackdm']
         ctx = objs['dummyctx']
 
@@ -946,11 +819,63 @@ def _create_set_point_bc_body_coupled(self, body):
                 i in fields if not isinstance(i.function, AbstractFunction)}
         for t in targets:
             tname = t.name
-            subs[constrain_bc[t].counter._C_symbol] = \
+            subs[constrain_bc_dict[t].counter._C_symbol] = \
                 sobjs[f'bcPointsArr_{tname}'].indexed[sobjs[f'k_iter_{tname}']]
 
         return Uxreplace(subs).visit(body)
 
+    def _make_user_struct_efunc(self):
+        """
+        This is the struct initialised inside the main kernel and
+        attached to the DM via DMSetApplicationContext.
+        """
+        mainctx = self.solver_objs['userctx'] = MainUserStruct(
+            name=self.sregistry.make_name(prefix='ctx'),
+            pname=self.sregistry.make_name(prefix='UserCtx'),
+            fields=self.filtered_struct_params,
+            liveness='lazy',
+            modifier=None
+        )
+        body = [
+            DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol)
+            for i in mainctx.callback_fields
+        ]
+        struct_callback_body = self._make_callable_body(body)
+        cb = Callable(
+            self.sregistry.make_name(prefix='PopulateUserContext'),
+            struct_callback_body, self.objs['err'],
+            parameters=[mainctx]
+        )
+        self._efuncs[cb.name] = cb
+        self._user_struct_efunc = cb
+
+    def _uxreplace_efuncs(self):
+        sobjs = self.solver_objs
+        callback_user_struct = CallbackUserStruct(
+            name=sobjs['userctx'].name,
+            pname=sobjs['userctx'].pname,
+            fields=self.filtered_struct_params,
+            liveness='lazy',
+            modifier=' *',
+            parent=sobjs['userctx']
+        )
+        mapper = {}
+        visitor = Uxreplace({self.objs['dummyctx']: callback_user_struct})
+        for k, v in self._efuncs.items():
+            mapper.update({k: visitor.visit(v)})
+        return mapper
+
+
+class CoupledCallbackBuilder(BaseCallbackBuilder):
+    def __init__(self, **kwargs):
+        self._submatrices_callback = None
+        self._destroy_submat_callback = None
+        super().__init__(**kwargs)
+
+    @property
+    def submatrices_callback(self):
+        return self._submatrices_callback
+
     @property
     def jacobian(self):
         return self.inject_solve.expr.rhs.field_data.jacobian
@@ -1175,9 +1100,10 @@ def _whole_formfunc_body(self, body):
         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
-
+        """
+        Each submatrix has a PetscMalloc'd context attached via MatShellSetContext
+        that PETSc's default MatDestroy won't free, so we register a custom destroy.
+        """
         objs = self.objs
 
         get_ctx = petsc_call(

devito/petsc/iet/passes.py

@@ -137,34 +137,12 @@ def lower_petsc(iet, **kwargs):
 
     iet = Transformer(subs).visit(iet)
     body = core + tuple(setup) + iet.body.body + tuple(clear_options)
-    # from IPython import embed; embed()
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
-    # from IPython import embed; embed()
     metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())}
     return iet, metadata
 
 
-@iet_pass
-def strip_petsc_callback_halos(iet, **kwargs):
-    """
-    Remove any HaloSpot nodes that `mpiize` may have injected into PETSc
-    callback functions (FormFunction, SetPointBCs, FormRHS, etc.).
-
-    HaloSpots should only appear in the main kernel.
-    """
-    if not isinstance(iet, PETScCallable):
-        return iet, {}
-
-    halos = FindNodes(HaloSpot).visit(iet)
-    if not halos:
-        return iet, {}
-
-    # Replace each HaloSpot with its body (unwrap it)
-    mapper = {hs: hs.body for hs in halos}
-    return Transformer(mapper).visit(iet), {}
-
-
 def lower_petsc_symbols(iet, **kwargs):
     """
     The `place_definitions` and `place_casts` passes may introduce new
@@ -185,6 +163,10 @@ def lower_petsc_symbols(iet, **kwargs):
 @iet_pass
 def linear_indices(iet, **kwargs):
     """
+    Convert multidimensional grid accesses in the callback `SetPointBCs` to flat
+    linear indices. DMDASetPointBC expects BC points as 1D offsets (e.g. i*ny + j),
+    so each u[x, y] access is linearised to its stride expression and written into
+    bcPointsArr.
     """
     if not iet.name.startswith("SetPointBCs"):
         return iet, {}
@@ -196,9 +178,7 @@ def linear_indices(iet, **kwargs):
 
     tracker = Tracker('basic', dtype, kwargs['sregistry'])
 
-    # TODO: CLEAN UP this is a hack
-    # Exclude SubDomainSet backing functions from linearization - in SETPOINTBCS
-    # I don't want to linearize the accesses to the SubDomainSet
+    # TODO: Rethink - bit of a hack to only linearise the relevant index accesses
     indexeds = [
         i for i in FindSymbols('indexeds').visit(iet)
         if not isinstance(i.function, LocalType)

devito/petsc/iet/type_builder.py

@@ -10,7 +10,7 @@
     PointerIS, PointerDM, VecScatter, JacobianStruct, SubMatrixStruct,
     PetscMPIInt, PetscErrorCode, PointerMat, MatReuse,
     DummyArg, NofSubMats, PetscSectionGlobal,
-    PetscSectionLocal, PetscSF, CallbackPetscInt, PointerPetscInt, SingleIS
+    PetscSectionLocal, PetscSF, CallbackPetscInt, PointerPetscInt
 )
 
 
@@ -223,8 +223,8 @@ def _extend_build(self, base_dict):
         base_dict[f'k_iter_{tname}'] = PostIncrementIndex(
             name='k_iter', initvalue=0
         )
-        base_dict['bcPointsIS'] = SingleIS(name='bcPointsIS')
-        base_dict['bcPoints'] = PointerIS(name='bcPoints')
+        base_dict['bcPointsIS'] = PointerIS(name='bcPointsIS', nindices=1)
+        base_dict['bcCompsIS'] = PointerIS(name='bcCompsIS', nindices=1)
         return base_dict
 
 

devito/petsc/logging.py

@@ -168,9 +168,9 @@ def fields(self):
     def _C_typedecl(self):
         """
         Override generated C struct declaration so that KSPType fields
-        are emitted as ``char name[KSPTYPE_MAX_LEN]`` rather than ``KSPType name``
-        (i.e. ``const char *``).  This avoids a segfault when Python reads the
-        profiler struct after SNESDestroy has freed the KSP.
+        are emitted as `char name[KSPTYPE_MAX_LEN]` rather than `KSPType name`
+        (i.e. `const char *`).  This provides a buffer for PetscStrncpy to copy
+        the type string before SNESDestroy frees the KSP.
         """
         entries = []
         for field in self._fields: