compiler: refine distributed dimension check for smarter halotouch

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -472,14 +472,14 @@ 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)
 
+            # 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)
@@ -792,7 +792,10 @@ def halo_write(c, hs):
 
     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.distributor.topology.get(d, 1) > 1
+                          for d in hs.dimensions)
+            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

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,8 @@ 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}))
+        if grid.distributor.comm.rank == 0:
+            print(op)
         assert_structure(op, ['txyz', 't', 'txyz', 'txyz'], 'txyzxyzz')
         op(time=100)
 
@@ -2754,7 +2755,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 +3157,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 +3532,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)