compiler: refine distributed dimension check for smarter halotouch

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -453,7 +453,7 @@ def callback(self, clusters, prefix, seen=None):
                not d._defines & hs.distributed_aindices:
                 continue
 
-            if any(halo_write(ci, hs) for ci in clusters[:i]):
+            if any(halo_write(ci, hs, prefix) for ci in clusters[:i]):
                 # If there's a halo write before `c`, then we cannot inject the HaloTouch
                 continue
 
@@ -472,14 +472,17 @@ def callback(self, clusters, prefix, seen=None):
             # the args is important because that's what search functions honor!
             points = sorted(points, key=str)
 
-            # Construct the HaloTouch Cluster
-            expr = Eq(self.B, HaloTouch(*points, halo_scheme=hs))
-
             key0 = lambda i: i in prefix[:-1] or i in hs.loc_indices  # noqa: B023
             key1 = lambda i: i not in hs.distributed_defined  # noqa: B023
             key = lambda i: key0(i) and key1(i)  # noqa: B023
             ispace = c.ispace.project(key)
 
+            # Reconstruct the HaloScheme with the new IterationSpace
+            hs = HaloScheme(c.exprs, ispace)
+
+            # Construct the HaloTouch Cluster
+            expr = Eq(self.B, HaloTouch(*points, halo_scheme=hs))
+
             properties = c.properties.sequentialize()
 
             halo_touch = c.rebuild(exprs=expr, ispace=ispace, properties=properties)
@@ -787,12 +790,15 @@ def normalize_reductions_sparse(cluster, sregistry):
     return cluster.rebuild(processed)
 
 
-def halo_write(c, hs):
+def halo_write(c, hs, prefix):
     loc_vals = hs.loc_values
+    hsdims = hs.dimensions & set(prefix.itdims)
 
     for f in hs.fmapper:
         for a in c.scope.getwrites(f):
-            if set(a.access.indices) & loc_vals:
+            is_write = set(a.access.indices) & loc_vals
+            is_dist = any(c.grid.is_distributed(d) for d in hsdims)
+            if is_write and is_dist:
                 return True
 
     return False

devito/mpi/distributed.py

@@ -261,7 +261,7 @@ def nprocs_local(self):
 
     @property
     def topology(self):
-        return self._topology
+        return DimensionTuple(*self._topology, getters=self.dimensions)
 
     @property
     def topology_logical(self):

devito/mpi/halo_scheme.py

@@ -619,6 +619,10 @@ def classify(exprs, ispace):
                         dims.add(ai.parent)
                     else:
                         dims.add(ai)
+                elif ai is not None and not ai._defines & set(f.dimensions):
+                    # Indirect dimension access, e.g., `u[t, rrecx]`.
+                    # Cannot assume it isn't distributed.
+                    dims.add(ai)
 
         if not halo_labels:
             continue
@@ -642,12 +646,11 @@ def classify(exprs, ispace):
                                           f"scheme for `{f}` along Dimension `{d}`")
             elif hl.pop() is STENCIL:
                 halos.append(Halo(d, s))
-            else:
+            elif d._defines & set(ispace.itdims):
                 raw_loc_indices[d].append(s)
 
         loc_indices, loc_dirs = process_loc_indices(raw_loc_indices,
                                                     ispace.directions)
-
         mapper[f] = HaloSchemeEntry(loc_indices, loc_dirs, halos, dims)
 
     return mapper

devito/types/grid.py

@@ -358,7 +358,7 @@ def is_distributed(self, dim):
         True if `dim` is a distributed Dimension for this CartesianDiscretization,
         False otherwise.
         """
-        return any(dim is d for d in self.distributor.dimensions)
+        return self.distributor.topology.get(dim, 1) > 1
 
     @cached_property
     def _arg_names(self):
@@ -530,7 +530,7 @@ def is_distributed(self, dim):
         False otherwise.
         """
         if self.grid:
-            return any(dim is d for d in self.distributor.dimensions)
+            return self.distributor.topology.get(dim, 1) > 1
         return False
 
     @property

tests/test_mpi.py

@@ -3156,8 +3156,8 @@ def test_fission_due_to_antidep(self, mode):
         # First, check the generated code
         assert_structure(op1, ['t',
                                't,x0_blk0,y0_blk0,x,y,z',
-                               't,x1_blk0,y1_blk0,x,y,z'],
-                         'tx0_blk0y0_blk0xyzx1_blk0y1_blk0xyz')
+                               't,x0_blk0,y0_blk0,x,y,z'],
+                         'tx0_blk0y0_blk0xyzz')
 
         def init(f, v=1):
             f.data[:] = np.indices(grid.shape).sum(axis=0) % (.004*v) + .01