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('post-arguments'):
+            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],)
@@ -146,7 +147,7 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
         for k in keys:
             if k not in kwargs:
                 continue
-            elif kwargs[k] is None:
+            elif kwargs[k] is None and inkwargs:
                 # In cases such as rebuild,
                 # the subfunction may be passed explicitly as None
                 return None
@@ -214,7 +215,6 @@ def __subfunc_setup__(self, suffix, keys, dtype=None, inkwargs=False, **kwargs):
         # 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,