Merge pull request #2626 from devitocodes/biharmonic

dsl/compiler: Fix&tidy PETSc mixed solver functionality

Author: ZoeLeibowitz

.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
+}