misc: Clean up

Author: ZoeLeibowitz

devito/petsc/iet/routines.py

@@ -16,7 +16,8 @@
 from devito.petsc.types import PETScArray, PetscBundle
 from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback,
                                     MatShellSetOp, PetscMetaData)
-from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector
+from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector,
+                                    dereference_funcs, residual_bundle)
 from devito.petsc.utils import solver_mapper
 from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES,
                                 PetscInt, StartPtr, PointerIS, PointerDM, VecScatter,
@@ -120,7 +121,7 @@ def _make_core(self):
         self._make_user_struct_callback()
 
     def _make_matvec(self, jacobian, prefix='MatMult'):
-        # Compile matvec `eqns` into an IET via recursive compilation
+        # Compile `matvecs` into an IET via recursive compilation
         matvecs = jacobian.matvecs
         irs, _ = self.rcompile(
             matvecs, options={'mpi': False}, sregistry=self.sregistry,
@@ -251,7 +252,7 @@ def _create_matvec_body(self, body, jacobian):
         )
 
         # Dereference function data in struct
-        derefs = self.dereference_funcs(ctx, fields)
+        derefs = dereference_funcs(ctx, fields)
 
         body = CallableBody(
             List(body=body),
@@ -390,7 +391,7 @@ def _create_formfunc_body(self, body):
         )
 
         # Dereference function data in struct
-        derefs = self.dereference_funcs(ctx, fields)
+        derefs = dereference_funcs(ctx, fields)
 
         body = CallableBody(
             List(body=body),
@@ -500,7 +501,7 @@ def _create_form_rhs_body(self, body):
         )
 
         # Dereference function data in struct
-        derefs = self.dereference_funcs(ctx, fields)
+        derefs = dereference_funcs(ctx, fields)
 
         body = CallableBody(
             List(body=[body]),
@@ -578,7 +579,7 @@ def _create_initial_guess_body(self, body):
         )
 
         # Dereference function data in struct
-        derefs = self.dereference_funcs(ctx, fields)
+        derefs = dereference_funcs(ctx, fields)
 
         body = CallableBody(
             List(body=[body]),
@@ -643,12 +644,6 @@ def _uxreplace_efuncs(self):
             mapper.update({k: visitor.visit(v)})
         return mapper
 
-    def dereference_funcs(self, struct, fields):
-        return tuple(
-            [Dereference(i, struct) for i in
-             fields if isinstance(i.function, AbstractFunction)]
-        )
-
 
 class CCBBuilder(CBBuilder):
     def __init__(self, **kwargs):
@@ -749,17 +744,17 @@ def _whole_matvec_body(self):
 
     def _make_whole_formfunc(self):
         F_exprs = self.fielddata.residual.F_exprs
-        # Compile formfunc `eqns` into an IET via recursive compilation
-        irs_formfunc, _ = self.rcompile(
+        # Compile `F_exprs` into an IET via recursive compilation
+        irs, _ = self.rcompile(
             F_exprs, options={'mpi': False}, sregistry=self.sregistry,
             concretize_mapper=self.concretize_mapper
         )
-        body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body))
+        body = self._whole_formfunc_body(List(body=irs.uiet.body))
 
         objs = self.objs
         cb = PETScCallable(
             self.sregistry.make_name(prefix='WholeFormFunc'),
-            body_formfunc,
+            body,
             retval=objs['err'],
             parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
         )
@@ -783,7 +778,8 @@ def _whole_formfunc_body(self, body):
         bundles = sobjs['bundles']
         fbundle = bundles['f']
         xbundle = bundles['x']
-        body = self.residual_bundle(body, bundles)
+
+        body = residual_bundle(body, bundles)
 
         dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True)
 
@@ -870,7 +866,7 @@ def _whole_formfunc_body(self, body):
         )
 
         # Dereference function data in struct
-        derefs = self.dereference_funcs(ctx, fields)
+        derefs = dereference_funcs(ctx, fields)
 
         f_soa = PointerCast(fbundle)
         x_soa = PointerCast(xbundle)
@@ -1034,21 +1030,6 @@ def _submat_callback_body(self):
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
-    def residual_bundle(self, body, bundles):
-        mapper = bundles['bundle_mapper']
-        indexeds = FindSymbols('indexeds').visit(body)
-        subs = {}
-
-        for i in indexeds:
-            if i.base in mapper:
-                bundle = mapper[i.base]
-                index = bundles['target_indices'][i.function.target]
-                index = (index,) + i.indices
-                subs[i] = bundle.__getitem__(index)
-
-        body = Uxreplace(subs).visit(body)
-        return body
-
 
 class BaseObjectBuilder:
     """

devito/petsc/iet/utils.py

@@ -1,5 +1,7 @@
 from devito.petsc.iet.nodes import PetscMetaData, PETScCall
 from devito.ir.equations import OpPetsc
+from devito.ir.iet import Dereference, FindSymbols, Uxreplace
+from devito.types.basic import AbstractFunction
 
 
 def petsc_call(specific_call, call_args):
@@ -19,9 +21,52 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None):
 
 
 def zero_vector(vec):
+    """
+    Set all entries of a PETSc vector to zero.
+    """
     return petsc_call('VecSet', [vec, 0.0])
 
 
+def dereference_funcs(struct, fields):
+    """
+    Dereference AbstractFunctions from a struct.
+    """
+    return tuple(
+        [Dereference(i, struct) for i in
+         fields if isinstance(i.function, AbstractFunction)]
+    )
+
+
+def residual_bundle(body, bundles):
+    """
+    Replaces PetscArrays in `body` with PetscBundle struct field accesses
+    (e.g., f_v[ix][iy] -> f_bundle[ix][iy].v).
+
+    Example:
+        f_v[ix][iy] = x_v[ix][iy];
+        f_u[ix][iy] = x_u[ix][iy];
+    becomes:
+        f_bundle[ix][iy].v = x_bundle[ix][iy].v;
+        f_bundle[ix][iy].u = x_bundle[ix][iy].u;
+
+    NOTE: This is used because the data is interleaved for
+    multi-component DMDAs in PETSc.
+    """
+    mapper = bundles['bundle_mapper']
+    indexeds = FindSymbols('indexeds').visit(body)
+    subs = {}
+
+    for i in indexeds:
+        if i.base in mapper:
+            bundle = mapper[i.base]
+            index = bundles['target_indices'][i.function.target]
+            index = (index,) + i.indices
+            subs[i] = bundle.__getitem__(index)
+
+    body = Uxreplace(subs).visit(body)
+    return body
+
+
 # Mapping special Eq operations to their corresponding IET Expression subclass types.
 # These operations correspond to subclasses of Eq utilised within PETScSolve.
 petsc_iet_mapper = {OpPetsc: PetscMetaData}