compiler/dsl: Add get_info functionality to petscsolve

Author: ZoeLeibowitz

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
 
 

devito/petsc/iet/logging.py

@@ -12,23 +12,27 @@ class PetscLogger:
     """
     Class for PETSc loggers that collect solver related statistics.
     """
-    def __init__(self, level, **kwargs):
+    def __init__(self, level, get_info=None, **kwargs):
+        self.function_list = get_info or []
+
         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',
                 'kspgettolerances',
                 'kspgetconvergedreason',
                 # SNES specific
                 'snesgetiterationnumber',
-            ])
+            ]
+            for f in funcs:
+                if f not in self.function_list:
+                    self.function_list.append(f)
 
         # TODO: To be extended with if level <= DEBUG: ...
 
@@ -56,7 +60,7 @@ def petsc_option_mapper(self):
         >>> self.petsc_option_mapper
         {
             'KSPGetIterationNumber': {'kspits': kspits0},
-            'KSPGetTolerances': {'rtol': rtol0, 'abstol': abstol0, ...}
+            'KSPGetTolerances': {'rtol': rtol0, 'atol': atol0, ...}
         }
         """
         opts = {}

devito/petsc/iet/passes.py

@@ -170,6 +170,7 @@ def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs):
         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 = {
@@ -217,7 +218,9 @@ def solve(self):
     @cached_property
     def logger(self):
         log_level = dl.logger.level
-        return PetscLogger(log_level, **self.common_kwargs)
+        return PetscLogger(
+            log_level, self.get_info, **self.common_kwargs
+        )
 
     @cached_property
     def calls(self):

devito/petsc/logging.py

@@ -17,6 +17,9 @@ def __init__(self, **kwargs):
     def __getitem__(self, key):
         return self._properties[key.lower()]
 
+    def __len__(self):
+        return len(self._properties)
+
     def __repr__(self):
         return f"PetscEntry({', '.join(f'{k}={v}' for k, v in self.kwargs.items())})"
 
@@ -63,7 +66,9 @@ def petsc_entry(self, petscinfo):
         Create a named tuple entry for the given PetscInfo object,
         containing the values for each PETSc function call.
         """
-        funcs = self._functions
+        funcs = [
+            petsc_return_variable_dict[f].name for f in petscinfo.function_list
+        ]
         values = tuple(getattr(petscinfo, c) for c in funcs)
         return PetscEntry(**{k: v for k, v in zip(funcs, values)})
 
@@ -171,7 +176,7 @@ def __getattr__(self, attr):
         #   return that value directly.
         # - If the function returns multiple values (e.g., KSPGetTolerances),
         #   return a dictionary mapping each output name to its value,
-        #   e.g., {'rtol': val0, 'abstol': val1, ...}.
+        #   e.g., {'rtol': val0, 'atol': 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()}
@@ -204,8 +209,7 @@ class PetscReturnVariable:
         name='KSPGetTolerances',
         variable_type=(PetscScalar, PetscScalar, PetscScalar, PetscInt),
         input_params='ksp',
-        # TODO: check if maxits is max_its in command line
-        output_param=('rtol', 'atol', 'dtol', 'maxits'),
+        output_param=('rtol', 'atol', 'divtol', 'max_it'),
     ),
     'kspgetconvergedreason': PetscReturnVariable(
         name='KSPGetConvergedReason',

devito/petsc/solve.py

@@ -13,7 +13,8 @@
 
 
 # TODO: Rename this to petsc_solve, petscsolve?
-def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix=None):
+def PETScSolve(target_exprs, target=None, solver_parameters=None,
+               options_prefix=None, get_info=None):
     """
     Returns a symbolic expression representing a linear PETSc solver,
     enriched with all the necessary metadata for execution within an `Operator`.
@@ -45,29 +46,37 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix
     solver_parameters : dict, optional
         PETSc solver options.
 
+    options_prefix : str, optional
+        Prefix for the solver, used to configure options via the command line. If not
+        provided, a default prefix is generated by Devito.
+
+    get_info : list[str], optional
+        A list of PETSc API functions to collect statistics from the solver.
+        For example, `['kspgetiterationnumber', 'kspgettolerances']`.
+
     Returns
     -------
     Eq:
         A symbolic expression that wraps the linear solver.
         This can be passed directly to a Devito Operator.
     """
     if target is not None:
-        return InjectSolve(
-            solver_parameters, {target: target_exprs}, options_prefix
-        ).build_expr()
+        return InjectSolve(solver_parameters, {target: target_exprs},
+                           options_prefix, get_info).build_expr()
     else:
-        return InjectMixedSolve(
-            solver_parameters, target_exprs, options_prefix
-        ).build_expr()
+        return InjectMixedSolve(solver_parameters, target_exprs,
+                                options_prefix, get_info).build_expr()
 
 
 class InjectSolve:
-    def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None):
+    def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None,
+                 get_info=None):
         self.solver_parameters = linear_solver_parameters(solver_parameters)
         self.time_mapper = None
         self.target_exprs = target_exprs
         self.user_prefix = options_prefix
         self.formatted_prefix = format_options_prefix(options_prefix)
+        self.get_info = get_info
 
     def build_expr(self):
         target, funcs, field_data = self.linear_solve_args()
@@ -80,7 +89,8 @@ def build_expr(self):
             time_mapper=self.time_mapper,
             localinfo=localinfo,
             user_prefix=self.user_prefix,
-            formatted_prefix=self.formatted_prefix
+            formatted_prefix=self.formatted_prefix,
+            get_info=self.get_info
         )
         return PetscEq(target, linear_solve)
 

devito/petsc/types/types.py

@@ -39,11 +39,13 @@ class SolveExpr(MetaData):
     """
     __rargs__ = ('expr',)
     __rkwargs__ = ('solver_parameters', 'field_data', 'time_mapper',
-                   'localinfo', 'user_prefix', 'formatted_prefix')
+                   'localinfo', 'user_prefix', 'formatted_prefix',
+                   'get_info')
 
     def __new__(cls, expr, solver_parameters=None,
                 field_data=None, time_mapper=None, localinfo=None,
-                user_prefix=None, formatted_prefix=None, **kwargs):
+                user_prefix=None, formatted_prefix=None,
+                get_info=None, **kwargs):
 
         with sympy_mutex:
             if isinstance(expr, tuple):
@@ -57,6 +59,7 @@ def __new__(cls, expr, solver_parameters=None,
         obj.localinfo = localinfo
         obj.user_prefix = user_prefix
         obj.formatted_prefix = formatted_prefix
+        obj.get_info = get_info
         return obj
 
     def __repr__(self):

tests/test_petsc.py

@@ -1448,8 +1448,8 @@ def test_logging(self, log_level):
         tols = entry0.KSPGetTolerances
         assert tols['rtol'] == linear_solve_defaults['ksp_rtol']
         assert tols['atol'] == linear_solve_defaults['ksp_atol']
-        assert tols['dtol'] == linear_solve_defaults['ksp_divtol']
-        assert tols['maxits'] == linear_solve_defaults['ksp_max_it']
+        assert tols['divtol'] == linear_solve_defaults['ksp_divtol']
+        assert tols['max_it'] == linear_solve_defaults['ksp_max_it']
 
     @skipif('petsc')
     @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG'])
@@ -1655,8 +1655,8 @@ def test_tolerances(self, log_level):
         # appear as expected in the `PetscSummary`.
         assert tolerances['rtol'] == params['ksp_rtol']
         assert tolerances['atol'] == params['ksp_atol']
-        assert tolerances['dtol'] == params['ksp_divtol']
-        assert tolerances['maxits'] == params['ksp_max_it']
+        assert tolerances['divtol'] == params['ksp_divtol']
+        assert tolerances['max_it'] == params['ksp_max_it']
 
     @skipif('petsc')
     def test_clearing_options(self):
@@ -1805,3 +1805,117 @@ def test_solveexpr(self):
             options_prefix='poisson3'
         )
         assert hash(petsc3.rhs) != hash(petsc4.rhs)
+
+
+class TestGetInfo:
+    """
+    Test the `get_info` optional argument to `PETScSolve`.
+
+    This argument can be used independently of the `log_level` to retrieve
+    specific information about the solve, such as the number of KSP
+    iterations to converge.
+    """
+    @skipif('petsc')
+    def test_get_info(self):
+
+        grid = Grid(shape=(11, 11), dtype=np.float64)
+        functions = [Function(name=n, grid=grid, space_order=2)
+                     for n in ['e', 'f']]
+        e, f = functions
+        eq = Eq(e.laplace, f)
+
+        get_info = ['kspgetiterationnumber', 'snesgetiterationnumber']
+
+        petsc = PETScSolve(
+            eq, target=e, options_prefix='pde1', get_info=get_info
+        )
+
+        with switchconfig(language='petsc'):
+            op = Operator(petsc)
+            summary = op.apply()
+
+        petsc_summary = summary.petsc
+        entry = petsc_summary.get_entry('section0', 'pde1')
+
+        # Verify that the entry contains only the requested info
+        # (since logging is not set)
+        assert len(entry) == 2
+        assert hasattr(entry, "KSPGetIterationNumber")
+        assert hasattr(entry, "SNESGetIterationNumber")
+
+    @skipif('petsc')
+    @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG'])
+    def test_get_info_with_logging(self, log_level):
+        """
+        Test that `get_info` works correctly when logging is enabled.
+        """
+        grid = Grid(shape=(11, 11), dtype=np.float64)
+        functions = [Function(name=n, grid=grid, space_order=2)
+                     for n in ['e', 'f']]
+        e, f = functions
+        eq = Eq(e.laplace, f)
+
+        get_info = ['kspgetiterationnumber']
+
+        petsc = PETScSolve(
+            eq, target=e, options_prefix='pde1', get_info=get_info
+        )
+
+        with switchconfig(language='petsc', log_level=log_level):
+            op = Operator(petsc)
+            summary = op.apply()
+
+        petsc_summary = summary.petsc
+        entry = petsc_summary.get_entry('section0', 'pde1')
+
+        # With logging enabled, the entry should include both the
+        # requested KSP iteration number and additional PETSc info
+        # (e.g., SNES iteration count logged at PERF/DEBUG).
+        assert len(entry) > 1
+        assert hasattr(entry, "KSPGetIterationNumber")
+        assert hasattr(entry, "SNESGetIterationNumber")
+
+    @skipif('petsc')
+    def test_different_solvers(self):
+        """
+        Test that `get_info` works correctly when multiple solvers are used
+        within the same Operator.
+        """
+        grid = Grid(shape=(11, 11), dtype=np.float64)
+        functions = [Function(name=n, grid=grid, space_order=2)
+                     for n in ['e', 'f', 'g', 'h']]
+        e, f, g, h = functions
+
+        eq1 = Eq(e.laplace, f)
+        eq2 = Eq(g.laplace, h)
+
+        # Create two PETScSolve instances with different get_info arguments
+
+        get_info_1 = ['kspgetiterationnumber']
+        get_info_2 = ['snesgetiterationnumber']
+
+        solver1 = PETScSolve(
+            eq1, target=e, options_prefix='pde1', get_info=get_info_1
+        )
+        solver2 = PETScSolve(
+            eq2, target=g, options_prefix='pde2', get_info=get_info_2
+        )
+
+        with switchconfig(language='petsc'):
+            op = Operator([solver1, solver2])
+            summary = op.apply()
+
+        petsc_summary = summary.petsc
+
+        assert len(petsc_summary) == 2
+        assert len(petsc_summary.KSPGetIterationNumber) == 1
+        assert len(petsc_summary.SNESGetIterationNumber) == 1
+
+        entry1 = petsc_summary.get_entry('section0', 'pde1')
+        entry2 = petsc_summary.get_entry('section1', 'pde2')
+
+        assert hasattr(entry1, "KSPGetIterationNumber")
+        assert not hasattr(entry1, "SNESGetIterationNumber")
+
+        assert not hasattr(entry2, "KSPGetIterationNumber")
+        assert hasattr(entry2, "SNESGetIterationNumber")