compiler: Patch halo exchange placement at stree construction time

devito/ir/clusters/algorithms.py

@@ -458,7 +458,7 @@ def callback(self, clusters, prefix, seen=None):
 
         # Construct a representation of the halo accesses
         processed = list(clusters)
-        for n, c in enumerate(clusters):
+        for c in clusters:
             if c.properties.is_sequential(d) or \
                c in seen:
                 continue
@@ -498,6 +498,7 @@ def callback(self, clusters, prefix, seen=None):
             # Insert `halo_touch` at the top of the IterationSpace within which
             # `c` is scheduled
             index = 0
+            n = processed.index(c)
             for i in reversed(range(n)):
                 if not processed[i].ispace.is_subset(c.ispace):
                     index = i + 1

devito/ir/stree/algorithms.py

@@ -192,24 +192,36 @@ def preprocess(clusters, options=None, **kwargs):
 
         else:
             dims = set(c.ispace.promote(lambda d: d.is_Block).itdims)
+            roots = {d.root for d in dims}
 
             found = []
             for c1 in list(queue):
-                distributed_aindices = c1.halo_scheme.distributed_aindices
-                h_indices = set().union(*[d._defines for d in c1.halo_scheme.loc_indices])
-
                 # Skip if the halo exchange would end up outside
                 # its iteration space
+                loc_indices = c1.halo_scheme.loc_indices
+                h_indices = set().union(*[d._defines for d in loc_indices])
                 if h_indices and not h_indices & dims:
                     continue
 
-                diff = dims - distributed_aindices
-                intersection = dims & distributed_aindices
+                # Ensure the guards are compatible
+                dist_aindices = c1.halo_scheme.distributed_aindices
+                diff = dims - dist_aindices
+                if not all(c1.guards.get(d) == c.guards.get(d) for d in diff):
+                    continue
+
+                # There must be at least one distributed Dimension
+                if not (dims & dist_aindices):
+                    continue
+
+                # Ensure we're inserting within a compatible IterationSpace
+                # E.g., if `dist_aindices` contains a SubDimension `yi`, we
+                # cannot proceed if `c` is over a different SubDimension `yi'`
+                if any(d.root in roots and d not in dims for d in dist_aindices):
+                    continue
 
-                if all(c1.guards.get(d) == c.guards.get(d) for d in diff) and \
-                   len(intersection) > 0:
-                    found.append(c1)
-                    queue.remove(c1)
+                # All good!
+                found.append(c1)
+                queue.remove(c1)
 
             syncs = normalize_syncs(*[c1.syncs for c1 in found])
             if syncs:

tests/test_dle.py

@@ -477,16 +477,10 @@ def test_cache_blocking_imperfect_nest(blockinner):
     op1 = Operator(eqns, opt=('advanced', {'blockinner': blockinner}))
 
     # First, check the generated code
-    bns, _ = assert_blocking(op1, {'x0_blk0'})
+    bns, _ = assert_blocking(op1, {'x0_blk0', 'x1_blk0'})
     trees = retrieve_iteration_tree(bns['x0_blk0'])
-    assert len(trees) == 2
-    assert len(trees[0]) == len(trees[1])
-    assert all(i is j for i, j in zip(trees[0][:4], trees[1][:4], strict=True))
-    assert trees[0][4] is not trees[1][4]
+    assert len(trees) == 1
     assert trees[0].root.dim.is_Block
-    assert trees[1].root.dim.is_Block
-    assert op1.parameters[7] is trees[0][0].step
-    assert op1.parameters[10] is trees[0][1].step
 
     u.data[:] = 0.2
     v.data[:] = 1.5
@@ -633,7 +627,7 @@ def test_nthreads_generation(self):
          (False, False)),
         # two nests, each nest: outermost parallel, innermost sequential
         (['Eq(fc[x,y], fc[x,y+1] + fd[x-1,y])', 'Eq(fd[x-1,y+1], fd[x-1,y] + fc[x,y+1])'],
-         (True, False, False)),
+         (True, False, True, False)),
         # outermost sequential, innermost parallel w/ mixed dimensions
         (['Eq(fc[x+1,y], fc[x,y+1] + fc[x,y])', 'Eq(fc[x+1,y], 2. + fc[x,y+1])'],
          (False, True)),

tests/test_operator.py

@@ -1630,7 +1630,7 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',
           'Eq(tw[t,x,y,z], tv[t,x,y,z-1] + 1.)'),
-         '++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzz'),
+         '++++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzxyz'),
         # 8) WAR 1->2; WAW 1->3
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x+2,y,z])',
@@ -1640,7 +1640,7 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z-2])',
           'Eq(tw[t,x,y,z], tv[t,x,y+1,z] + 1.)'),
-         '+++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzyz'),
+         '++++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzxyz'),
         # 10) WAR 1->2; WAW 1->3
         (('Eq(tu[t-1,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',