misc: Drop petsc/iet/utils.py

Author: ZoeLeibowitz

devito/petsc/iet/callback_builder.py

@@ -3,18 +3,17 @@
 from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody,
                            Dereference, DummyExpr, BlankLine, Callable, Iteration,
                            PointerCast, Definition)
-from devito.symbolics import (Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null)
+from devito.symbolics import (Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID)
 from devito.symbolics.unevaluation import Mul
 from devito.types.basic import AbstractFunction
-from devito.types import Dimension
+from devito.types import Dimension, Temp, TempArray
 from devito.tools import filter_ordered
 
-from devito.petsc.iet.nodes import PETScCallable, MatShellSetOp
-from devito.petsc.iet.utils import (petsc_call, void, get_user_struct_fields)
+from devito.petsc.iet.nodes import PETScCallable, MatShellSetOp, petsc_call
 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
+from devito.petsc.types.modes import InsertMode
 
 
 class BaseCallback:
@@ -226,12 +225,11 @@ def _create_matvec_body(self, body, jacobian):
         )
 
         global_to_local_begin = petsc_call(
-            'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     InsertMode.insert_values, xlocal]
+            'DMGlobalToLocalBegin', [dmda, objs['X'], insert_values, xlocal]
         )
 
         global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], InsertMode.insert_values, xlocal
+            dmda, objs['X'], insert_values, xlocal
         ])
 
         dm_get_local_yvec = petsc_call(
@@ -261,11 +259,11 @@ def _create_matvec_body(self, body, jacobian):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, ylocal, InsertMode.add_values, objs['Y']
+            dmda, ylocal, add_values, objs['Y']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, ylocal, InsertMode.add_values, objs['Y']
+            dmda, ylocal, add_values, objs['Y']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -373,13 +371,12 @@ def _create_formfunc_body(self, body):
         )
 
         global_to_local_begin = petsc_call(
-            'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     InsertMode.insert_values, objs['xloc']]
+            'DMGlobalToLocalBegin', [dmda, objs['X'], insert_values, objs['xloc']]
         )
 
-        global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], InsertMode.insert_values, objs['xloc']
-        ])
+        global_to_local_end = petsc_call(
+            'DMGlobalToLocalEnd', [dmda, objs['X'], insert_values, objs['xloc']]
+        )
 
         dm_get_local_yvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(objs['floc'])]
@@ -406,11 +403,11 @@ def _create_formfunc_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, objs['floc'], InsertMode.add_values, objs['F']
+            dmda, objs['floc'], add_values, objs['F']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, objs['floc'], InsertMode.add_values, objs['F']
+            dmda, objs['floc'], add_values, objs['F']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -490,14 +487,12 @@ def _create_form_rhs_body(self, body):
         )
 
         dm_global_to_local_begin = petsc_call(
-            'DMGlobalToLocalBegin', [dmda, objs['B'],
-                                     InsertMode.insert_values, sobjs['blocal']]
+            'DMGlobalToLocalBegin', [dmda, objs['B'], insert_values, sobjs['blocal']]
         )
 
-        dm_global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['B'], InsertMode.insert_values,
-            sobjs['blocal']
-        ])
+        dm_global_to_local_end = petsc_call(
+            'DMGlobalToLocalEnd', [dmda, objs['B'], insert_values, sobjs['blocal']]
+        )
 
         b_arr = self.field_data.arrays[target]['b']
 
@@ -519,13 +514,11 @@ def _create_form_rhs_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, sobjs['blocal'], InsertMode.insert_values,
-            objs['B']
+            dmda, sobjs['blocal'], insert_values, objs['B']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, sobjs['blocal'], InsertMode.insert_values,
-            objs['B']
+            dmda, sobjs['blocal'], insert_values, objs['B']
         ])
 
         vec_restore_array = petsc_call(
@@ -822,13 +815,12 @@ def _whole_formfunc_body(self, body):
             'DMGetLocalVector', [dmda, Byref(objs['xloc'])]
         )
 
-        global_to_local_begin = petsc_call(
-            'DMGlobalToLocalBegin', [dmda, objs['X'],
-                                     InsertMode.insert_values, objs['xloc']]
-        )
+        global_to_local_begin = petsc_call('DMGlobalToLocalBegin', [
+            dmda, objs['X'], insert_values, objs['xloc']
+        ])
 
         global_to_local_end = petsc_call('DMGlobalToLocalEnd', [
-            dmda, objs['X'], InsertMode.insert_values, objs['xloc']
+            dmda, objs['X'], insert_values, objs['xloc']
         ])
 
         dm_get_local_yvec = petsc_call(
@@ -856,11 +848,11 @@ def _whole_formfunc_body(self, body):
         )
 
         dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [
-            dmda, objs['floc'], InsertMode.add_values, objs['F']
+            dmda, objs['floc'], add_values, objs['F']
         ])
 
         dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [
-            dmda, objs['floc'], InsertMode.add_values, objs['F']
+            dmda, objs['floc'], add_values, objs['F']
         ])
 
         dm_restore_local_xvec = petsc_call(
@@ -1033,7 +1025,7 @@ def _submat_callback_body(self):
                 [
                     objs['submat_arr'].indexed[sb.linear_idx],
                     'MATOP_MULT',
-                    MatShellSetOp(matvec_lookup[sb.name].name, void, void),
+                    MatShellSetOp(matvec_lookup[sb.name].name, VOID._dtype, VOID._dtype),
                 ],
             )
             for sb in nonzero_submats if sb.name in matvec_lookup
@@ -1120,3 +1112,18 @@ def zero_vector(vec):
     Set all entries of a PETSc vector to zero.
     """
     return petsc_call('VecSet', [vec, 0.0])
+
+
+def get_user_struct_fields(iet):
+    fields = [f.function for f in FindSymbols('basics').visit(iet)]
+    from devito.types.basic import LocalType
+    avoid = (Temp, TempArray, LocalType)
+    fields = [f for f in fields if not isinstance(f.function, avoid)]
+    fields = [
+        f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo))
+    ]
+    return fields
+
+
+insert_values = InsertMode.insert_values
+add_values = InsertMode.add_values

devito/petsc/iet/logging.py

@@ -5,7 +5,7 @@
 from devito.logger import PERF
 from devito.tools import frozendict
 
-from devito.petsc.iet.utils import petsc_call
+from devito.petsc.iet.nodes import petsc_call
 from devito.petsc.logging import petsc_return_variable_dict, PetscInfo
 
 

devito/petsc/iet/nodes.py

@@ -31,6 +31,10 @@ class PETScCall(Call):
     pass
 
 
+def petsc_call(specific_call, call_args):
+    return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)])
+
+
 # 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/passes.py

@@ -16,17 +16,16 @@
     CallbackUserStruct
 )
 from devito.petsc.types.macros import petsc_func_begin_user
-from devito.petsc.iet.nodes import PetscMetaData
+from devito.petsc.iet.nodes import PetscMetaData, petsc_call
 from devito.petsc.utils import core_metadata, petsc_languages
 from devito.petsc.iet.callback_builder import (
-    BaseCallback, CoupledCallback, populate_matrix_context
+    BaseCallback, CoupledCallback, populate_matrix_context, get_user_struct_fields
 )
 from devito.petsc.iet.object_builder import BaseObjectBuilder, CoupledObjectBuilder, objs
 from devito.petsc.iet.pre_solver import BaseSetup, CoupledSetup, make_core_petsc_calls
 from devito.petsc.iet.run_solver import Solver, CoupledSolver
 from devito.petsc.iet.time_dependence import TimeDependent, TimeIndependent
 from devito.petsc.iet.logging import PetscLogger
-from devito.petsc.iet.utils import petsc_call, get_user_struct_fields
 
 
 @iet_pass

devito/petsc/iet/pre_solver.py

@@ -4,8 +4,9 @@
 from devito.symbolics import (Byref, FieldFromPointer, VOID,
                               FieldFromComposite, Null)
 
-from devito.petsc.iet.nodes import FormFunctionCallback, MatShellSetOp, PETScCall
-from devito.petsc.iet.utils import petsc_call
+from devito.petsc.iet.nodes import (
+    FormFunctionCallback, MatShellSetOp, PETScCall, petsc_call
+)
 
 
 def make_core_petsc_calls(objs, comm):
@@ -82,13 +83,11 @@ def _setup(self):
                                   [sobjs['snes'], Byref(sobjs['ksp'])])
 
         matvec = self.cbbuilder.main_matvec_callback
-        matvec_operation = petsc_call(
-            'MatShellSetOperation',
+        matvec_operation = petsc_call('MatShellSetOperation',
             [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)]
         )
         formfunc = self.cbbuilder._F_efunc
-        formfunc_operation = petsc_call(
-            'SNESSetFunction',
+        formfunc_operation = petsc_call('SNESSetFunction',
             [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void),
              self.snes_ctx]
         )
@@ -337,4 +336,4 @@ def petsc_call_mpi(specific_call, call_args):
     return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)])
 
 
-void = 'void'
+void = VOID._dtype

devito/petsc/iet/run_solver.py

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

devito/petsc/iet/time_dependence.py

@@ -3,7 +3,8 @@
 from devito.ir.iet import Uxreplace, DummyExpr
 from devito.symbolics import FieldFromPointer, cast, FieldFromComposite
 from devito.symbolics.unevaluation import Mul
-from devito.petsc.iet.utils import petsc_call
+
+from devito.petsc.iet.nodes import petsc_call
 
 
 class TimeIndependent:

devito/petsc/iet/utils.py

@@ -14,4 +14,4 @@ class ScatterMode:
     Reference - https://petsc.org/release/manualpages/Vec/ScatterMode/
     """
     scatter_reverse = 'SCATTER_REVERSE'
-    scatter_forward = 'SCATTER_FORWARD'
+    scatter_forward = 'SCATTER_FORWARD'

devito/petsc/types/modes.py

@@ -6,7 +6,7 @@
 from devito.symbolics import Byref, cast
 from devito.types.basic import DataSymbol, LocalType
 
-from devito.petsc.iet.utils import petsc_call
+from devito.petsc.iet.nodes import petsc_call
 
 
 class PetscMixin: