dsl/compiler: Add PETSc logging instrastructure and tests

Author: ZoeLeibowitz

devito/operator/operator.py

@@ -194,7 +194,7 @@ def _sanitize_exprs(cls, expressions, **kwargs):
     @classmethod
     def _build(cls, expressions, **kwargs):
         # Python- (i.e., compile-) and C-level (i.e., run-time) performance
-        profiler = create_profile('timers')
+        profiler = create_profile('timers', kwargs['language'])
 
         # Lower the input expressions into an IET
         irs, byproduct = cls._lower(expressions, profiler=profiler, **kwargs)
@@ -1004,7 +1004,9 @@ def _emit_apply_profiling(self, args):
         elapsed = fround(self._profiler.py_timers['apply'])
         info(f"Operator `{self.name}` ran in {elapsed:.2f} s")
 
-        summary = self._profiler.summary(args, self._dtype, reduce_over=elapsed)
+        summary = self._profiler.summary(
+            args, self._dtype, self.parameters, reduce_over=elapsed
+        )
 
         if not is_log_enabled_for('PERF'):
             # Do not waste time

devito/operator/profiling.py

@@ -17,6 +17,7 @@
 from devito.parameters import configuration
 from devito.symbolics import subs_op_args
 from devito.tools import DefaultOrderedDict, flatten
+from devito.petsc.logging import PetscSummary
 
 __all__ = ['create_profile']
 
@@ -42,7 +43,7 @@ class Profiler:
 
     _attempted_init = False
 
-    def __init__(self, name):
+    def __init__(self, name, language):
         self.name = name
 
         # Operation reductions observed in sections
@@ -55,6 +56,9 @@ def __init__(self, name):
         # Python-level timers
         self.py_timers = OrderedDict()
 
+        # For language specific summaries
+        self.language = language
+
         self._attempted_init = True
 
     def analyze(self, iet):
@@ -179,7 +183,7 @@ def record_ops_variation(self, initial, final):
     def all_sections(self):
         return list(self._sections) + flatten(self._subsections.values())
 
-    def summary(self, args, dtype, reduce_over=None):
+    def summary(self, args, dtype, params, reduce_over=None):
         """
         Return a PerformanceSummary of the profiled sections.
 
@@ -277,7 +281,7 @@ def _allgather_from_comm(self, comm, time, ops, points, traffic, sops, itershape
         return list(zip(times, opss, pointss, traffics, sops, itershapess))
 
     # Override basic summary so that arguments other than runtime are computed.
-    def summary(self, args, dtype, reduce_over=None):
+    def summary(self, args, dtype, params, reduce_over=None):
         grid = args.grid
         comm = args.comm
 
@@ -338,6 +342,11 @@ def summary(self, args, dtype, reduce_over=None):
                     # data transfers)
                     summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup)
 
+        # Add the language specific summary
+        summary.add_language_summary(
+            self.language,
+            language_summary_mapper[self.language](params)
+        )
         return summary
 
 
@@ -478,6 +487,16 @@ def add_glb_fdlike(self, key, points, time):
 
         self.globals[key] = PerfEntry(time, None, gpointss, None, None, None)
 
+    def add_language_summary(self, lang, summary):
+        """
+        Register a language specific summary (e.g., PetscSummary)
+        and dynamically add a property to access it via perf_summary.<language_name>.
+        """
+        # TODO: Consider renaming `PerformanceSummary` to something more generic
+        # (e.g., `Summary`), or separating `PetscSummary` entirely from
+        # `PerformanceSummary`.
+        setattr(self, lang, summary)
+
     @property
     def globals_all(self):
         v0 = self.globals['vanilla']
@@ -503,21 +522,21 @@ def timings(self):
         return OrderedDict([(k, v.time) for k, v in self.items()])
 
 
-def create_profile(name):
+def create_profile(name, language):
     """Create a new Profiler."""
     if configuration['log-level'] in ['DEBUG', 'PERF'] and \
        configuration['profiling'] == 'basic':
         # Enforce performance profiling in DEBUG mode
         level = 'advanced'
     else:
         level = configuration['profiling']
-    profiler = profiler_registry[level](name)
+    profiler = profiler_registry[level](name, language)
 
     if profiler._attempted_init:
         return profiler
     else:
         warning(f"Couldn't set up `{level}` profiler; reverting to 'advanced'")
-        profiler = profiler_registry['advanced'](name)
+        profiler = profiler_registry['advanced'](name, language)
         # We expect the `advanced` profiler to always initialize successfully
         assert profiler._attempted_init
         return profiler
@@ -533,3 +552,8 @@ def create_profile(name):
     'advisor': AdvisorProfiler
 }
 """Profiling levels."""
+
+
+language_summary_mapper = {
+    'petsc': PetscSummary
+}

devito/petsc/iet/logging.py

@@ -0,0 +1,77 @@
+from functools import cached_property
+
+from devito.symbolics import Byref, FieldFromPointer
+from devito.ir.iet import DummyExpr
+from devito.logger import PERF
+
+from devito.petsc.iet.utils import petsc_call
+from devito.petsc.logging import petsc_return_variable_dict, PetscInfo
+
+
+class PetscLogger:
+    """
+    Class for PETSc loggers that collect solver related statistics.
+    """
+    def __init__(self, level, **kwargs):
+        self.sobjs = kwargs.get('solver_objs')
+        self.sreg = kwargs.get('sregistry')
+        self.section_mapper = kwargs.get('section_mapper', {})
+        self.injectsolve = kwargs.get('injectsolve', None)
+
+        self.function_list = []
+
+        if level <= PERF:
+            self.function_list.extend([
+                'kspgetiterationnumber',
+                'snesgetiterationnumber'
+            ])
+
+        # 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,
+            self.section_mapper, self.injectsolve,
+            self.function_list
+        )
+
+    @cached_property
+    def logobjs(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]]
+        }
+
+    @cached_property
+    def calls(self):
+        """
+        Generate the PETSc calls that will be injected into the C code to
+        extract solver statistics.
+        """
+        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])
+
+            logobj = self.logobjs[param.name]
+
+            calls.append(
+                petsc_call(param.name, inputs + [Byref(logobj)])
+            )
+            # TODO: Perform a PetscCIntCast here?
+            expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol)
+            calls.append(expr)
+
+        return tuple(calls)

devito/petsc/iet/passes.py

@@ -5,7 +5,7 @@
 from devito.passes.iet.engine import iet_pass
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
                            DummyExpr, CallableBody, List, Call, Callable,
-                           FindNodes)
+                           FindNodes, Section)
 from devito.symbolics import Byref, Macro, FieldFromPointer
 from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
@@ -22,8 +22,11 @@
                                        CoupledObjectBuilder, BaseSetup, CoupledSetup,
                                        Solver, CoupledSolver, TimeDependent,
                                        NonTimeDependent)
+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):
@@ -63,22 +66,25 @@ def lower_petsc(iet, **kwargs):
     subs = {}
     efuncs = {}
 
+    # Map PETScSolve to its Section (for logging)
+    section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
+
     for iters, (injectsolve,) in injectsolve_mapper.items():
 
-        builder = Builder(injectsolve, objs, iters, comm, **kwargs)
+        builder = Builder(injectsolve, objs, iters, comm, section_mapper, **kwargs)
 
         setup.extend(builder.solversetup.calls)
 
         # Transform the spatial iteration loop with the calls to execute the solver
-        subs.update({builder.solve.spatial_body: builder.solve.calls})
+        subs.update({builder.solve.spatial_body: builder.calls})
 
         efuncs.update(builder.cbbuilder.efuncs)
 
     populate_matrix_context(efuncs, objs)
 
     iet = Transformer(subs).visit(iet)
 
-    body = core + tuple(setup) + (BlankLine,) + iet.body.body
+    body = core + tuple(setup) + iet.body.body
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
     metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())}
@@ -125,49 +131,64 @@ class Builder:
     returning subclasses of the objects initialised in __init__,
     depending on the properties of `injectsolve`.
     """
-    def __init__(self, injectsolve, objs, iters, comm, **kwargs):
+    def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs):
         self.injectsolve = injectsolve
         self.objs = objs
         self.iters = iters
         self.comm = comm
+        self.section_mapper = section_mapper
         self.kwargs = kwargs
         self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData)
-        self.args = {
+        self.common_kwargs = {
             'injectsolve': self.injectsolve,
             'objs': self.objs,
             'iters': self.iters,
             'comm': self.comm,
+            'section_mapper': self.section_mapper,
             **self.kwargs
         }
-        self.args['solver_objs'] = self.objbuilder.solver_objs
-        self.args['timedep'] = self.timedep
-        self.args['cbbuilder'] = self.cbbuilder
+        self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs
+        self.common_kwargs['timedep'] = self.timedep
+        self.common_kwargs['cbbuilder'] = self.cbbuilder
+        self.common_kwargs['logger'] = self.logger
 
     @cached_property
     def objbuilder(self):
         return (
-            CoupledObjectBuilder(**self.args)
+            CoupledObjectBuilder(**self.common_kwargs)
             if self.coupled else
-            BaseObjectBuilder(**self.args)
+            BaseObjectBuilder(**self.common_kwargs)
         )
 
     @cached_property
     def timedep(self):
         time_mapper = self.injectsolve.expr.rhs.time_mapper
         timedep_class = TimeDependent if time_mapper else NonTimeDependent
-        return timedep_class(**self.args)
+        return timedep_class(**self.common_kwargs)
 
     @cached_property
     def cbbuilder(self):
-        return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args)
+        return CCBBuilder(**self.common_kwargs) \
+            if self.coupled else CBBuilder(**self.common_kwargs)
 
     @cached_property
     def solversetup(self):
-        return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args)
+        return CoupledSetup(**self.common_kwargs) \
+            if self.coupled else BaseSetup(**self.common_kwargs)
 
     @cached_property
     def solve(self):
-        return CoupledSolver(**self.args) if self.coupled else Solver(**self.args)
+        return CoupledSolver(**self.common_kwargs) \
+            if self.coupled else Solver(**self.common_kwargs)
+
+    @cached_property
+    def logger(self):
+        log_level = dl.logger.level
+        return PetscLogger(log_level, **self.common_kwargs)
+
+    @cached_property
+    def calls(self):
+        return List(body=self.solve.calls+self.logger.calls)
 
 
 def populate_matrix_context(efuncs, objs):