diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 86170eafa1..374536fa62 100644 --- a/devito/ir/clusters/algorithms.py +++ b/devito/ir/clusters/algorithms.py @@ -130,8 +130,61 @@ class Schedule(Queue): @timed_pass(name='schedule') def process(self, clusters): + clusters = self._schedule_adjacent(clusters) + clusters = self._constrain_nonadjacent(clusters) + + return clusters + + def _schedule_adjacent(self, clusters): + """ + Schedule Clusters that are adjacent in program order. + """ return self._process_fatd(clusters, 1) + def _constrain_nonadjacent(self, clusters): + """ + Clusters that are non-adjacent in program order may require additional + constraints on their IterationSpaces. + + For example, consider the following statements: + + - `u[x, y] = f(w[x, y])` + - `f[x] = 3` + - `v[x, y] = g(u[x, y ± 1])` + + The first and third statements are non-adjacent in program order, have + the same IterationSpace, and have a flow- or anti-dependence along `y` + due to `y ± 1`. Hence, those two statements will never be fusible, so + their IterationSpace must be constrained to prevent such a possibility + by any of the later passes (e.g., topofusion). + """ + # A family collects the consecutive groups of Clusters with the same + # IterationSpace + key = lambda c: c.ispace + families = DefaultOrderedDict(list) + for ispace, group in groupby(clusters, key=key): + families[ispace].append(list(group)) + + mapper = {} + for family in families.values(): + if len(family) < 2: + continue + + candidates = flatten(family) + scheduled = self._schedule_adjacent(candidates) + for c, c1 in zip(candidates, scheduled, strict=True): + if c1.ispace.intervals == c.ispace.intervals: + continue + + # Only retain the new Interval constraints. Directions were + # inferred without the intervening Clusters. + ispace1 = IterationSpace(c1.ispace.intervals, + c.ispace.sub_iterators, + c.ispace.directions) + mapper[c] = c.rebuild(ispace=ispace1) + + return [mapper.get(c, c) for c in clusters] + def callback(self, clusters, prefix, backlog=None, known_break=None): if not prefix: return clusters @@ -458,7 +511,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 +551,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..cae263bc81 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 + # There must be at least one distributed Dimension + dist_aindices = c1.halo_scheme.distributed_aindices + if not (dims & dist_aindices): + continue + + # Ensure the guards are compatible + diff = dims - dist_aindices + if not all(c1.guards.get(d) == c.guards.get(d) for d in diff): + 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/devito/ir/support/basic.py b/devito/ir/support/basic.py index 52799a52af..bd8e3927e2 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -379,6 +379,10 @@ def distance(self, other, logical=False): # E.g., `self=R` and `self.itintervals=(y,)` => `sai=None` pass + if skippable_interval(sai, self.ispace, sit) and \ + skippable_interval(oai, other.ispace, oit): + continue + # In some cases, the distance degenerates because `self` and # `other` never intersect, which essentially means there's no # dependence between them. In this case, we set the distance to a @@ -406,8 +410,10 @@ def distance(self, other, logical=False): # Case 3: `self` and `other` have some special form such that # it's provable that they never intersect - if sai and sit == oit and disjoint_test(self[n], other[n], sai, sit): - return Vector(S.ImaginaryUnit) + if sit == oit: + dims = {sai, oai} - {None} + if any(disjoint_test(self[n], other[n], d, sit) for d in dims): + return Vector(S.ImaginaryUnit) # Compute the distance along the current IterationInterval if self.function._mem_shared: @@ -426,10 +432,6 @@ def distance(self, other, logical=False): ret.append(other[n] - self[n]) else: ret.append(self[n] - other[n]) - elif sai in self.ispace and oai in other.ispace: - # E.g., `self=R`, `sai=time`, - # `self.itintervals=(time, x, y)`, `n=0` - continue elif not sai and not oai: if self[n] - other[n] == 0: # E.g., `self=R` and `other=W` @@ -1505,6 +1507,26 @@ def vinf(entries): return Vector(*(entries + [S.Infinity])) +def skippable_interval(d, ispace, it): + """ + Return True if the IterationInterval `it` can be skipped while matching + the access Dimension `d` to the IterationSpace `ispace`. + + Constant accesses, represented by `d=None`, do not consume an + IterationInterval. A non-constant access can skip `it` only if `d` + occurs elsewhere in `ispace`. Otherwise the access is irregular and + must be handled conservatively. + + Examples + -------- + Given `W` and `R` over `(r, i)`, `r` is skippable + for both accesses, while `i` is not skippable for the write. By contrast, + `x` is not skippable for `a[y]` over `(x)` because `y` does not + occur in the IterationSpace. + """ + return d is None or (d in ispace and not d._defines & it.dim._defines) + + def disjoint_test(e0, e1, d, it): """ A rudimentary test to check if two accesses `e0` and `e1` along `d` within 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_ir.py b/tests/test_ir.py index 9f26ecb8c5..cd7bbee4b9 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -25,7 +25,9 @@ from devito.symbolics import DefFunction, FieldFromPointer from devito.tools import prod from devito.tools.data_structures import frozendict -from devito.types import Array, Bundle, CriticalRegion, Jump, Scalar, Symbol +from devito.types import ( + Array, Bundle, CriticalRegion, CustomDimension, Jump, Scalar, Symbol +) class TestVectorHierarchy: @@ -912,6 +914,25 @@ def test_nodep(self, eqns): scope = Scope(eqns) assert len(scope.d_all) == 0 + def test_nodep_with_nonaccessed_implicit_dim(self): + """ + An outer implicit Dimension must not hide that an indexed write and a + constant read touch disjoint points of a CustomDimension. + """ + r = Dimension(name='r') + i = CustomDimension(name='i', symbolic_min=0, symbolic_max=3) + + a = Array(name='a', dimensions=(i,)) + f = Function(name='f', dimensions=(r, i), shape=(3, 4)) + + eqns = [Eq(a[i], r, implicit_dims=(r, i)), + Eq(f, a[4], implicit_dims=(r, i))] + eqns = [LoweredEq(e) for e in eqns] + + scope = Scope(eqns) + deps = [d for d in scope.d_all if d.function is a] + assert not deps + @pytest.mark.parametrize('eqns', [ ['Eq(a0[4], 1)', 'Eq(s, a0[4])'], ['Eq(a1[x+1, 4], 1)', 'Eq(s, a1[x, 4])'], diff --git a/tests/test_mpi.py b/tests/test_mpi.py index d5ad76323a..5b78305670 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -3289,7 +3289,7 @@ def test_overriding_from_different_grid(self, mode): assert np.all(u3.data[0, 3:-3, 3:-3] == 1.) @pytest.mark.parallel(mode=4) - def test_fission_due_to_antidep(self, mode): + def test_fission_due_to_antidep_along_innermost_dim(self, mode): grid = Grid(shape=(16, 16, 64), dtype=np.float64) u = TimeFunction(name='u', grid=grid, space_order=4) @@ -3305,8 +3305,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,x0_blk0,y0_blk0,x,y,z'], - 'tx0_blk0y0_blk0xyzz') + 't,x1_blk0,y1_blk0,x,y,z'], + 'tx0_blk0y0_blk0xyzx1_blk0y1_blk0xyz') def init(f, v=1): f.data[:] = np.indices(grid.shape).sum(axis=0) % (.004*v) + .01 diff --git a/tests/test_operator.py b/tests/test_operator.py index a6090b54be..218c5e6710 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -1650,7 +1650,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])', @@ -1660,7 +1660,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])', @@ -2192,6 +2192,53 @@ def test_topofuse_w_numeric_dim(self): assert_structure(op, ['r,i', 'r'], 'r,i') + def test_topofuse_preserves_nonadjacent_fission(self): + """ + A shallower IterationSpace may separate a stencil producer and consumer + during adjacent scheduling. Topofusion must preserve their fission. + """ + x, y = dimensions('x y') + + u = Function(name='u_nonadj', dimensions=(x, y), shape=(16, 16), + space_order=1) + v = Function(name='v_nonadj', dimensions=(x, y), shape=(16, 16), + space_order=1) + w = Function(name='w_nonadj', dimensions=(x,), shape=(16,)) + + eqns = [Eq(u, 1), + Eq(w, 2), + Eq(v, u[x, y - 1] + u[x, y + 1])] + + op = Operator(eqns, opt=('topofuse', {'openmp': False})) + + assert_structure(op, ['x,y', 'x', 'x,y']) + + def test_nonadjacent_constraints_preserve_directions(self): + """ + Constraints inferred from non-adjacent Clusters must not replace the + directions inferred while the intervening Clusters were still visible. + """ + x, y = dimensions('x y') + + u = Function(name='u_direction', dimensions=(x, y), shape=(16, 16), + space_order=1) + p = Function(name='p_direction', dimensions=(x, y), shape=(16, 16), + space_order=1) + q = Function(name='q_direction', dimensions=(x,), shape=(16,), + space_order=1) + + eqns = [Eq(u[x, y], 1), + Eq(q[x], p[x + 1, 0]), + Eq(u[x, y + 1], 2), + Eq(p[x, y], 3)] + + op = Operator(eqns, opt=('topofuse', {'openmp': False})) + + assert_structure(op, ['x,y', 'x', 'x,y'], 'x,y,y') + + iterations = FindNodes(Iteration).visit(op) + assert iterations[0].direction is Backward + @pytest.mark.parametrize('eqns, expected, exp_trees, exp_iters', [ (['Eq(u[0, x], 1)', 'Eq(u[1, x], u[0, x + h_x] + u[0, x - h_x] - 2*u[0, x])'], diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index c3e9e68a7a..e6cab67311 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -10,7 +10,9 @@ SparseFunction, SparseTimeFunction, SubDomain, SubDomainSet, TensorFunction, TimeFunction, VectorFunction, solve ) -from devito.ir import Expression, FindNodes, FindSymbols, Iteration, SymbolRegistry +from devito.ir import ( + Expression, FindNodes, FindSymbols, Iteration, SymbolRegistry, retrieve_iteration_tree +) from devito.tools import timed_region @@ -263,6 +265,77 @@ def define(self, dimensions): assert_structure(op, ['t', 'txyz', 'txyz'], 'txyzyz') +class TestSubDomainScheduling: + """Tests scheduling across different SubDomains.""" + + def test_topofusion_preserves_fission(self): + """ + Ensure topofusion does not merge a stencil producer and consumer after + orthogonal SubDomains kept them apart during initial scheduling. + """ + + class Slab(SubDomain): + + def __init__(self, axis, side, **kwargs): + self.name = f'{kwargs["grid"].dimensions[axis].name}_{side}' + self.axis = axis + self.side = side + super().__init__(**kwargs) + + def define(self, dimensions): + return { + d: (self.side, 8) if i == self.axis else d + for i, d in enumerate(dimensions) + } + + grid = Grid(shape=(16, 16), extent=(15., 15.)) + slabs = [Slab(axis, 'left', grid=grid) for axis in (0, 1)] + + src = Function(name='src', grid=grid) + src.data[:] = np.arange(src.data.size, dtype=grid.dtype).reshape( + src.data.shape + ) + + producer_eqs = [] + consumer_eqs = [] + reference_eqs = [] + states = [] + + for axis, slab in zip((0, 1), slabs, strict=True): + d = grid.dimensions[axis] + u = TimeFunction(name=f'u_{d.name}', grid=grid, + space_order=2, time_order=1) + v = TimeFunction(name=f'v_{d.name}', grid=grid, + space_order=2, time_order=1) + u_ref = u.func(name=f'u_{d.name}_ref') + v_ref = v.func(name=f'v_{d.name}_ref') + du = getattr(u, f'd{d.name}') + du_ref = getattr(u_ref, f'd{d.name}') + + producer_eqs.append(Eq(u, u.backward + src + 1, subdomain=slab)) + consumer_eqs.append(Eq(v, v.backward + du, subdomain=slab)) + reference_eqs.extend([ + Eq(u_ref, u_ref.backward + src + 1, subdomain=slab), + Eq(v_ref, v_ref.backward + du_ref, subdomain=slab) + ]) + states.append((v, v_ref)) + + op = Operator(producer_eqs + consumer_eqs, + opt=('topofuse', {'openmp': False})) + reference = Operator(reference_eqs, opt=('noop', {'openmp': False})) + + # Each consumer reads neighboring values of its newly written producer, + # requiring a producer and a consumer loop for each slab. + assert len(retrieve_iteration_tree(reference)) == 2 * len(slabs) + assert len(retrieve_iteration_tree(op)) == 2 * len(slabs) + + reference(time_M=1) + op(time_M=1) + + for v, v_ref in states: + np.testing.assert_allclose(v.data, v_ref.data) + + class TestMultiSubDomain: @pytest.mark.parametrize('opt', opts_tiling)