misc: Docs

Author: ZoeLeibowitz

devito/petsc/equations.py

@@ -19,6 +19,11 @@ def lower_exprs_petsc(expressions, **kwargs):
 
 def constrain_essential_bcs(expressions, **kwargs):
     """
+    Expand loop bounds for `ConstrainBC` expressions so that each MPI rank
+    iterates over all locally visible constrained points, including those in
+    the halo. PETSc requires each rank to report all constrained nodes in its
+    local data region. The loops are not used for data access — only to
+    identify which local indices are constrained.
     """
     constrain_expressions = [e for e in expressions if isinstance(e, ConstrainBC)]
     if not constrain_expressions:

devito/petsc/iet/builder.py

@@ -252,12 +252,6 @@ def _setup(self):
 
         dmda_calls = self._create_dmda_calls(dmda)
 
-        mainctx = sobjs['userctx']
-
-        call_struct_callback = petsc_call(
-            self.callback_builder.user_struct_efunc.name, [Byref(mainctx)]
-        )
-
         # TODO: maybe don't need to explictly set this
         mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda])
 
@@ -331,7 +325,6 @@ def _setup(self):
             matvec_operation,
             formfunc_operation,
             snes_set_options,
-            call_struct_callback,
             mat_set_dm,
             create_field_decomp,
             matop_create_submats_op,
@@ -420,7 +413,6 @@ class ConstrainedBCBuilder(ConstrainedBCMixin, BuilderBase):
     pass
 
 
-# TODO: Implement this class properly
 class CoupledConstrainedBCBuilder(ConstrainedBCMixin, CoupledBuilder):
     pass
 

devito/petsc/iet/callbacks.py

@@ -697,7 +697,7 @@ def _make_constrain_bc(self):
 
     def _create_count_bc_body(self, body, pairs):
         """
-        Generic CountBCs body. `pairs` is a list of (numBCPtr, counter) tuples,
+        Build the `CountBCs` body. `pairs` is a list of (numBCPtr, counter) tuples,
         one per field. All fields are handled in a single callback body.
         """
         objs = self.objs
@@ -737,7 +737,7 @@ def _create_count_bc_body(self, body, pairs):
 
     def _create_set_point_bc_body(self, body, constrain_bc_dict):
         """
-        Generic SetPointBCs body, handles single- and multi-field.
+        Build the `SetPointBCs` body.
         """
         targets = list(constrain_bc_dict.keys())
         nfields = len(targets)

devito/petsc/iet/passes.py

@@ -70,13 +70,12 @@ def lower_petsc(iet, **kwargs):
         )
     grid = unique_grids.pop()
 
-    # Protect PETSc solve targets from being dropped by `_drop_if_unwritten`.
     # `lower_petsc` runs before `mpiize`, replacing `PetscMetaData` (an
     # `Expression` subclass whose `.write` reveals the target function) with
     # `Call` nodes to run the PETSc solver.  Once that happens,
     # `_drop_if_unwritten` can no longer see the target as written and
     # incorrectly discards its `HaloSpot`. So we compose `dist-drop-unwritten`
-    # with a guard that always returns False for PETSc targets.
+    # with a guard that always returns False for `petsc_targets`.
     options = kwargs['options']
     petsc_targets = {n.write for n in data if n.write is not None}
     if petsc_targets:
@@ -164,9 +163,7 @@ def lower_petsc_symbols(iet, **kwargs):
 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.
+    linear indices. DMDASetPointBC expects BC points as 1D offsets (e.g. i*ny + j).
     """
     if not iet.name.startswith("SetPointBCs"):
         return iet, {}

devito/petsc/types/dimension.py

@@ -1,99 +1,73 @@
 from devito.types.dimension import Thickness
 
 
-class SubDimMax(Thickness):
-    """
-    Local index of a SubDimension's right global boundary, which may not correspond
-    to a locally owned point. Not used for indexing into data.
-    """
+class _SubDimBound(Thickness):
     def __init_finalize__(self, *args, **kwargs):
         self._subdim = kwargs.pop('subdim')
         self._dtype = self._subdim.dtype
-
         super().__init_finalize__(*args, **kwargs)
 
     @property
     def subdim(self):
         return self._subdim
 
+
+class _SpaceDimBound(Thickness):
+    def __init_finalize__(self, *args, **kwargs):
+        self._space_dim = kwargs.pop('space_dim')
+        self._dtype = self._space_dim.dtype
+        super().__init_finalize__(*args, **kwargs)
+
+    @property
+    def space_dim(self):
+        return self._space_dim
+
+
+class SubDimMax(_SubDimBound):
+    """
+    Local index of a SubDimension's right global boundary, which may not correspond
+    to a locally owned point. Not used for indexing into data.
+    """
     def _arg_values(self, grid=None, **kwargs):
         dist = grid.distributor
-        # global rtkn
         grtkn = kwargs.get(self.subdim.rtkn.name, self.subdim.rtkn.value)
-        # decomposition info
         decomp = dist.decomposition[self.subdim.parent]
-        glb_max = decomp.glb_max
-        val = decomp.index_glb_to_loc_unsafe(glb_max - grtkn)
+        val = decomp.index_glb_to_loc_unsafe(decomp.glb_max - grtkn)
         return {self.name: int(val)}
 
 
-class SubDimMin(Thickness):
+class SubDimMin(_SubDimBound):
     """
     Local index of a SubDimension's left global boundary, which may not correspond
     to a locally owned point. Not used for indexing into data.
     """
-    def __init_finalize__(self, *args, **kwargs):
-        self._subdim = kwargs.pop('subdim')
-        self._dtype = self._subdim.dtype
-
-        super().__init_finalize__(*args, **kwargs)
-
-    @property
-    def subdim(self):
-        return self._subdim
-
     def _arg_values(self, grid=None, **kwargs):
         dist = grid.distributor
-        # global ltkn
         gltkn = kwargs.get(self.subdim.ltkn.name, self.subdim.ltkn.value)
-        # decomposition info
         decomp = dist.decomposition[self.subdim.parent]
-        glb_min = decomp.glb_min
-        val = decomp.index_glb_to_loc_unsafe(glb_min + gltkn)
+        val = decomp.index_glb_to_loc_unsafe(decomp.glb_min + gltkn)
         return {self.name: int(val)}
 
 
-class SpaceDimMax(Thickness):
+class SpaceDimMax(_SpaceDimBound):
     """
     Local index of a SpaceDimension's right global boundary, which may not correspond
     to a locally owned point. Not used for indexing into data.
     """
-    def __init_finalize__(self, *args, **kwargs):
-        self._space_dim = kwargs.pop('space_dim')
-        self._dtype = self._space_dim.dtype
-
-        super().__init_finalize__(*args, **kwargs)
-
-    @property
-    def space_dim(self):
-        return self._space_dim
-
     def _arg_values(self, grid=None, **kwargs):
         dist = grid.distributor
         decomp = dist.decomposition[self.space_dim]
-        glb_max = decomp.glb_max
-        val = decomp.index_glb_to_loc_unsafe(glb_max)
+        val = decomp.index_glb_to_loc_unsafe(decomp.glb_max)
         return {self.name: int(val)}
 
 
-class SpaceDimMin(Thickness):
+class SpaceDimMin(_SpaceDimBound):
     """
     Local index of a SpaceDimension's left global boundary, which may not correspond
     to a locally owned point. Not used for indexing into data.
     """
-    def __init_finalize__(self, *args, **kwargs):
-        self._space_dim = kwargs.pop('space_dim')
-        self._dtype = self._space_dim.dtype
-
-        super().__init_finalize__(*args, **kwargs)
-
-    @property
-    def space_dim(self):
-        return self._space_dim
-
     def _arg_values(self, grid=None, **kwargs):
         dist = grid.distributor
         decomp = dist.decomposition[self.space_dim]
-        glb_min = decomp.glb_min
-        val = decomp.index_glb_to_loc_unsafe(glb_min)
+        val = decomp.index_glb_to_loc_unsafe(decomp.glb_min)
         return {self.name: int(val)}