compiler: Improve kernel error check with MPI

Author: FabioLuporini

devito/operator/operator.py

@@ -706,9 +706,14 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
 
         return args
 
-    def _postprocess_errors(self, retval):
+    def _postprocess_errors(self, retval, comm=None):
         if retval == 0:
             return
+        elif comm and comm is not MPI.COMM_NULL:
+            # A rank-local Python exception can leave peer ranks blocked inside
+            # generated MPI waits/exchanges. Abort the communicator instead of
+            # risking an indefinite hang.
+            comm.Abort(retval)
         elif retval == error_mapper['Stability']:
             raise ExecutionError("Detected nan/inf in some output Functions")
         elif retval == error_mapper['KernelLaunch']:
@@ -1008,7 +1013,7 @@ def apply(self, **kwargs):
 
         with self._profiler.timer_on('arguments-postprocess'):
             # Perform error checking
-            self._postprocess_errors(retval)
+            self._postprocess_errors(retval, comm=args.comm)
             # Post-process runtime arguments
             self._postprocess_arguments(args, **kwargs)
 

devito/passes/iet/errors.py

@@ -8,6 +8,7 @@
     Break, Call, Conditional, DummyExpr, EntryFunction, FindNodes, FindSymbols, Iteration,
     KernelLaunch, List, Return, Transformer, make_callable
 )
+from devito.mpi.distributed import MPICommObject
 from devito.passes.iet.engine import iet_pass
 from devito.symbolics import CondEq, MathFunction
 from devito.tools import dtype_to_ctype
@@ -104,28 +105,34 @@ def _check_stability(iet, wmovs=(), rcompile=None, sregistry=None):
 
 def check_launch(graph, options=None, **kwargs):
     """
-    Insert the CHECK_LAUNCH macro if errctl is set to ensure graceful handling of
-    failed kernel launches. This macro should only be inserted if the kernel is
-    directly within a loop, as compilation will fail otherwise.
+    Insert the CHECK_LAUNCH* macros if errctl is set to ensure graceful
+    handling of failed kernel launches. This macro is only inserted if the
+    kernel launch is directly within a loop, as compilation would fail
+    otherwise.
     """
     if options is None or not options.get('errctl', False):
         return
 
     langbb = kwargs['langbb']
 
-    definition = make_launch_macros(langbb)
+    definition = make_launch_macros(langbb, options=options)
     if not definition:
         return
 
-    macro = [langbb['check-launch']]
-
-    _check_launch(graph, definition=definition, macro=macro, **kwargs)
+    _check_launch(graph, definition=definition, options=options, **kwargs)
 
 
 @iet_pass
-def _check_launch(iet, definition=None, macro=None, **kwargs):
+def _check_launch(iet, definition=None, options=None, **kwargs):
     iterations = FindNodes(Iteration).visit(iet)
 
+    if options['mpi'] and \
+       isinstance(iet, EntryFunction) and \
+       any(isinstance(i, MPICommObject) for i in iet.parameters):
+        check = [List(body=c.Statement('CHECK_LAUNCH_RETURN'))]
+    else:
+        check = [List(body=c.Statement('CHECK_LAUNCH'))]
+
     mapper = {}
     for i in iterations:
         # Two stages of substitution to account for the edge case
@@ -135,7 +142,7 @@ def _check_launch(iet, definition=None, macro=None, **kwargs):
         launches = FindNodes(KernelLaunch).visit(i)
 
         for launch in launches:
-            launch_mapper[launch] = List(body=[launch] + macro)
+            launch_mapper[launch] = List(body=[launch] + check)
 
         if launch_mapper:
             mapper[i] = Transformer(launch_mapper).visit(i)
@@ -148,7 +155,7 @@ def _check_launch(iet, definition=None, macro=None, **kwargs):
     return iet, extras
 
 
-def make_launch_macros(langbb):
+def make_launch_macros(langbb, options=None):
     """
     Define macros to check for errors to ensure graceful handling of failed kernel
     launches.
@@ -158,7 +165,20 @@ def make_launch_macros(langbb):
     with contextlib.suppress(NotImplementedError):
         peek = langbb['peek-error']
         success = langbb['error-none']
-        return [('CHECK_LAUNCH', f'if ({peek().name}() != {success}) {{break;}}')]
+
+        headers = [
+            ('CHECK_LAUNCH',
+             f'if ({peek().name}() != {success}) {{break;}}')
+        ]
+
+        if options['mpi']:
+            headers.append(
+                ('CHECK_LAUNCH_RETURN',
+                 f'if ({peek().name}() != {success}) '
+                 f'{{return {error_mapper["KernelLaunch"]};}}')
+            )
+
+        return headers
 
     return []