misc: Separate more files

Author: ZoeLeibowitz

devito/ir/iet/algorithms.py

@@ -4,7 +4,7 @@
                            Section, HaloSpot, ExpressionBundle)
 from devito.tools import timed_pass
 from devito.petsc.types import MetaData
-from devito.petsc.iet.utils import petsc_iet_mapper
+from devito.petsc.iet.nodes import petsc_iet_mapper
 
 __all__ = ['iet_build']
 

devito/petsc/iet/callback_builder.py

@@ -10,12 +10,11 @@
 from devito.tools import filter_ordered
 
 from devito.petsc.iet.nodes import PETScCallable, MatShellSetOp
-from devito.petsc.iet.utils import (petsc_call, zero_vector,
-                                    dereference_funcs, void,
-                                    insert_vals, add_vals, get_user_struct_fields)
+from devito.petsc.iet.utils import (petsc_call, void, get_user_struct_fields)
 from devito.petsc.types import DMCast, MainUserStruct, CallbackUserStruct
 from devito.petsc.iet.object_builder import objs
 from devito.petsc.types.macros import petsc_func_begin_user
+from devito.petsc.types.strings import InsertMode
 
 
 class BaseCallback:
@@ -228,11 +227,11 @@ def _create_matvec_body(self, body, jacobian):
 
         global_to_local_begin = petsc_call(
             'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     insert_vals, xlocal]
+                                     InsertMode.insert_values, xlocal]
         )
 
         global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], insert_vals, xlocal
+            dmda, objs['X'], InsertMode.insert_values, xlocal
         ])
 
         dm_get_local_yvec = petsc_call(
@@ -262,11 +261,11 @@ def _create_matvec_body(self, body, jacobian):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, ylocal, add_vals, objs['Y']
+            dmda, ylocal, InsertMode.add_values, objs['Y']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, ylocal, add_vals, objs['Y']
+            dmda, ylocal, InsertMode.add_values, objs['Y']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -375,11 +374,11 @@ def _create_formfunc_body(self, body):
 
         global_to_local_begin = petsc_call(
             'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     insert_vals, objs['xloc']]
+                                     InsertMode.insert_values, objs['xloc']]
         )
 
         global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], insert_vals, objs['xloc']
+            dmda, objs['X'], InsertMode.insert_values, objs['xloc']
         ])
 
         dm_get_local_yvec = petsc_call(
@@ -407,11 +406,11 @@ def _create_formfunc_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, objs['floc'], add_vals, objs['F']
+            dmda, objs['floc'], InsertMode.add_values, objs['F']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, objs['floc'], add_vals, objs['F']
+            dmda, objs['floc'], InsertMode.add_values, objs['F']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -492,11 +491,11 @@ def _create_form_rhs_body(self, body):
 
         dm_global_to_local_begin = petsc_call(
             'DMGlobalToLocalBegin', [dmda, objs['B'],
-                                     insert_vals, sobjs['blocal']]
+                                     InsertMode.insert_values, sobjs['blocal']]
         )
 
         dm_global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['B'], insert_vals,
+            dmda, objs['B'], InsertMode.insert_values,
             sobjs['blocal']
         ])
 
@@ -520,12 +519,12 @@ def _create_form_rhs_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, sobjs['blocal'], insert_vals,
+            dmda, sobjs['blocal'], InsertMode.insert_values,
             objs['B']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, sobjs['blocal'], insert_vals,
+            dmda, sobjs['blocal'], InsertMode.insert_values,
             objs['B']
         ])
 
@@ -825,11 +824,11 @@ def _whole_formfunc_body(self, body):
 
         global_to_local_begin = petsc_call(
             'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     insert_vals, objs['xloc']]
+                                     InsertMode.insert_values, objs['xloc']]
         )
 
         global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], insert_vals, objs['xloc']
+            dmda, objs['X'], InsertMode.insert_values, objs['xloc']
         ])
 
         dm_get_local_yvec = petsc_call(
@@ -857,11 +856,11 @@ def _whole_formfunc_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, objs['floc'], add_vals, objs['F']
+            dmda, objs['floc'], InsertMode.add_values, objs['F']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, objs['floc'], add_vals, objs['F']
+            dmda, objs['floc'], InsertMode.add_values, objs['F']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -1104,3 +1103,20 @@ def populate_matrix_context(efuncs):
         name, body, objs['err'],
         parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']]
     )
+
+
+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 zero_vector(vec):
+    """
+    Set all entries of a PETSc vector to zero.
+    """
+    return petsc_call('VecSet', [vec, 0.0])

devito/petsc/iet/nodes.py

@@ -29,3 +29,8 @@ def callback_form(self):
 
 class PETScCall(Call):
     pass
+
+
+# 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}

devito/petsc/iet/pre_solver.py

@@ -4,8 +4,8 @@
 from devito.symbolics import (Byref, FieldFromPointer, VOID,
                               FieldFromComposite, Null)
 
-from devito.petsc.iet.nodes import FormFunctionCallback, MatShellSetOp
-from devito.petsc.iet.utils import petsc_call, void, petsc_call_mpi
+from devito.petsc.iet.nodes import FormFunctionCallback, MatShellSetOp, PETScCall
+from devito.petsc.iet.utils import petsc_call
 
 
 def make_core_petsc_calls(objs, comm):
@@ -331,3 +331,10 @@ def _setup(self):
             create_submats) + \
             tuple(deref_dms) + tuple(xglobals) + tuple(xlocals) + (BlankLine,)
         return coupled_setup
+
+
+def petsc_call_mpi(specific_call, call_args):
+    return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)])
+
+
+void = 'void'

devito/petsc/iet/run_solver.py

@@ -6,7 +6,8 @@
 from devito.symbolics import (Byref, Null)
 
 from devito.petsc.iet.nodes import PetscMetaData
-from devito.petsc.iet.utils import petsc_call, insert_vals, sreverse, sforward
+from devito.petsc.types.strings import InsertMode, ScatterMode
+from devito.petsc.iet.utils import petsc_call
 
 
 class Solver:
@@ -44,7 +45,7 @@ def _execute_solve(self):
             initguess_call = None
 
         dm_local_to_global_x = petsc_call(
-            'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals,
+            'DMLocalToGlobal', [dmda, sobjs['xlocal'], InsertMode.insert_values,
                                 sobjs['xglobal']]
         )
 
@@ -53,7 +54,7 @@ def _execute_solve(self):
         )
 
         dm_global_to_local_x = petsc_call('DMGlobalToLocal', [
-            dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']]
+            dmda, sobjs['xglobal'], InsertMode.insert_values, sobjs['xlocal']]
         )
 
         vec_reset_array = self.time_dependence.reset_array(target)
@@ -113,35 +114,35 @@ def _execute_solve(self):
                 self.time_dependence.place_array(t),
                 petsc_call(
                     'DMLocalToGlobal',
-                    [dm, target_xloc, insert_vals, target_xglob]
+                    [dm, target_xloc, InsertMode.insert_values, target_xglob]
                 ),
                 petsc_call(
                     'VecScatterCreate',
                     [xglob, field, target_xglob, Null, Byref(s)]
                 ),
                 petsc_call(
                     'VecScatterBegin',
-                    [s, target_xglob, xglob, insert_vals, sreverse]
+                    [s, target_xglob, xglob, InsertMode.insert_values, ScatterMode.scatter_reverse]
                 ),
                 petsc_call(
                     'VecScatterEnd',
-                    [s, target_xglob, xglob, insert_vals, sreverse]
+                    [s, target_xglob, xglob, InsertMode.insert_values, ScatterMode.scatter_reverse]
                 ),
                 BlankLine,
             )
 
             post_solve += (
                 petsc_call(
                     'VecScatterBegin',
-                    [s, xglob, target_xglob, insert_vals, sforward]
+                    [s, xglob, target_xglob, InsertMode.insert_values, ScatterMode.scatter_forward]
                 ),
                 petsc_call(
                     'VecScatterEnd',
-                    [s, xglob, target_xglob, insert_vals, sforward]
+                    [s, xglob, target_xglob, InsertMode.insert_values, ScatterMode.scatter_forward]
                 ),
                 petsc_call(
                     'DMGlobalToLocal',
-                    [dm, target_xglob, insert_vals, target_xloc]
+                    [dm, target_xglob, InsertMode.insert_values, target_xloc]
                 )
             )
 

devito/petsc/iet/utils.py

@@ -1,36 +1,13 @@
-from devito.ir.equations import OpPetsc
-from devito.ir.iet import Dereference, FindSymbols
-from devito.types.basic import AbstractFunction
+from devito.ir.iet import FindSymbols
 from devito.types import Temp, TempArray
 
-from devito.petsc.iet.nodes import PetscMetaData, PETScCall
+from devito.petsc.iet.nodes import PETScCall
 
 
 def petsc_call(specific_call, call_args):
     return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)])
 
 
-def petsc_call_mpi(specific_call, call_args):
-    return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)])
-
-
-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 get_user_struct_fields(iet):
     fields = [f.function for f in FindSymbols('basics').visit(iet)]
     from devito.types.basic import LocalType
@@ -40,15 +17,3 @@ def get_user_struct_fields(iet):
         f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo))
     ]
     return fields
-
-
-# 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}
-
-
-void = 'void'
-insert_vals = 'INSERT_VALUES'
-add_vals = 'ADD_VALUES'
-sreverse = 'SCATTER_REVERSE'
-sforward = 'SCATTER_FORWARD'

devito/petsc/types/__init__.py

@@ -1,5 +1,6 @@
 from .array import *  # noqa
-from .types import *  # noqa
+from .metadata import *  # noqa
 from .object import *  # noqa
 from .equation import *  # noqa
 from .macros import *  # noqa
+from .modes import *  # noqa

devito/petsc/types/macros.py

@@ -1,5 +1,4 @@
 import cgen as c
 
-
 # TODO: Don't use c.Line here?
 petsc_func_begin_user = c.Line('PetscFunctionBeginUser;')

devito/petsc/types/types.py devito/petsc/types/metadata.pydevito/petsc/types/types.py renamed to devito/petsc/types/metadata.py

@@ -0,0 +1,17 @@
+
+class InsertMode:
+    """
+    How the entries are combined with the current values in the vectors or matrices.
+    Reference - https://petsc.org/main/manualpages/Sys/InsertMode/
+    """
+    insert_values = 'INSERT_VALUES'
+    add_values = 'ADD_VALUES'
+
+
+class ScatterMode:
+    """
+    Determines the direction of a scatter in `VecScatterBegin()` and `VecScatterEnd()`.
+    Reference - https://petsc.org/release/manualpages/Vec/ScatterMode/
+    """
+    scatter_reverse = 'SCATTER_REVERSE'
+    scatter_forward = 'SCATTER_FORWARD'