compiler/dsl: Extend get_info functions and add tests

Author: ZoeLeibowitz

devito/ir/iet/visitors.py

@@ -292,7 +292,7 @@ def _gen_struct_decl(self, obj, masked=()):
         for i, (n, ct) in zip(fields, ctype._fields_):
             try:
                 entries.append(self._gen_value(i, 0, masked=('const',)))
-                
+
             except AttributeError:
                 cstr = self.ccode(ct)
                 if ct is c_restrict_void_p:

devito/petsc/iet/logging.py

@@ -12,8 +12,8 @@ class PetscLogger:
     """
     Class for PETSc loggers that collect solver related statistics.
     """
-    def __init__(self, level, get_info=None, **kwargs):
-        self.function_list = get_info or []
+    def __init__(self, level, get_info=[], **kwargs):
+        self.function_list = get_info
 
         self.sobjs = kwargs.get('solver_objs')
         self.sreg = kwargs.get('sregistry')
@@ -28,6 +28,8 @@ def __init__(self, level, get_info=None, **kwargs):
                 'kspgettolerances',
                 'kspgetconvergedreason',
                 'kspgettype',
+                'kspgettype',
+                'kspgetnormtype',
                 # SNES specific
                 'snesgetiterationnumber',
             ]

devito/petsc/logging.py

@@ -1,11 +1,12 @@
-from collections import namedtuple, OrderedDict
+import os
+from collections import namedtuple
 from dataclasses import dataclass
-from functools import cached_property
-from cgen import Struct, Value
 
 from devito.types import CompositeObject
 
-from devito.petsc.types import PetscInt, PetscScalar, KSPType
+from devito.petsc.types import (
+    PetscInt, PetscScalar, KSPType, KSPConvergedReason, KSPNormType
+)
 from devito.petsc.utils import petsc_type_mappings, fixed_petsc_type_mappings
 
 
@@ -79,18 +80,18 @@ def _add_properties(self):
         For each function name in `self._functions` (e.g., 'KSPGetIterationNumber'),
         dynamically add a property to the class with the same name.
 
-        Each property returns an OrderedDict that maps each PetscKey to the
+        Each property returns a dict mapping each PetscKey to the
         result of looking up that function on the corresponding PetscEntry,
         if the function exists on that entry.
         """
         def make_property(function):
             def getter(self):
-                return OrderedDict(
-                    (k, getattr(v, function))
+                return {
+                    k: getattr(v, function)
                     for k, v in self.items()
                     # Only include entries that have the function
                     if hasattr(v, function)
-                )
+                }
             return property(getter)
 
         for f in self._functions:
@@ -141,27 +142,25 @@ def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
         self.formatted_prefix = inject_solve.expr.rhs.formatted_prefix
 
         mapper = {v: k for k, v in petsc_type_mappings.items()}
-
         pfields = []
-        self._tmp_fields = []
-
-        # obj_mapper is e.g {'kspits': kspits0}
-        # TODO: think this can just be for (x,y) in ....
-        for obj_mapper in petsc_option_mapper.values():
-            for petsc_option in obj_mapper.values():
-                self._tmp_fields.append(petsc_option)
-                # petsc_type is e.g. 'PetscInt', 'PetscScalar', 'KSPType'
-                petsc_type = str(petsc_option.dtype)
-                if petsc_type in mapper:
-                    ctype = mapper[petsc_type]
-                else:
-                    ctype = fixed_petsc_type_mappings[petsc_type] 
-                pfields.append((petsc_option.name, ctype))
+
+        # All petsc options needed to form the PetscInfo struct
+        # e.g (kspits0, rtol0, atol0, ...)
+        self._fields = [i for j in petsc_option_mapper.values() for i in j.values()]
+
+        for petsc_option in self._fields:
+            # petsc_type is e.g. 'PetscInt', 'PetscScalar', 'KSPType'
+            petsc_type = str(petsc_option.dtype)
+            if petsc_type in mapper:
+                ctype = mapper[petsc_type]
+            else:
+                ctype = fixed_petsc_type_mappings[petsc_type]
+            pfields.append((petsc_option.name, ctype))
         super().__init__(name, pname, pfields)
 
     @property
     def fields(self):
-        return self._tmp_fields
+        return self._fields
 
     @property
     def section(self):
@@ -184,8 +183,10 @@ 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 profiling struct
-        get_val = lambda v: getattr(self.value._obj, v.name)
+        # Decode the value if it is a bytes object
+        def decode_if_bytes(val):
+            return str(os.fsdecode(val)) if isinstance(val, bytes) else val
+        get_val = lambda v: decode_if_bytes(getattr(self.value._obj, v.name))
 
         # - If the function returns a single value (e.g., KSPGetIterationNumber),
         #   return that value directly.
@@ -228,7 +229,7 @@ class PetscReturnVariable:
     ),
     'kspgetconvergedreason': PetscReturnVariable(
         name='KSPGetConvergedReason',
-        variable_type=(PetscInt,),
+        variable_type=(KSPConvergedReason,),
         input_params='ksp',
         output_param=('reason',),
     ),
@@ -238,7 +239,13 @@ class PetscReturnVariable:
         input_params='ksp',
         output_param=('ksptype',),
     ),
-    # SNES specific
+    'kspgetnormtype': PetscReturnVariable(
+        name='KSPGetNormType',
+        variable_type=(KSPNormType,),
+        input_params='ksp',
+        output_param=('kspnormtype',),
+    ),
+    # SNES specific -> will be extended when non-linear solvers are supported
     'snesgetiterationnumber': PetscReturnVariable(
         name='SNESGetIterationNumber',
         variable_type=(PetscInt,),

devito/petsc/solve.py

@@ -14,7 +14,7 @@
 
 # TODO: Rename this to petsc_solve, petscsolve?
 def PETScSolve(target_exprs, target=None, solver_parameters=None,
-               options_prefix=None, get_info=None):
+               options_prefix=None, get_info=[]):
     """
     Returns a symbolic expression representing a linear PETSc solver,
     enriched with all the necessary metadata for execution within an `Operator`.
@@ -68,13 +68,13 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None,
 
     get_info : list[str], optional
         A list of PETSc API functions to collect statistics from the solver.
-        For example, `['kspgetiterationnumber', 'kspgettolerances']`. 
+        For example, `['kspgetiterationnumber', 'kspgettolerances']`.
         Capitalisation does not matter; e.g. `'KSPGetIterationNumber'` and
         `'kspgetiterationnumber'` are treated the same.
 
         List of available functions:
         - ['kspgetiterationnumber', 'kspgettolerances', 'kspgetconvergedreason',
-           'snesgetiterationnumber']
+           'kspgettype', 'kspgetnormtype', 'snesgetiterationnumber']
 
     Returns
     -------
@@ -92,7 +92,7 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None,
 
 class InjectSolve:
     def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None,
-                 get_info=None):
+                 get_info=[]):
         self.solver_parameters = linear_solver_parameters(solver_parameters)
         self.time_mapper = None
         self.target_exprs = target_exprs

devito/petsc/types/object.py

@@ -117,6 +117,10 @@ class KSPType(PetscObject):
     dtype = CustomDtype('KSPType')
 
 
+class KSPNormType(PetscObject):
+    dtype = CustomDtype('KSPNormType')
+
+
 class CallbackSNES(PetscObject):
     """
     PETSc SNES : Non-Linear Systems Solvers.

devito/petsc/utils.py

@@ -95,8 +95,13 @@ def get_petsc_type_mappings():
 petsc_type_mappings = get_petsc_type_mappings()
 
 
+# NOTE: These mappings are only used when constructing ctypes.Structures
+# that wrap PETSc objects. In the generated C code, the fields will still
+# appear as the actual PETSc types.
 fixed_petsc_type_mappings = {
     'KSPType': ctypes.c_char_p,
+    'KSPConvergedReason': ctypes.c_int,
+    'KSPNormType': ctypes.c_int,
 }
 
 

tests/test_petsc.py

@@ -3,7 +3,6 @@
 import numpy as np
 import os
 import re
-from collections import OrderedDict
 
 from conftest import skipif
 from devito import (Grid, Function, TimeFunction, Eq, Operator,
@@ -1437,7 +1436,6 @@ def test_logging(self, log_level):
 
         assert snesits0 == snesits1 == snesits2 == snesits3
 
-        assert isinstance(snesits0, OrderedDict)
         assert len(snesits0) == 1
         key, value = next(iter(snesits0.items()))
         assert str(key) == "PetscKey(name='section0', options_prefix='poisson')"
@@ -1813,20 +1811,25 @@ class TestGetInfo:
     iterations to converge.
     """
     @skipif('petsc')
-    def test_get_info(self):
-
+    def setup_class(self):
+        """
+        Setup grid, functions and equations shared across
+        tests in this class
+        """
         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)
+        self.e, self.f, self.g, self.h = [
+            Function(name=n, grid=grid, space_order=2)
+            for n in ['e', 'f', 'g', 'h']
+        ]
+        self.eq1 = Eq(self.e.laplace, self.f)
+        self.eq2 = Eq(self.g.laplace, self.h)
 
+    @skipif('petsc')
+    def test_get_info(self):
         get_info = ['kspgetiterationnumber', 'snesgetiterationnumber']
-
         petsc = PETScSolve(
-            eq, target=e, options_prefix='pde1', get_info=get_info
+            self.eq1, target=self.e, options_prefix='pde1', get_info=get_info
         )
-
         with switchconfig(language='petsc'):
             op = Operator(petsc)
             summary = op.apply()
@@ -1846,18 +1849,10 @@ 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
+            self.eq1, target=self.e, options_prefix='pde1', get_info=get_info
         )
-
         with switchconfig(language='petsc', log_level=log_level):
             op = Operator(petsc)
             summary = op.apply()
@@ -1878,26 +1873,17 @@ 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
+            self.eq1, target=self.e, options_prefix='pde1', get_info=get_info_1
         )
         solver2 = PETScSolve(
-            eq2, target=g, options_prefix='pde2', get_info=get_info_2
+            self.eq2, target=self.g, options_prefix='pde2', get_info=get_info_2
         )
-
         with switchconfig(language='petsc'):
             op = Operator([solver1, solver2])
             summary = op.apply()
@@ -1916,3 +1902,59 @@ def test_different_solvers(self):
 
         assert not hasattr(entry2, "KSPGetIterationNumber")
         assert hasattr(entry2, "SNESGetIterationNumber")
+
+    @skipif('petsc')
+    def test_case_insensitive(self):
+        """
+        Test that `get_info` is case insensitive
+        """
+        # Create a list with mixed cases
+        get_info = ['KSPGetIterationNumber', 'snesgetiterationnumber']
+        petsc = PETScSolve(
+            self.eq1, target=self.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')
+
+        assert hasattr(entry, "KSPGetIterationNumber")
+        assert hasattr(entry, "SNESGetIterationNumber")
+
+    @skipif('petsc')
+    def test_get_ksp_type(self):
+        """
+        Test that `get_info` can retrieve the KSP type as
+        a string.
+        """
+        get_info = ['kspgettype']
+        solver1 = PETScSolve(
+            self.eq1, target=self.e, options_prefix='poisson1', get_info=get_info
+        )
+        solver2 = PETScSolve(
+            self.eq1, target=self.e, options_prefix='poisson2',
+            solver_parameters={'ksp_type': 'cg'}, get_info=get_info
+        )
+        with switchconfig(language='petsc'):
+            op = Operator([solver1, solver2])
+            summary = op.apply()
+
+        petsc_summary = summary.petsc
+        entry1 = petsc_summary.get_entry('section0', 'poisson1')
+        entry2 = petsc_summary.get_entry('section1', 'poisson2')
+
+        assert hasattr(entry1, "KSPGetType")
+        # Check the type matches the default in linear_solve_defaults
+        # since it has not been overridden
+        assert entry1.KSPGetType == linear_solve_defaults['ksp_type']
+        assert entry1['KSPGetType'] == linear_solve_defaults['ksp_type']
+        assert entry1['kspgettype'] == linear_solve_defaults['ksp_type']
+
+        # Test that the KSP type default is correctly overridden by the
+        # solver_parameters dictionary passed to solver2
+        assert hasattr(entry2, "KSPGetType")
+        assert entry2.KSPGetType == 'cg'
+        assert entry2['KSPGetType'] == 'cg'
+        assert entry2['kspgettype'] == 'cg'