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) for ci in clusters[:i]):
                 # If there's a halo write before `c`, then we cannot inject the HaloTouch
                 continue
 
@@ -787,12 +787,19 @@ def normalize_reductions_sparse(cluster, sregistry):
     return cluster.rebuild(processed)
 
 
-def halo_write(c, hs):
-    loc_vals = hs.loc_values
-
+def _halo_write(c, hs):
+    """
+    Check if the cluster `c` writes into any of the local values read by `hs`.
+    """
     for f in hs.fmapper:
-        for a in c.scope.getwrites(f):
-            if set(a.access.indices) & loc_vals:
-                return True
+        if not any(f.grid.distributor.topology.get(d, 1) > 1
+                   for d in hs.dimensions):
+            # Not distributed halo dimension, write does not impact the halo exchange
+            continue
+
+        if any(set(a.access.indices) & hs.loc_values for a in c.scope.getwrites(f)):
+            # Writing into a local value, which is read by the halo exchange,
+            # creates a write dependency
+            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

@@ -642,12 +642,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

tests/test_mpi.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 import pytest
-from test_dse import TestTTI
 
 from conftest import _R, assert_blocking, assert_structure, body0
 from devito import (
@@ -2216,6 +2215,7 @@ def test_halo_inner_dim(self, mode):
         eq = [Eq(v.forward, v + 1), Eq(e, v.forward.dydz)]
 
         op = Operator(eq, opt=('advanced', {'blocklevels': 0}))
+
         assert_structure(op, ['txyz', 't', 'txyz', 'txyz'], 'txyzxyzz')
         op(time=100)
 
@@ -2754,7 +2754,7 @@ def test_haloupdate_same_timestep_v2(self, mode):
         assert titer.dim is grid.time_dim
         assert titer.nodes[0].body[0].body[0].is_List
         assert len(titer.nodes[0].body[0].body[0].body[0].body) == 1
-        assert not titer.nodes[0].body[0].body[0].body[0].body[0].is_Call
+        assert titer.nodes[0].body[0].body[0].body[0].body[0].is_Call
 
         op.apply(time=0)
 
@@ -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
@@ -3531,9 +3531,9 @@ def test_issue_2448_backward(self, mode):
 
 class TestTTIOp:
 
-    @pytest.mark.skipif(TestTTI is None, reason="Requires installing the tests")
     @pytest.mark.parallel(mode=1)
     def test_halo_structure(self, mode):
+        from test_dse import TestTTI
         solver = TestTTI().tti_operator(opt='advanced', space_order=8)
         op = solver.op_fwd(save=False)