dsl/compiler: Re-factor solver params, add solver_parameters.py file

Author: ZoeLeibowitz

devito/finite_differences/finite_difference.py

@@ -9,7 +9,7 @@
 
 # Number of digits for FD coefficients to avoid roundup errors and non-deterministic
 # code generation
-_PRECISION = 18
+_PRECISION = 9
 
 
 @check_input

devito/petsc/clusters.py

@@ -1,5 +1,5 @@
 from devito.tools import timed_pass
-from devito.petsc.types import LinearSolveExpr
+from devito.petsc.types import SolveExpr
 
 
 @timed_pass()
@@ -19,7 +19,7 @@ def petsc_lift(clusters):
     """
     processed = []
     for c in clusters:
-        if isinstance(c.exprs[0].rhs, LinearSolveExpr):
+        if isinstance(c.exprs[0].rhs, SolveExpr):
             ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions)
             processed.append(c.rebuild(ispace=ispace))
         else:

devito/petsc/iet/passes.py

@@ -71,10 +71,8 @@ def lower_petsc(iet, **kwargs):
     # Map PETScSolve to its Section (for logging)
     section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
 
-    # Utility callback for setting PETSc options.
-    # This generic function can be reused across all PETScSolves.
+    # Callback used to set a PetscOption - used by all PETScSolves
     set_solver_option(efuncs)
-    # from IPython import embed; embed()
 
     for iters, (inject_solve,) in inject_solve_mapper.items():
 
@@ -94,7 +92,6 @@ def lower_petsc(iet, **kwargs):
     body = core + tuple(setup) + iet.body.body
     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
 
@@ -228,7 +225,7 @@ def set_solver_option(efuncs):
     option = ConstCharPtr(name='option', is_const=True)
     value = ConstCharPtr(name='value', is_const=True)
     set = PetscBool(name='set')
-    
+
     body = List(body=[
         petsc_call('PetscOptionsHasName', [Null, Null, option, Byref(set)]),
         Conditional(Not(set), petsc_call('PetscOptionsSetValue', [Null, option, value]))
@@ -246,7 +243,6 @@ def set_solver_option(efuncs):
         retval=objs['err'],
         parameters=(option, value)
     )
-    # from IPython import embed; embed()
     efuncs[cb.name] = cb
 
 

devito/petsc/iet/routines.py

@@ -17,11 +17,10 @@
                                     MatShellSetOp, PetscMetaData)
 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 (PETScArray, PetscBundle, DM, Mat, CallbackVec, Vec,
                                 KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM,
                                 VecScatter, DMCast, JacobianStruct, SubMatrixStruct,
-                                CallbackDM, PetscBool)
+                                CallbackDM)
 from devito.petsc.types.macros import petsc_func_begin_user, Null
 
 
@@ -119,14 +118,14 @@ def _make_core(self):
     def _petsc_options_callback(self):
         objs = self.objs
         params = self.solver_parameters
-        options_prefix = self.inject_solve.expr.rhs.options_prefix
+        prefix = self.inject_solve.expr.rhs.formatted_prefix
 
-        body = []
-
-        # from IPython import embed; embed()
-        # TODO: improve
-        for k, v in params.items():
-            body.append(petsc_call('SetPetscOption', [String("-"+options_prefix+k), String(v)]))
+        body = [
+            petsc_call(
+                'SetPetscOption', [String(f"-{prefix}{k}"), String(str(v))]
+            )
+            for k, v in params.items()
+        ]
 
         body = CallableBody(
             List(body=body),
@@ -695,6 +694,7 @@ def _make_core(self):
         for sm in self.field_data.jacobian.nonzero_submatrices:
             self._make_matvec(sm, prefix=f'{sm.name}_MatMult')
 
+        self._petsc_options_callback()
         self._make_whole_matvec()
         self._make_whole_formfunc()
         self._make_user_struct_callback()
@@ -1089,7 +1089,7 @@ def _build(self):
         targets = self.field_data.targets
 
         snes_name = sreg.make_name(prefix='snes')
-        options_prefix = self.inject_solve.expr.rhs.options_prefix
+        formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix
 
         base_dict = {
             'Jac': Mat(sreg.make_name(prefix='J')),
@@ -1103,9 +1103,9 @@ def _build(self):
             'localsize': PetscInt(sreg.make_name(prefix='localsize')),
             'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)),
             'callbackdm': CallbackDM(sreg.make_name(prefix='dm')),
-            'snesprefix': String(options_prefix or ''),
-            'options_prefix': options_prefix,
+            'snes_prefix': String(formatted_prefix),
         }
+
         base_dict['comm'] = self.comm
         self._target_dependent(base_dict)
         return self._extend_build(base_dict)
@@ -1244,6 +1244,7 @@ def __init__(self, **kwargs):
         self.solver_objs = kwargs.get('solver_objs')
         self.cbbuilder = kwargs.get('cbbuilder')
         self.field_data = self.inject_solve.expr.rhs.field_data
+        self.formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix
         self.calls = self._setup()
 
     @property
@@ -1255,18 +1256,14 @@ def snes_ctx(self):
         return VOID(self.solver_objs['dmda'], stars='*')
 
     def _setup(self):
-        objs = self.objs
         sobjs = self.solver_objs
-
         dmda = sobjs['dmda']
 
-        # solver_params = self.inject_solve.expr.rhs.solver_parameters
-
         snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])])
 
         snes_options_prefix = petsc_call(
-            'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']]
-        ) if sobjs['options_prefix'] else None
+            'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']]
+        ) if self.formatted_prefix else None
 
         set_options = petsc_call(
             self.cbbuilder._options_efunc.name, []
@@ -1276,9 +1273,6 @@ def _setup(self):
 
         create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])])
 
-        # NOTE: Assuming all solves are linear for now
-        snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY'])
-
         snes_set_jac = petsc_call(
             'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'],
                                 sobjs['Jac'], 'MatMFFDComputeJacobian', Null]
@@ -1295,6 +1289,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, maybe this should be VecCreateSeqWithArray
         local_x = petsc_call('VecCreateMPIWithArray',
                              [sobjs['comm'], 1, local_size, 'PETSC_DECIDE',
                               field_from_ptr, Byref(sobjs['xlocal'])])
@@ -1310,26 +1305,6 @@ def _setup(self):
         snes_get_ksp = petsc_call('SNESGetKSP',
                                   [sobjs['snes'], Byref(sobjs['ksp'])])
 
-        # ksp_set_tols = petsc_call(
-        #     'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'],
-        #                          solver_params['ksp_atol'], solver_params['ksp_divtol'],
-        #                          solver_params['ksp_max_it']]
-        # )
-
-        # ksp_set_type = petsc_call(
-        #     'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]]
-        # )
-
-        # TODO: can drop this
-        ksp_get_pc = petsc_call(
-            'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])]
-        )
-
-        # Even though the default will be jacobi, set to PCNONE for now
-        pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE'])
-
-        ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']])
-
         matvec = self.cbbuilder.main_matvec_callback
         matvec_operation = petsc_call(
             'MatShellSetOperation',
@@ -1366,16 +1341,11 @@ def _setup(self):
             snes_set_dm,
             create_matrix,
             snes_set_jac,
-            snes_set_type,
             global_x,
             local_x,
             get_local_size,
             global_b,
             snes_get_ksp,
-            # ksp_set_tols,
-            ksp_get_pc,
-            pc_set_type,
-            ksp_set_from_ops,
             matvec_operation,
             formfunc_operation,
             snes_set_options,
@@ -1400,7 +1370,6 @@ def _create_dmda_calls(self, dmda):
         return dmda_create, dm_setup, dm_mat_type
 
     def _create_dmda(self, dmda):
-        objs = self.objs
         sobjs = self.solver_objs
         grid = self.field_data.grid
         nspace_dims = len(grid.dimensions)
@@ -1445,23 +1414,22 @@ def _setup(self):
         # TODO: minimise code duplication with superclass
         objs = self.objs
         sobjs = self.solver_objs
-
         dmda = sobjs['dmda']
-        solver_params = self.inject_solve.expr.rhs.solver_parameters
 
         snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])])
 
         snes_options_prefix = petsc_call(
-            'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']]
-        ) if sobjs['options_prefix'] else None
+            'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']]
+        ) if self.formatted_prefix else None
+
+        set_options = petsc_call(
+            self.cbbuilder._options_efunc.name, []
+        )
 
         snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda])
 
         create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])])
 
-        # NOTE: Assuming all solves are linear for now
-        snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY'])
-
         snes_set_jac = petsc_call(
             'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'],
                                 sobjs['Jac'], 'MatMFFDComputeJacobian', Null]
@@ -1478,25 +1446,6 @@ def _setup(self):
         snes_get_ksp = petsc_call('SNESGetKSP',
                                   [sobjs['snes'], Byref(sobjs['ksp'])])
 
-        ksp_set_tols = petsc_call(
-            'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'],
-                                 solver_params['ksp_atol'], solver_params['ksp_divtol'],
-                                 solver_params['ksp_max_it']]
-        )
-
-        ksp_set_type = petsc_call(
-            'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]]
-        )
-
-        ksp_get_pc = petsc_call(
-            'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])]
-        )
-
-        # Even though the default will be jacobi, set to PCNONE for now
-        pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE'])
-
-        ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']])
-
         matvec = self.cbbuilder.main_matvec_callback
         matvec_operation = petsc_call(
             'MatShellSetOperation',
@@ -1569,19 +1518,14 @@ def _setup(self):
         coupled_setup = dmda_calls + (
             snes_create,
             snes_options_prefix,
+            set_options,
             snes_set_dm,
             create_matrix,
             snes_set_jac,
-            snes_set_type,
             global_x,
             local_x,
             get_local_size,
             snes_get_ksp,
-            ksp_set_tols,
-            ksp_set_type,
-            ksp_get_pc,
-            pc_set_type,
-            ksp_set_from_ops,
             matvec_operation,
             formfunc_operation,
             snes_set_options,
@@ -1678,7 +1622,6 @@ def _execute_solve(self):
         Assigns the required time iterators to the struct and executes
         the necessary calls to execute the SNES solver.
         """
-        objs = self.objs
         sobjs = self.solver_objs
         xglob = sobjs['xglobal']
 
@@ -1794,10 +1737,10 @@ class TimeDependent(NonTimeDependent):
       for each `SNESSolve` at every time step, don't require the time loop, but
       may still need access to data from other time steps.
     - All `Function` objects are passed through the initial lowering via the
-      `LinearSolveExpr` object, ensuring the correct time loop is generated
+      `SolveExpr` object, ensuring the correct time loop is generated
       in the main kernel.
     - Another mapper is created based on the modulo dimensions
-      generated by the `LinearSolveExpr` object in the main kernel
+      generated by the `SolveExpr` object in the main kernel
       (e.g., {time: time, t: t0, t + 1: t1}).
     - These two mappers are used to generate a final mapper `symb_to_moddim`
       (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to

devito/petsc/logging.py

@@ -142,7 +142,8 @@ def section(self):
 
     @property
     def summary_key(self):
-        return (self.section, self.sobjs['options_prefix'])
+        user_prefix = self.inject_solve.expr.rhs.user_prefix
+        return (self.section, user_prefix)
 
     def __getattr__(self, attr):
         if attr in self.logobjs.keys():