diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 5daffb78b0..e10afdcda2 100644 --- a/devito/ir/clusters/algorithms.py +++ b/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 diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index fb313f7640..132c793f2a 100644 --- a/devito/ir/stree/algorithms.py +++ b/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: diff --git a/tests/test_dle.py b/tests/test_dle.py index dd3ebeb947..83f0c89319 100644 --- a/tests/test_dle.py +++ b/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)), diff --git a/tests/test_operator.py b/tests/test_operator.py index de445cde01..d6c760efc6 100644 --- a/tests/test_operator.py +++ b/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])',