api: ensure SparseFunction distributes according to input distributor

Author: mloubout

devito/data/data.py

@@ -368,6 +368,9 @@ def __setitem__(self, glb_idx, val, comm_type):
                 # `val` is decomposed, `self` is replicated -> gatherall-like
                 raise NotImplementedError
         elif isinstance(val, np.ndarray):
+            if val.shape == self.view().shape:
+                super().__setitem__(tuple(slice(None) for _ in val.shape), val)
+                return
             if self._is_decomposed:
                 # `val` is replicated, `self` is decomposed -> `val` gets decomposed
                 glb_idx = self._normalize_index(glb_idx)

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]
@@ -1020,6 +1022,11 @@ def apply(self, **kwargs):
 
     # Performance profiling
 
+    def _emit_args_profiling(self):
+        fround = lambda i, n=100: ceil(i * n) / n
+        elapsed = fround(self._profiler.py_timers['arguments'])
+        info(f"Operator `{self.name}` arguments 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],)
@@ -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,