compiler: Emit arguments used to invoke kernels and add test for specialization with MPI

Author: EdCaunt

devito/operator/operator.py

@@ -989,27 +989,22 @@ def apply(self, **kwargs):
         # Get items expected to be specialized
         specialize = as_tuple(kwargs.pop('specialize', []))
 
-        # In the case of specialization, arguments must be processed before
-        # the operator can be compiled
         if specialize:
             # FIXME: Cannot cope with things like sizes/strides yet since it only
             # looks at the parameters
 
             # Build the arguments list for specialization
-            with self._profiler.timer_on('specialized-arguments-preprocess'):
+            with self._profiler.timer_on('specialization'):
                 args = self.arguments(**kwargs)
-            with switch_log_level(comm=args.comm):
-                self._emit_args_profiling('specialized-arguments-preprocess')
+                # Uses parameters here since Specializer needs {symbol: sympy value}
+                specialized_values = {p: sympify(args[p.name])
+                                      for p in self.parameters if p.name in specialize}
 
-            # Uses parameters here since Specializer needs {symbol: sympy value} mapper
-            specialized_values = {p: sympify(args[p.name])
-                                  for p in self.parameters if p.name in specialize}
+                op = Specializer(specialized_values).visit(self)
 
-            op = Specializer(specialized_values).visit(self)
+            with switch_log_level(comm=args.comm):
+                self._emit_args_profiling('specialization')
 
-            # TODO: Does this cause problems for profilers?
-            # FIXME: Need some way to inspect this Operator for testing
-            # FIXME: Perhaps this should use some separate method
             unspecialized_kwargs = {k: v for k, v in kwargs.items()
                                     if k not in specialize}
 
@@ -1025,9 +1020,7 @@ def apply(self, **kwargs):
         with switch_log_level(comm=args.comm):
             self._emit_args_profiling('arguments-preprocess')
 
-        args_string = ", ".join([f"{p.name}={args[p.name]}"
-                                 for p in self.parameters if p.is_Symbol])
-        debug(f"Invoking `{self.name}` with scalar arguments: {args_string}")
+        self._emit_arguments(args)
 
         # Invoke kernel function with args
         arg_values = [args[p.name] for p in self.parameters]
@@ -1064,6 +1057,28 @@ def _emit_args_profiling(self, tag=''):
         tagstr = ' '.join(tag.split('-'))
         debug(f"Operator `{self.name}` {tagstr}: {elapsed:.2f} s")
 
+    def _emit_arguments(self, args):
+        comm = args.comm
+        scalar_args = ", ".join([f"{p.name}={args[p.name]}"
+                                 for p in self.parameters
+                                 if p.is_Symbol])
+
+        rank = f"[rank{args.comm.Get_rank()}] " if comm is not MPI.COMM_NULL else ""
+
+        msg = f"* {rank}{scalar_args}"
+
+        with switch_log_level(comm=comm):
+            debug(f"Scalar arguments used to invoke `{self.name}`")
+
+            if comm is not MPI.COMM_NULL:
+                # With MPI enabled, we add one entry per rank
+                allmsg = comm.allgather(msg)
+                if comm.Get_rank() == 0:
+                    for m in allmsg:
+                        debug(m)
+            else:
+                debug(msg)
+
     def _emit_build_profiling(self):
         if not is_log_enabled_for('PERF'):
             return

tests/test_specialization.py

@@ -228,12 +228,15 @@ def test_basic(self, caplog, override):
 
             # Ensure that the specialized operator was run
             assert all(s not in caplog.text for s in specialize)
-            assert "specialized arguments preprocess" in caplog.text
+            assert "specialization" in caplog.text
 
         check = np.array(f.data[:])
         f.data[:] = 0
         op.apply(**kwargs)
 
-        assert np.all(check == f.data)
+        assert np.all(check == f.data[:])
 
-# Need to test specialization with MPI (both at)
+    @pytest.mark.parallel(mode=[2, 4])
+    @pytest.mark.parametrize('override', [False, True])
+    def test_basic_mpi(self, caplog, mode, override):
+        self.test_basic(caplog, override)