api: fix numpy broadcasting with mpi

Author: mloubout

devito/builtins/initializers.py

@@ -278,18 +278,20 @@ def _extend_padding_mpi(function, nbl):
     distributed `Data` instead resolves global indices on whichever rank owns
     them.
     """
+    distributor = function.grid.distributor
+
     # `TimeFunction`s are rejected upstream, so `axis` indexes `glb_shape` too
     for axis, (nl, nr) in enumerate(as_tuple(nbl)):
-        glb_max = function.grid.distributor.glb_shape[axis] - 1
+        glb_max = distributor.glb_shape[axis] - 1
 
         # One Dimension at a time, each reading the planes left by the previous
         # ones, so that corners come out filled
         for nb, src, lo in ((nl, nl, 0), (nr, glb_max - nr, glb_max - nr + 1)):
             if nb <= 0:
                 continue
+            # `plane` is size 1 along `axis`, so it broadcasts over the slot
             plane = _global_plane(function, axis, src)
-            function.data[axis_slice(function.ndim, axis, lo, lo + nb)] = \
-                np.repeat(plane, nb, axis=axis)
+            function.data[axis_slice(function.ndim, axis, lo, lo + nb)] = plane
 
 
 def _write_interior(function, data, slices):

devito/data/data.py

@@ -347,7 +347,11 @@ def __setitem__(self, glb_idx, val):
                         processed.append(slice(j.start, j.stop, 1))
                     else:
                         processed.append(j)
-                val_idx = as_tuple(processed)
+                # An axis `val` broadcasts along carries the same value for every
+                # rank, so it must be taken whole rather than cut down to this
+                # rank's share of the destination
+                val_idx = tuple(slice(None) if n == 1 else j
+                                for j, n in zip(processed, val.shape, strict=True))
             elif is_windowed(val):
                 # Read all of it, there is no decomposition to restrict it to
                 val_idx = tuple(slice(0, s) for s in val.shape)

tests/test_data.py

@@ -1217,6 +1217,26 @@ def test_big_steps(self, mode):
                        (slice(6, None, -3), slice(6, None, -3))]:
             self._assert_induced(f.data[gslice], a[gslice])
 
+    @pytest.mark.parallel(mode=4)
+    def test_setitem_broadcast(self, mode):
+        """
+        An RHS that is size 1 along an axis broadcasts over that axis, rather
+        than being cut down to each rank's share of the destination.
+        """
+        grid = Grid(shape=(4, 8), topology=(1, 4))
+        f = Function(name='f', grid=grid, space_order=0, dtype=np.int32)
+        f.data[:] = 0
+
+        col = np.array([10, 20, 30, 40], dtype=np.int32)[:, None]
+        f.data[:, 0:3] = col
+
+        expected = np.zeros((4, 8), dtype=np.int32)
+        expected[:, 0:3] = col
+
+        glb = f.data_gather(rank=0)
+        if grid.distributor.myrank == 0:
+            assert np.array_equal(np.asarray(glb), expected)
+
     @pytest.mark.parallel(mode=4)
     def test_setitem(self, mode):
         # __setitem__ mpi slicing tests