api: ensure SparseFunction distributes according to input distributor

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/operator/operator.py

@@ -991,6 +991,8 @@ def apply(self, **kwargs):
         # Build the arguments list to invoke the kernel function
         with self._profiler.timer_on('arguments'):
             args = self.arguments(**kwargs)
+        with switch_log_level(comm=args.comm):
+            self._emit_args_profiling()
 
         # Invoke kernel function with args
         arg_values = [args[p.name] for p in self.parameters]
@@ -1009,7 +1011,10 @@ def apply(self, **kwargs):
                 raise
 
         # Perform error checking
-        self._postprocess_errors(retval)
+        with self._profiler.timer_on('postarguments'):
+            self._postprocess_errors(retval)
+        with switch_log_level(comm=args.comm):
+            self._emit_args_profiling('post')
 
         # Post-process runtime arguments
         self._postprocess_arguments(args, **kwargs)
@@ -1020,6 +1025,11 @@ def apply(self, **kwargs):
 
     # Performance profiling
 
+    def _emit_args_profiling(self, tag=''):
+        fround = lambda i, n=100: ceil(i * n) / n
+        elapsed = fround(self._profiler.py_timers[f'{tag}arguments'])
+        debug(f"Operator `{self.name}` arguments {tag} processed in {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,17 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
             else:
                 dtype = dtype or self.dtype
 
+        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:

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