compiler: Add utility function inside petsc routines

Author: ZoeLeibowitz

devito/petsc/iet/logging.py

@@ -77,10 +77,10 @@ def calls(self):
 
             input = self.sobjs[return_variable.input_params]
             output_params = self.petsc_option_mapper[return_variable.name].values()
-            outputs = [Byref(i) for i in output_params]
+            by_ref_output = [Byref(i) for i in output_params]
 
             calls.append(
-                petsc_call(return_variable.name, [input] + outputs)
+                petsc_call(return_variable.name, [input] + by_ref_output)
             )
             # TODO: Perform a PetscCIntCast here?
             exprs = [

devito/petsc/iet/passes.py

@@ -69,19 +69,22 @@ def lower_petsc(iet, **kwargs):
     # Map PETScSolve to its Section (for logging)
     section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
 
-    # prefixes within a single Operator should not be duplicated
+    # Prefixes within a single Operator should not be duplicated
     prefixes = [d.expr.rhs.user_prefix for d in data if d.expr.rhs.user_prefix]
     duplicates = {p for p in prefixes if prefixes.count(p) > 1}
 
-    # TODO: How to avoid the other exception raised given it is an iet_pass?
+    # TODO: Avoid the other exception raised given it is an iet_pass?
     if duplicates:
-        dup_list = ", ".join(sorted(duplicates))
+        dup_list = ", ".join(repr(p) for p in sorted(duplicates))
         raise ValueError(
             f"The following `options_prefix` values are duplicated "
             f"among your PETScSolves. Ensure each one is unique: {dup_list}"
         )
 
+    # List of calls to clear options from the global PETSc options database.
+    # These are executed at the end of the Operator.
     clear_options = []
+
     for iters, (inject_solve,) in inject_solve_mapper.items():
 
         builder = Builder(inject_solve, iters, comm, section_mapper, **kwargs)
@@ -98,7 +101,6 @@ def lower_petsc(iet, **kwargs):
         ),))
 
     populate_matrix_context(efuncs)
-
     iet = Transformer(subs).visit(iet)
     body = core + tuple(setup) + iet.body.body + tuple(clear_options)
     body = iet.body._rebuild(body=body)

devito/petsc/iet/routines.py

@@ -117,18 +117,40 @@ def _make_core(self):
             self._make_initial_guess()
         self._make_user_struct_callback()
 
+    def _make_petsc_callable(self, prefix, body, parameters=()):
+        return PETScCallable(
+            self.sregistry.make_name(prefix=prefix),
+            body,
+            retval=self.objs['err'],
+            parameters=parameters
+        )
+
+    def _make_callable_body(self, body, stacks=(), casts=()):
+        return CallableBody(
+            List(body=body),
+            init=(petsc_func_begin_user,),
+            stacks=stacks,
+            casts=casts,
+            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+        )
+
     def _make_options_callback(self):
-        objs = self.objs
+        """
+        Create two callbacks: one to set PETSc options and one for
+        to clear them.
+
+        Options are only set/cleared if they were not specifed via
+        command line arguments.
+        """
         params = self.solver_parameters
         prefix = self.inject_solve.expr.rhs.formatted_prefix
 
-        set_body = []
-        clear_body = []
+        set_body, clear_body = [], []
 
         for k, v in params.items():
             option = f'-{prefix}{k}'
             if option in sys.argv:
-                # Ensures that the command line options take priority
+                # Ensures that the command line args take priority
                 continue
             option_name = String(option)
             option_value = Null if v is None else String(str(v))
@@ -139,31 +161,12 @@ def _make_options_callback(self):
                 petsc_call('PetscOptionsClearValue', [Null, option_name])
             )
 
-        set_body = CallableBody(
-            List(body=set_body),
-            init=(petsc_func_begin_user,),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        set_body = self._make_callable_body(set_body)
+        clear_body = self._make_callable_body(clear_body)
 
-        clear_body = CallableBody(
-            List(body=clear_body),
-            init=(petsc_func_begin_user,),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        set_callback = self._make_petsc_callable('SetPetscOptions', set_body)
+        clear_callback = self._make_petsc_callable('ClearPetscOptions', clear_body)
 
-        set_callback = PETScCallable(
-            self.sregistry.make_name(prefix='SetPetscOptions'),
-            set_body,
-            retval=objs['err'],
-            parameters=()
-        )
-
-        clear_callback = PETScCallable(
-            self.sregistry.make_name(prefix='ClearPetscOptions'),
-            clear_body,
-            retval=objs['err'],
-            parameters=()
-        )
         self._set_options_efunc = set_callback
         self._efuncs[set_callback.name] = set_callback
         self._clear_options_efunc = clear_callback
@@ -179,13 +182,9 @@ def _make_matvec(self, jacobian, prefix='MatMult'):
         body = self._create_matvec_body(
             List(body=irs.uiet.body), jacobian
         )
-
         objs = self.objs
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix=prefix),
-            body,
-            retval=objs['err'],
-            parameters=(objs['J'], objs['X'], objs['Y'])
+        cb = self._make_petsc_callable(
+            prefix, body, parameters=(objs['J'], objs['X'], objs['Y'])
         )
         self._J_efuncs.append(cb)
         self._efuncs[cb.name] = cb
@@ -303,12 +302,7 @@ def _create_matvec_body(self, body, jacobian):
         # Dereference function data in struct
         derefs = dereference_funcs(ctx, fields)
 
-        body = CallableBody(
-            List(body=body),
-            init=(petsc_func_begin_user,),
-            stacks=stacks+derefs,
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        body = self._make_callable_body(body, stacks=stacks+derefs)
 
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields}
@@ -318,6 +312,7 @@ def _create_matvec_body(self, body, jacobian):
         return body
 
     def _make_formfunc(self):
+        objs = self.objs
         F_exprs = self.field_data.residual.F_exprs
         # Compile `F_exprs` into an IET via recursive compilation
         irs, _ = self.rcompile(
@@ -327,13 +322,9 @@ def _make_formfunc(self):
         body_formfunc = self._create_formfunc_body(
             List(body=irs.uiet.body)
         )
-        objs = self.objs
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='FormFunction'),
-            body_formfunc,
-            retval=objs['err'],
-            parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
-        )
+        parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
+        cb = self._make_petsc_callable('FormFunction', body_formfunc, parameters)
+
         self._F_efunc = cb
         self._efuncs[cb.name] = cb
 
@@ -442,12 +433,7 @@ def _create_formfunc_body(self, body):
         # Dereference function data in struct
         derefs = dereference_funcs(ctx, fields)
 
-        body = CallableBody(
-            List(body=body),
-            init=(petsc_func_begin_user,),
-            stacks=stacks+derefs,
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),))
-
+        body = self._make_callable_body(body, stacks=stacks+derefs)
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields}
 
@@ -466,11 +452,8 @@ def _make_formrhs(self):
             List(body=irs.uiet.body)
         )
         objs = self.objs
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='FormRHS'),
-            body,
-            retval=objs['err'],
-            parameters=(sobjs['callbackdm'], objs['B'])
+        cb = self._make_petsc_callable(
+            'FormRHS', body, parameters=(sobjs['callbackdm'], objs['B'])
         )
         self._b_efunc = cb
         self._efuncs[cb.name] = cb
@@ -552,12 +535,7 @@ def _create_form_rhs_body(self, body):
         # Dereference function data in struct
         derefs = dereference_funcs(ctx, fields)
 
-        body = CallableBody(
-            List(body=[body]),
-            init=(petsc_func_begin_user,),
-            stacks=stacks+derefs,
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        body = self._make_callable_body([body], stacks=stacks+derefs)
 
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
@@ -568,6 +546,7 @@ def _create_form_rhs_body(self, body):
     def _make_initial_guess(self):
         exprs = self.field_data.initial_guess.exprs
         sobjs = self.solver_objs
+        objs = self.objs
 
         # Compile initital guess `eqns` into an IET via recursive compilation
         irs, _ = self.rcompile(
@@ -577,12 +556,8 @@ def _make_initial_guess(self):
         body = self._create_initial_guess_body(
             List(body=irs.uiet.body)
         )
-        objs = self.objs
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='FormInitialGuess'),
-            body,
-            retval=objs['err'],
-            parameters=(sobjs['callbackdm'], objs['xloc'])
+        cb = self._make_petsc_callable(
+            'FormInitialGuess', body, parameters=(sobjs['callbackdm'], objs['xloc'])
         )
         self._initial_guesses.append(cb)
         self._efuncs[cb.name] = cb
@@ -629,13 +604,7 @@ def _create_initial_guess_body(self, body):
 
         # Dereference function data in struct
         derefs = dereference_funcs(ctx, fields)
-
-        body = CallableBody(
-            List(body=[body]),
-            init=(petsc_func_begin_user,),
-            stacks=stacks+derefs,
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        body = self._make_callable_body(body, stacks=stacks+derefs)
 
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
@@ -658,10 +627,7 @@ def _make_user_struct_callback(self):
             DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol)
             for i in mainctx.callback_fields
         ]
-        struct_callback_body = CallableBody(
-            List(body=body), init=(petsc_func_begin_user,),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        struct_callback_body = self._make_callable_body(body)
         cb = Callable(
             self.sregistry.make_name(prefix='PopulateUserContext'),
             struct_callback_body, self.objs['err'],
@@ -731,11 +697,9 @@ def _make_whole_matvec(self):
         objs = self.objs
         body = self._whole_matvec_body()
 
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='WholeMatMult'),
-            List(body=body),
-            retval=objs['err'],
-            parameters=(objs['J'], objs['X'], objs['Y'])
+        parameters = (objs['J'], objs['X'], objs['Y'])
+        cb = self._make_petsc_callable(
+            'WholeMatMult', List(body=body), parameters=parameters
         )
         self._main_matvec_callback = cb
         self._efuncs[cb.name] = cb
@@ -782,13 +746,11 @@ def _whole_matvec_body(self):
                     [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)]
                 ),
             )
-        return CallableBody(
-            List(body=(ctx_main, zero_y_memory, BlankLine) + calls),
-            init=(petsc_func_begin_user,),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        body = (ctx_main, zero_y_memory, BlankLine) + calls
+        return self._make_callable_body(body)
 
     def _make_whole_formfunc(self):
+        objs = self.objs
         F_exprs = self.field_data.residual.F_exprs
         # Compile `F_exprs` into an IET via recursive compilation
         irs, _ = self.rcompile(
@@ -797,13 +759,11 @@ def _make_whole_formfunc(self):
         )
         body = self._whole_formfunc_body(List(body=irs.uiet.body))
 
-        objs = self.objs
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='WholeFormFunc'),
-            body,
-            retval=objs['err'],
-            parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
+        parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
+        cb = self._make_petsc_callable(
+            'WholeFormFunc', body, parameters=parameters
         )
+
         self._F_efunc = cb
         self._efuncs[cb.name] = cb
 
@@ -917,14 +877,10 @@ def _whole_formfunc_body(self, body):
         f_soa = PointerCast(fbundle)
         x_soa = PointerCast(xbundle)
 
-        formfunc_body = CallableBody(
-            List(body=body),
-            init=(petsc_func_begin_user,),
-            stacks=stacks+derefs,
+        formfunc_body = self._make_callable_body(
+            body, stacks=stacks+derefs,
             casts=(f_soa, x_soa),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
-
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields}
 
@@ -941,12 +897,9 @@ def _create_submatrices(self):
             objs['matreuse'],
             objs['Submats'],
         )
-        cb = PETScCallable(
-            self.sregistry.make_name(prefix='MatCreateSubMatrices'),
-            body,
-            retval=objs['err'],
-            parameters=params
-        )
+        cb = self._make_petsc_callable(
+            'MatCreateSubMatrices', body, parameters=params)
+
         self._submatrices_callback = cb
         self._efuncs[cb.name] = cb
 
@@ -1068,12 +1021,7 @@ def _submat_callback_body(self):
             iteration,
         ] + matmult_op
 
-        return CallableBody(
-            List(body=tuple(body)),
-            init=(petsc_func_begin_user,),
-            stacks=(get_ctx, deref_subdm),
-            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-        )
+        return self._make_callable_body(tuple(body), stacks=(get_ctx, deref_subdm))
 
 
 class BaseObjectBuilder:

devito/petsc/logging.py

@@ -125,7 +125,6 @@ class PetscInfo(CompositeObject):
     def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
                  inject_solve, function_list):
 
-        # TODO: change name to match new name elsewehere
         self.petsc_option_mapper = petsc_option_mapper
         self.sobjs = sobjs
         self.section_mapper = section_mapper
@@ -162,12 +161,15 @@ def __getattr__(self, attr):
         # Maps the petsc_option to its generated variable name e.g {'its': its0}
         obj_mapper = self.petsc_option_mapper[attr]
 
-        # Helper to get the value from the petsc profiling struct
+        # Helper to get the value from the profiling struct
         get_val = lambda v: getattr(self.value._obj, v.name)
 
-        # If there's only one value to retrieve for the given attribute, for example, KSPGetIterationNumber
-        # we return it directly, otherwise we return a dictionary of all values e.g for KSPGetTolerances
-        # we return {'rtol': val0, 'abstol': val1, ...}
+        # Return the value(s) for the given PETSc attribute:
+        # - If there's only one output (e.g., for KSPGetIterationNumber),
+        #   return it directly.
+        # - If there are multiple outputs (e.g., for KSPGetTolerances),
+        #   return a dictionary mapping each output name to
+        #   its value, e.g., {'rtol': val0, 'abstol': val1, ...}.
         if len(obj_mapper) == 1:
             return get_val(next(iter(obj_mapper.values())))
         return {k: get_val(v) for k, v in obj_mapper.items()}