dsl/compiler: Extend solver_parameters handling in PETScSolve and update logging (#2718)

* compiler: Utilise petsctools and start extending solver parameters functionality

examples: Add example using argparse in conjunction with PetscInitialize

* compiler: Edit argument to PetscInitialize

* compiler: Use petsctools to process solver params and start callback to set petsc options

* compiler: Progress on petscoptions callbacks

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

* misc: Add tests

* compiler: Start extending the PetscSummary

* dsl: Fix hashing for solveexpr

* misc: Add tests and clean up

* compiler: Add utility function inside petsc routines

* misc: Clean up

* tests: Add command line tests with random prefixes

* compiler: Create dummy op for petscgetargs

* misc: Update requirements

* misc: Create getargs Op inside function

* compiler/dsl: Add get_info functionality to petscsolve

* compiler/misc: Drop petscgetargs callback, add functions to logging/get_info, clean up and more tests

* workflows: Fix serial wf run

Author: ZoeLeibowitz

.github/workflows/pytest-petsc.yml

@@ -72,9 +72,13 @@ jobs:
       run: |
         ${{ env.RUN_CMD }} python3 -c "from devito import configuration; print(''.join(['%s: %s \n' % (k, v) for (k, v) in configuration.items()]))"
 
-    - name: Test with pytest
+    - name: Test with pytest - serial
       run: |
-        ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+        ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+
+    - name: Test with pytest - parallel
+      run: |
+        ${{ env.RUN_CMD }} python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml -m parallel ${{ env.TESTS }}
 
     - name: Test examples
       run: |

devito/operator/profiling.py

@@ -213,6 +213,11 @@ def summary(self, args, dtype, params, reduce_over=None):
             else:
                 summary.add(name, None, time)
 
+        # Add the language-specific summary if necessary
+        mapper_func = language_summary_mapper.get(self.language)
+        if mapper_func:
+            summary.add_language_summary(self.language, mapper_func(params))
+
         return summary
 
 
@@ -342,7 +347,7 @@ def summary(self, args, dtype, params, reduce_over=None):
                     # data transfers)
                     summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup)
 
-        # Add the language specific summary if necessary
+        # Add the language-specific summary if necessary
         mapper_func = language_summary_mapper.get(self.language)
         if mapper_func:
             summary.add_language_summary(self.language, mapper_func(params))

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 SolverMetaData
 
 
 @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, SolverMetaData):
             ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions)
             processed.append(c.rebuild(ispace=ispace))
         else:

devito/petsc/iet/logging.py

@@ -3,6 +3,7 @@
 from devito.symbolics import Byref, FieldFromPointer
 from devito.ir.iet import DummyExpr
 from devito.logger import PERF
+from devito.tools import frozendict
 
 from devito.petsc.iet.utils import petsc_call
 from devito.petsc.logging import petsc_return_variable_dict, PetscInfo
@@ -12,43 +13,66 @@ class PetscLogger:
     """
     Class for PETSc loggers that collect solver related statistics.
     """
+    # TODO: Update docstring with kwargs
     def __init__(self, level, **kwargs):
+
+        self.query_functions = kwargs.get('get_info', [])
         self.sobjs = kwargs.get('solver_objs')
         self.sreg = kwargs.get('sregistry')
         self.section_mapper = kwargs.get('section_mapper', {})
         self.inject_solve = kwargs.get('inject_solve', None)
 
-        self.function_list = []
-
         if level <= PERF:
-            self.function_list.extend([
+            funcs = [
+                # KSP specific
                 'kspgetiterationnumber',
-                'snesgetiterationnumber'
-            ])
+                'kspgettolerances',
+                'kspgetconvergedreason',
+                'kspgettype',
+                'kspgetnormtype',
+                # SNES specific
+                'snesgetiterationnumber',
+            ]
+            self.query_functions = set(self.query_functions)
+            self.query_functions.update(funcs)
+            self.query_functions = sorted(list(self.query_functions))
 
         # TODO: To be extended with if level <= DEBUG: ...
 
         name = self.sreg.make_name(prefix='petscinfo')
         pname = self.sreg.make_name(prefix='petscprofiler')
 
         self.statstruct = PetscInfo(
-            name, pname, self.logobjs, self.sobjs,
+            name, pname, self.petsc_option_mapper, self.sobjs,
             self.section_mapper, self.inject_solve,
-            self.function_list
+            self.query_functions
         )
 
     @cached_property
-    def logobjs(self):
+    def petsc_option_mapper(self):
         """
-        Create PETSc objects specifically needed for logging solver statistics.
-        """
-        return {
-            info.name: info.variable_type(
-                self.sreg.make_name(prefix=info.output_param)
-            )
-            for func_name in self.function_list
-            for info in [petsc_return_variable_dict[func_name]]
+        For each function in `self.query_functions`, look up its metadata in
+        `petsc_return_variable_dict` and instantiate the corresponding PETSc logging
+        variables with names from the symbol registry.
+
+        Example:
+        --------
+        >>> self.query_functions
+        ['kspgetiterationnumber', 'snesgetiterationnumber', 'kspgettolerances']
+
+        >>> self.petsc_option_mapper
+        {
+            'KSPGetIterationNumber': {'kspits': kspits0},
+            'KSPGetTolerances': {'rtol': rtol0, 'atol': atol0, ...}
         }
+        """
+        opts = {}
+        for func_name in self.query_functions:
+            info = petsc_return_variable_dict[func_name]
+            opts[info.name] = {}
+            for vtype, out in zip(info.variable_type, info.output_param, strict=True):
+                opts[info.name][out] = vtype(self.sreg.make_name(prefix=out))
+        return frozendict(opts)
 
     @cached_property
     def calls(self):
@@ -58,20 +82,21 @@ def calls(self):
         """
         struct = self.statstruct
         calls = []
-        for param in self.function_list:
-            param = petsc_return_variable_dict[param]
-
-            inputs = []
-            for i in param.input_params:
-                inputs.append(self.sobjs[i])
+        for func_name in self.query_functions:
+            return_variable = petsc_return_variable_dict[func_name]
 
-            logobj = self.logobjs[param.name]
+            input = self.sobjs[return_variable.input_params]
+            output_params = self.petsc_option_mapper[return_variable.name].values()
+            by_ref_output = [Byref(i) for i in output_params]
 
             calls.append(
-                petsc_call(param.name, inputs + [Byref(logobj)])
+                petsc_call(return_variable.name, [input] + by_ref_output)
             )
             # TODO: Perform a PetscCIntCast here?
-            expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol)
-            calls.append(expr)
+            exprs = [
+                DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol)
+                for i in output_params
+            ]
+            calls.extend(exprs)
 
         return tuple(calls)

devito/petsc/iet/passes.py

@@ -6,10 +6,12 @@
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
                            DummyExpr, CallableBody, List, Call, Callable,
                            FindNodes, Section)
-from devito.symbolics import Byref, Macro, FieldFromPointer
+from devito.symbolics import Byref, FieldFromPointer, Macro, Null
 from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
 from devito.tools import frozendict
+import devito.logger
+
 from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData,
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
@@ -25,8 +27,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 +69,24 @@ def lower_petsc(iet, **kwargs):
     # Map PETScSolve to its Section (for logging)
     section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
 
+    # Prefixes within the same `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}
+
+    if 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 `Call`s to clear options from the global PETSc options database,
+    # executed at the end of the Operator.
+    clear_options = []
+
     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,11 +95,13 @@ def lower_petsc(iet, **kwargs):
 
         efuncs.update(builder.cbbuilder.efuncs)
 
-    populate_matrix_context(efuncs, objs)
+        clear_options.extend((petsc_call(
+            builder.cbbuilder._clear_options_efunc.name, []
+        ),))
 
+    populate_matrix_context(efuncs)
     iet = Transformer(subs).visit(iet)
-
-    body = core + tuple(setup) + iet.body.body
+    body = core + tuple(setup) + iet.body.body + tuple(clear_options)
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
     metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())}
@@ -131,12 +148,13 @@ 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
         self.comm = comm
         self.section_mapper = section_mapper
+        self.get_info = inject_solve.expr.rhs.get_info
         self.kwargs = kwargs
         self.coupled = isinstance(inject_solve.expr.rhs.field_data, MultipleFieldData)
         self.common_kwargs = {
@@ -183,15 +201,17 @@ def solve(self):
 
     @cached_property
     def logger(self):
-        log_level = dl.logger.level
-        return PetscLogger(log_level, **self.common_kwargs)
+        log_level = devito.logger.logger.level
+        return PetscLogger(
+            log_level, get_info=self.get_info, **self.common_kwargs
+        )
 
     @cached_property
     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 +225,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'
@@ -262,13 +282,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')