dsl/compiler: Fix&tidy PETSc mixed solver functionality

.github/workflows/examples-mpi.yml

@@ -17,11 +17,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   build:

.github/workflows/examples.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   tutorials:

.github/workflows/flake8.yml

@@ -10,11 +10,11 @@ on:
    push:
      branches:
        - main
-       - master
+       - petsc
    pull_request:
      branches:
        - main
-       - master
+       - petsc
 
 jobs:
   flake8:

.github/workflows/pytest-core-mpi.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   test-mpi-basic:

.github/workflows/pytest-core-nompi.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   pytest:

.github/workflows/pytest-petsc.yml

@@ -10,11 +10,13 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
+      - biharmonic
   pull_request:
     branches:
       - main
-      - master
+      - petsc
+      - biharmonic
 
 jobs:
   pytest:
@@ -83,6 +85,7 @@ jobs:
         ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py
         ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py
         ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/02_biharmonic.py
 
     - name: Upload coverage to Codecov
       if: "!contains(matrix.name, 'docker')"

.github/workflows/tutorials.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   tutorials:

devito/logger.py

@@ -76,9 +76,10 @@ def set_log_level(level, comm=None):
         used, for example, if one wants to log to one file per rank.
     """
     from devito import configuration
+    from devito.mpi.distributed import MPI
 
     if comm is not None and configuration['mpi']:
-        if comm.rank != 0:
+        if comm != MPI.COMM_NULL and comm.rank != 0:
             logger.removeHandler(stream_handler)
             logger.addHandler(logging.NullHandler())
     else:

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 if necessary
+        mapper_func = language_summary_mapper.get(self.language)
+        if mapper_func:
+            summary.add_language_summary(self.language, mapper_func(params))
+
         return summary
 
 
@@ -366,11 +375,11 @@ class AdvisorProfiler(AdvancedProfiler):
     _default_libs = ['ittnotify']
     _ext_calls = [_api_resume, _api_pause]
 
-    def __init__(self, name):
+    def __init__(self, name, language):
         if self._attempted_init:
             return
 
-        super().__init__(name)
+        super().__init__(name, language)
 
         path = get_advisor_path()
         if path:
@@ -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/passes/iet/engine.py

@@ -4,6 +4,7 @@
 import numpy as np
 from sympy import Mul
 
+from devito.finite_differences.differentiable import Differentiable
 from devito.ir.iet import (
     Call, ExprStmt, Expression, Iteration, SyncSpot, AsyncCallable, FindNodes,
     FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction,
@@ -24,6 +25,7 @@
 from devito.types.args import ArgProvider
 from devito.types.dense import DiscreteFunction
 from devito.types.dimension import AbstractIncrDimension, BlockDimension
+from devito.types.array import ArrayBasic
 
 __all__ = ['Graph', 'iet_pass', 'iet_visit']
 
@@ -492,7 +494,8 @@ def abstract_objects(objects0, sregistry=None):
 
     # Precedence rules make it possible to reconstruct objects that depend on
     # higher priority objects
-    keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension]
+    keys = [Bundle, Array, Differentiable, DiscreteFunction,
+            AbstractIncrDimension, BlockDimension]
     priority = {k: i for i, k in enumerate(keys, start=1)}
     objects = sorted_priority(objects, priority)
 
@@ -527,7 +530,7 @@ def _(i, mapper, sregistry):
     })
 
 
-@abstract_object.register(Array)
+@abstract_object.register(ArrayBasic)
 def _(i, mapper, sregistry):
     if isinstance(i, Lock):
         name = sregistry.make_name(prefix='lock')

devito/petsc/__init__.py

@@ -1 +1,2 @@
 from devito.petsc.solve import *  # noqa
+from devito.petsc.types.equation import *  # noqa

devito/petsc/clusters.py

@@ -20,7 +20,7 @@ def petsc_lift(clusters):
     processed = []
     for c in clusters:
         if isinstance(c.exprs[0].rhs, LinearSolveExpr):
-            ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions)
+            ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions)
             processed.append(c.rebuild(ispace=ispace))
         else:
             processed.append(c)

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.inject_solve = kwargs.get('inject_solve', 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.inject_solve,
+            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)