Merge pull request #2739 from devitocodes/decoup-sparse-perf

api: Fix SparseFunction distributor setup

Author: mloubout

devito/data/data.py

@@ -396,9 +396,6 @@ def __setitem__(self, glb_idx, val, comm_type):
                         processed.append(j)
                 val_idx = as_tuple(processed)
                 val = val[val_idx]
-            else:
-                # `val` is replicated`, `self` is replicated -> plain ndarray.__setitem__
-                pass
             super().__setitem__(glb_idx, val)
         elif isinstance(val, Iterable):
             if self._is_decomposed:

devito/mpi/distributed.py

@@ -705,6 +705,7 @@ class SparseDistributor(AbstractDistributor):
     distributor : Distributor
         The domain decomposition the SparseDistributor depends on.
     """
+    _is_local = False
 
     def __init__(self, npoint, dimension, distributor):
         super().__init__(npoint, dimension)

devito/operator/operator.py

@@ -989,8 +989,10 @@ def apply(self, **kwargs):
         cfunction = self.cfunction
 
         # Build the arguments list to invoke the kernel function
-        with self._profiler.timer_on('arguments'):
+        with self._profiler.timer_on('arguments-preprocess'):
             args = self.arguments(**kwargs)
+        with switch_log_level(comm=args.comm):
+            self._emit_args_profiling('arguments-preprocess')
 
         # Invoke kernel function with args
         arg_values = [args[p.name] for p in self.parameters]
@@ -1008,18 +1010,25 @@ def apply(self, **kwargs):
             else:
                 raise
 
-        # Perform error checking
-        self._postprocess_errors(retval)
-
-        # Post-process runtime arguments
-        self._postprocess_arguments(args, **kwargs)
+        with self._profiler.timer_on('arguments-postprocess'):
+            # Perform error checking
+            self._postprocess_errors(retval)
+            # Post-process runtime arguments
+            self._postprocess_arguments(args, **kwargs)
 
         # In case MPI is used restrict result logging to one rank only
         with switch_log_level(comm=args.comm):
+            self._emit_args_profiling('arguments-postprocess')
             return self._emit_apply_profiling(args)
 
     # Performance profiling
 
+    def _emit_args_profiling(self, tag=''):
+        fround = lambda i, n=100: ceil(i * n) / n
+        elapsed = fround(self._profiler.py_timers[tag])
+        tagstr = ' '.join(tag.split('-'))
+        debug(f"Operator `{self.name}` {tagstr}: {elapsed:.2f} s")
+
     def _emit_build_profiling(self):
         if not is_log_enabled_for('PERF'):
             return

devito/types/sparse.py

@@ -98,7 +98,8 @@ def __shape_setup__(cls, **kwargs):
         shape = kwargs.get('shape', kwargs.get('shape_global'))
         dimensions = kwargs.get('dimensions')
         npoint = kwargs.get('npoint', kwargs.get('npoint_global'))
-        glb_npoint = SparseDistributor.decompose(npoint, grid.distributor)
+        distributor = kwargs.get('distributor', SparseDistributor)
+        glb_npoint = distributor.decompose(npoint, grid.distributor)
         # Plain SparseFunction construction with npoint.
         if shape is None:
             loc_shape = (glb_npoint[grid.distributor.myrank],)
@@ -184,7 +185,6 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
 
         # Given an array or nothing, create dimension and SubFunction
         if key is not None:
-            dimensions = (self._sparse_dim, Dimension(name='d'))
             if key.ndim > 2:
                 dimensions = (self._sparse_dim, Dimension(name='d'),
                               *mkdims("i", n=key.ndim-2))
@@ -211,14 +211,21 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
             else:
                 dtype = dtype or self.dtype
 
+        # Wether to initialize the subfunction with the provided data
+        # Useful when rebuilding with a placeholder array only used to
+        # infer shape and dtype and set the actual data later
+        if kwargs.get('init_subfunc', True):
+            init = {'initializer': key}
+        else:
+            init = {}
+
         # Complex coordinates are not valid, so fall back to corresponding
         # real floating point type if dtype is complex.
         dtype = dtype(0).real.__class__
-
         sf = SparseSubFunction(
             name=name, dtype=dtype, dimensions=dimensions,
-            shape=shape, space_order=0, initializer=key, alias=self.alias,
-            distributor=self._distributor, parent=self
+            shape=shape, space_order=0, alias=self.alias,
+            distributor=self._distributor, parent=self, **init
         )
 
         if self.npoint == 0:
@@ -230,6 +237,10 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
 
         return sf
 
+    @property
+    def is_local(self):
+        return self._distributor._is_local
+
     @property
     def sparse_position(self):
         return self._sparse_position
@@ -534,7 +545,7 @@ def _dist_data_scatter(self, data=None):
         data = data if data is not None else self.data._local
 
         # If not using MPI, don't waste time
-        if self._distributor.nprocs == 1:
+        if self._distributor.nprocs == 1 or self.is_local:
             return data
 
         # Compute dist map only once
@@ -556,8 +567,13 @@ def _dist_data_scatter(self, data=None):
 
     def _dist_subfunc_scatter(self, subfunc):
         # If not using MPI, don't waste time
-        if self._distributor.nprocs == 1:
-            return {subfunc: subfunc.data}
+        if self._distributor.nprocs == 1 or self.is_local:
+            if self.is_local and self.dist_origin[subfunc] is not None:
+                shift = np.array(self.dist_origin[subfunc], dtype=subfunc.dtype)
+                subfuncd = subfunc.data._local - shift
+            else:
+                subfuncd = subfunc.data
+            return {subfunc: subfuncd}
 
         # Compute dist map only once
         dmap = self._dist_datamap
@@ -581,7 +597,7 @@ def _dist_subfunc_scatter(self, subfunc):
 
     def _dist_data_gather(self, data):
         # If not using MPI, don't waste time
-        if self._distributor.nprocs == 1:
+        if self._distributor.nprocs == 1 or self.is_local:
             return
 
         # Compute dist map only once
@@ -612,7 +628,7 @@ def _dist_subfunc_gather(self, sfuncd, subfunc):
         except AttributeError:
             pass
         # If not using MPI, don't waste time
-        if self._distributor.nprocs == 1:
+        if self._distributor.nprocs == 1 or self.is_local:
             return
 
         # Compute dist map only once

tests/test_rebuild.py

@@ -99,6 +99,6 @@ def test_none_subfunc(self, sfunc):
 
         assert s.coordinates is not None
 
-        # Explicity set coordinates to None
+        # Explicitly set coordinates to None
         sr = s._rebuild(function=None, initializer=None, coordinates=None)
         assert sr.coordinates is None