compiler: Progress on petscoptions callbacks

Author: ZoeLeibowitz

devito/petsc/iet/passes.py

@@ -1,22 +1,26 @@
 import cgen as c
 import numpy as np
 from functools import cached_property
+from sympy import Not
 
 from devito.passes.iet.engine import iet_pass
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
                            DummyExpr, CallableBody, List, Call, Callable,
-                           FindNodes, Section)
-from devito.symbolics import Byref, Macro, FieldFromPointer
+                           FindNodes, Section, Conditional)
+from devito.symbolics import Byref, FieldFromPointer, Macro
 from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
 from devito.tools import frozendict
+import devito.logger as dl
+
 from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData,
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
-                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
-from devito.petsc.types.macros import petsc_func_begin_user
-from devito.petsc.iet.nodes import PetscMetaData
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol,
+                                ConstCharPtr, PetscBool)
+from devito.petsc.types.macros import petsc_func_begin_user, Null
+from devito.petsc.iet.nodes import PetscMetaData, PETScCallable
 from devito.petsc.utils import core_metadata, petsc_languages
 from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder,
                                        CoupledObjectBuilder, BaseSetup, CoupledSetup,
@@ -25,8 +29,6 @@
 from devito.petsc.iet.logging import PetscLogger
 from devito.petsc.iet.utils import petsc_call, petsc_call_mpi
 
-import devito.logger as dl
-
 
 @iet_pass
 def lower_petsc(iet, **kwargs):
@@ -69,9 +71,14 @@ 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.
+    set_solver_option(efuncs)
+    # from IPython import embed; embed()
+
     for iters, (inject_solve,) in inject_solve_mapper.items():
 
-        builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs)
+        builder = Builder(inject_solve, iters, comm, section_mapper, **kwargs)
 
         setup.extend(builder.solver_setup.calls)
 
@@ -80,13 +87,14 @@ def lower_petsc(iet, **kwargs):
 
         efuncs.update(builder.cbbuilder.efuncs)
 
-    populate_matrix_context(efuncs, objs)
+    populate_matrix_context(efuncs)
 
     iet = Transformer(subs).visit(iet)
 
     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
 
@@ -131,7 +139,7 @@ class Builder:
     returning subclasses of the objects initialised in __init__,
     depending on the properties of `inject_solve`.
     """
-    def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs):
+    def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs):
         self.inject_solve = inject_solve
         self.objs = objs
         self.iters = iters
@@ -191,7 +199,7 @@ def calls(self):
         return List(body=self.solve.calls+self.logger.calls)
 
 
-def populate_matrix_context(efuncs, objs):
+def populate_matrix_context(efuncs):
     if not objs['dummyefunc'] in efuncs.values():
         return
 
@@ -205,7 +213,7 @@ def populate_matrix_context(efuncs, objs):
     )
     body = CallableBody(
         List(body=[subdms_expr, fields_expr]),
-        init=(objs['begin_user'],),
+        init=(petsc_func_begin_user,),
         retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])
     )
     name = 'PopulateMatContext'
@@ -215,6 +223,33 @@ def populate_matrix_context(efuncs, objs):
     )
 
 
+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]))
+    ])
+
+    body = CallableBody(
+        body,
+        init=(petsc_func_begin_user,),
+        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+    )
+
+    cb = PETScCallable(
+        'SetPetscOption',
+        body,
+        retval=objs['err'],
+        parameters=(option, value)
+    )
+    # from IPython import embed; embed()
+    efuncs[cb.name] = cb
+
+
 subdms = PointerDM(name='subdms')
 fields = PointerIS(name='fields')
 submats = PointerMat(name='submats')
@@ -262,13 +297,8 @@ def populate_matrix_context(efuncs, objs):
         fields=[subdms, fields, submats], modifier=' *'
     ),
     'subctx': SubMatrixStruct(fields=[rows, cols]),
-    'Null': Macro('NULL'),
     'dummyctx': Symbol('lctx'),
     'dummyptr': DummyArg('dummy'),
     'dummyefunc': Symbol('dummyefunc'),
     'dof': PetscInt('dof'),
-    'begin_user': c.Line('PetscFunctionBeginUser;'),
 })
-
-# Move to macros file?
-Null = Macro('NULL')

devito/petsc/iet/routines.py

@@ -22,6 +22,7 @@
                                 KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM,
                                 VecScatter, DMCast, JacobianStruct, SubMatrixStruct,
                                 CallbackDM, PetscBool)
+from devito.petsc.types.macros import petsc_func_begin_user, Null
 
 
 class CBBuilder:
@@ -118,20 +119,18 @@ def _make_core(self):
     def _petsc_options_callback(self):
         objs = self.objs
         params = self.solver_parameters
-        Null = objs['Null']
+        options_prefix = self.inject_solve.expr.rhs.options_prefix
 
-        has_names = ()
+        body = []
 
+        # from IPython import embed; embed()
         # TODO: improve
         for k, v in params.items():
-            is_set = PetscBool(self.sregistry.make_name(prefix='set'))
-            has_name = petsc_call('PetscOptionsHasName', [
-                Null, Null, String(k), Byref(is_set)])
-            has_names += (has_name,)
+            body.append(petsc_call('SetPetscOption', [String("-"+options_prefix+k), String(v)]))
 
         body = CallableBody(
-            List(body=has_names),
-            init=(objs['begin_user'],),
+            List(body=body),
+            init=(petsc_func_begin_user,),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -281,7 +280,7 @@ def _create_matvec_body(self, body, jacobian):
 
         body = CallableBody(
             List(body=body),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
@@ -420,7 +419,7 @@ def _create_formfunc_body(self, body):
 
         body = CallableBody(
             List(body=body),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),))
 
@@ -530,7 +529,7 @@ def _create_form_rhs_body(self, body):
 
         body = CallableBody(
             List(body=[body]),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
@@ -608,7 +607,7 @@ def _create_initial_guess_body(self, body):
 
         body = CallableBody(
             List(body=[body]),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
@@ -635,7 +634,7 @@ def _make_user_struct_callback(self):
             for i in mainctx.callback_fields
         ]
         struct_callback_body = CallableBody(
-            List(body=body), init=(self.objs['begin_user'],),
+            List(body=body), init=(petsc_func_begin_user,),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
         cb = Callable(
@@ -759,7 +758,7 @@ def _whole_matvec_body(self):
             )
         return CallableBody(
             List(body=(ctx_main, zero_y_memory, BlankLine) + calls),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -894,7 +893,7 @@ def _whole_formfunc_body(self, body):
 
         formfunc_body = CallableBody(
             List(body=body),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=stacks+derefs,
             casts=(f_soa, x_soa),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
@@ -943,7 +942,6 @@ def _submat_callback_body(self):
 
         get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])])
 
-        Null = objs['Null']
         dm_get_info = petsc_call(
             'DMDAGetInfo', [
                 sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']),
@@ -1046,7 +1044,7 @@ def _submat_callback_body(self):
 
         return CallableBody(
             List(body=tuple(body)),
-            init=(objs['begin_user'],),
+            init=(petsc_func_begin_user,),
             stacks=(get_ctx, deref_subdm),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
@@ -1105,7 +1103,7 @@ 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 '') + '_'),
+            'snesprefix': String(options_prefix or ''),
             'options_prefix': options_prefix,
         }
         base_dict['comm'] = self.comm
@@ -1262,7 +1260,7 @@ def _setup(self):
 
         dmda = sobjs['dmda']
 
-        solver_params = self.inject_solve.expr.rhs.solver_parameters
+        # solver_params = self.inject_solve.expr.rhs.solver_parameters
 
         snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])])
 
@@ -1283,7 +1281,7 @@ def _setup(self):
 
         snes_set_jac = petsc_call(
             'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'],
-                                sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']]
+                                sobjs['Jac'], 'MatMFFDComputeJacobian', Null]
         )
 
         global_x = petsc_call('DMCreateGlobalVector',
@@ -1312,16 +1310,17 @@ 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_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'])]
         )
@@ -1339,7 +1338,7 @@ def _setup(self):
         formfunc = self.cbbuilder._F_efunc
         formfunc_operation = petsc_call(
             'SNESSetFunction',
-            [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void),
+            [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void),
              self.snes_ctx]
         )
 
@@ -1373,7 +1372,7 @@ def _setup(self):
             get_local_size,
             global_b,
             snes_get_ksp,
-            ksp_set_tols,
+            # ksp_set_tols,
             ksp_get_pc,
             pc_set_type,
             ksp_set_from_ops,
@@ -1430,7 +1429,7 @@ def _create_dmda(self, dmda):
         stencil_width = self.field_data.space_order
 
         args.append(stencil_width)
-        args.extend([objs['Null']]*nspace_dims)
+        args.extend([Null]*nspace_dims)
 
         # The distributed array object
         args.append(Byref(dmda))
@@ -1465,7 +1464,7 @@ def _setup(self):
 
         snes_set_jac = petsc_call(
             'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'],
-                                sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']]
+                                sobjs['Jac'], 'MatMFFDComputeJacobian', Null]
         )
 
         global_x = petsc_call('DMCreateGlobalVector',
@@ -1506,7 +1505,7 @@ def _setup(self):
         formfunc = self.cbbuilder._F_efunc
         formfunc_operation = petsc_call(
             'SNESSetFunction',
-            [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void),
+            [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void),
              self.snes_ctx]
         )
 
@@ -1529,7 +1528,7 @@ def _setup(self):
 
         create_field_decomp = petsc_call(
             'DMCreateFieldDecomposition',
-            [dmda, Byref(sobjs['nfields']), objs['Null'], Byref(sobjs['fields']),
+            [dmda, Byref(sobjs['nfields']), Null, Byref(sobjs['fields']),
              Byref(sobjs['subdms'])]
         )
         submat_cb = self.cbbuilder.submatrices_callback
@@ -1710,7 +1709,7 @@ def _execute_solve(self):
                 ),
                 petsc_call(
                     'VecScatterCreate',
-                    [xglob, field, target_xglob, self.objs['Null'], Byref(s)]
+                    [xglob, field, target_xglob, Null, Byref(s)]
                 ),
                 petsc_call(
                     'VecScatterBegin',
@@ -1738,7 +1737,7 @@ def _execute_solve(self):
                 )
             )
 
-        snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),)
+        snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], Null, xglob]),)
 
         return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,)
 

devito/petsc/iet/utils.py

@@ -1,8 +1,9 @@
-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
 
+from devito.petsc.iet.nodes import PetscMetaData, PETScCall
+
 
 def petsc_call(specific_call, call_args):
     return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)])