mpi: Localise SubDimension thicknesses covering a whole rank exactly

devito/data/decomposition.py

@@ -358,15 +358,15 @@ def index_glb_to_loc(self, *args, rel=True):
                 if abs_ofs >= base and abs_ofs <= top:
                     return rel_ofs
                 elif abs_ofs > top:
-                    return top + 1
+                    return top - base
                 else:
                     return None
             else:
                 rel_ofs = abs_ofs - (glb_max - top)
                 if abs_ofs >= glb_max - top and abs_ofs <= glb_max - base:
                     return rel_ofs
                 elif abs_ofs > glb_max - base:
-                    return glb_max - base + 1
+                    return top - base
                 else:
                     return None
         else:

devito/types/dimension.py

@@ -607,19 +607,10 @@ def _arg_values(self, grid=None, **kwargs):
             root = next((d for d in grid.dimensions if d.name == root.name), root)
 
         rtkn = self.value
-        if grid is not None and grid.is_distributed(root):
-            # Get local thickness
-            if self.local:
-                # Dimension is of type `left`/`right` - compute the offset
-                # and then add 1 to get the appropriate thickness
-                if self.value is not None:
-                    tkn = grid.distributor.glb_to_loc(root, rtkn-1, self.side)
-                    tkn = tkn+1 if tkn is not None else 0
-                else:
-                    tkn = 0
-            else:
-                # Dimension is of type `middle`
-                tkn = grid.distributor.glb_to_loc(root, rtkn, self.side) or 0
+        if rtkn and grid is not None and grid.is_distributed(root):
+            # Convert the last boundary index to a local count
+            tkn = grid.distributor.glb_to_loc(root, rtkn-1, self.side)
+            tkn = tkn+1 if tkn is not None else 0
         else:
             tkn = rtkn or 0
 

tests/test_data.py

@@ -503,6 +503,10 @@ def test_glb_to_loc_w_side(self):
         assert d.index_glb_to_loc(2, LEFT) is None
         assert d.index_glb_to_loc(3, RIGHT) is None
 
+        # An offset covering the whole subdomain gives its last index
+        assert d.index_glb_to_loc(8, LEFT) == 2
+        assert d.index_glb_to_loc(7, RIGHT) == 2
+
     def test_loc_to_glb_index_conversions(self):
         d = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 2)
 

tests/test_mpi.py

@@ -3448,6 +3448,41 @@ def test_interpolation_at_uforward(self, mode):
         assert t2.name == 't2'
         assert t2.origin == t + 1
 
+    @pytest.mark.parallel(mode=4)
+    @pytest.mark.parametrize('side', ['left', 'right', 'middle'])
+    def test_thickness_covering_ranks(self, side, mode):
+        """
+        A thickness of 20 points over four ranks of 16 points each covers a rank
+        entirely, so the local thickness of the latter is its whole extent.
+        """
+        class Band(SubDomain):
+            name = 'band'
+
+            def define(self, dimensions):
+                x, y = dimensions
+                return {x: (side, 20, 20) if side == 'middle' else (side, 20), y: y}
+
+        grid = Grid(shape=(64, 4), topology=(4, 1))
+        band = Band(grid=grid)
+        f = Function(name='f', grid=grid, space_order=0)
+        g = Function(name='g', grid=grid, space_order=0)
+
+        op = Operator(Eq(f, 1, subdomain=band), name='Thickness')
+
+        args = op.arguments()
+        rank = grid.distributor.myrank
+        if side in ['left', 'middle']:
+            assert args['x_ltkn0'] == [16, 4, 0, 0][rank]
+        if side in ['right', 'middle']:
+            assert args['x_rtkn0'] == [0, 0, 4, 16][rank]
+
+        op.apply()
+
+        inside = {'left': slice(0, 20), 'right': slice(44, 64),
+                  'middle': slice(20, 44)}[side]
+        g.data[inside] = 1
+        assert np.all(f.data == g.data)
+
 
 def gen_serial_norms(shape, so):
     """