From bde23935a6713cc6a6f28d04f328f55e6e7b6bde Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 8 Oct 2026 13:28:08 +0100 Subject: [PATCH 1/4] compiler: Cache TimedAccess.distance on access geometry --- devito/ir/support/basic.py | 29 ++++++++++++++--- tests/test_ir.py | 64 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 85 insertions(+), 8 deletions(-) diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index 0bce1b9387..78947c0b08 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -339,7 +339,6 @@ def rebuild(self, **kwargs): return TimedAccess(access, mode, timestamp, ispace) - @memoized_meth def distance(self, other, logical=False): """ Compute the distance from ``self`` to ``other``. @@ -351,7 +350,19 @@ def distance(self, other, logical=False): logical : bool Compute a logical distance rather than true distance (i.e. ignoring degenerating indices created by size 1 buffers etc). + + Notes + ----- + The distance depends on the two accesses (including their Functions' + metadata), their IterationSpaces and `logical`, but never on the + timestamps or access modes. Hence it is cached on the former, and thus + shared by all TimedAccesses with the same geometry, such as the + timestamp-shifted copies created by `Scope.from_scopes`. """ + return _cached_distance(self.access, self.ispace, other.access, + other.ispace, logical) + + def _distance(self, other, logical=False): if isinstance(self.access, ComponentAccess) and \ isinstance(other.access, ComponentAccess) and \ self.access.index != other.access.index: @@ -1604,6 +1615,16 @@ def skippable_interval(d, ispace, it): return d is None or (d in ispace and not d._defines & it.dim._defines) +@memoized_func(scope='build') +def _cached_distance(access0, ispace0, access1, ispace1, logical): + """ + The geometry-keyed cache behind `TimedAccess.distance`. + """ + a0 = TimedAccess(access0, 'R', 0, ispace0) + a1 = TimedAccess(access1, 'R', 0, ispace1) + return a0._distance(a1, logical) + + def disjoint_subdims(a0, a1): """ Determine whether two TimedAccesses touch disjoint SubDimension regions @@ -1748,10 +1769,8 @@ def disjoint_test(e0, e1, d, it): if any(not i.is_Number for i in [p00, p01, p10, p11]): return False - i0 = sympy.Interval(min(p00, p01), max(p00, p01)) - i1 = sympy.Interval(min(p10, p11), max(p10, p11)) - - return not bool(i0.intersect(i1)) + # Closed intervals, hence touching endpoints still overlap + return bool(max(p00, p01) < min(p10, p11) or max(p10, p11) < min(p00, p01)) def degenerating_indices(i0, i1, function, logical=False): diff --git a/tests/test_ir.py b/tests/test_ir.py index 9a7b2e9fea..a4c8274e7e 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -16,7 +16,7 @@ from devito.ir.stree import stree_build from devito.ir.support.basic import ( AFFINE, IRREGULAR, REGULAR, IterationInstance, Relation, Scope, TimedAccess, Vector, - mocksym0, mocksym1 + _cached_distance, disjoint_test, mocksym0, mocksym1 ) from devito.ir.support.guards import GuardOverflow from devito.ir.support.space import ( @@ -28,8 +28,8 @@ from devito.tools import prod from devito.tools.data_structures import frozendict from devito.types import ( - Array, BlockDimension, Bundle, ConditionalDimension, CriticalRegion, CustomDimension, - Jump, Scalar, StencilDimension, Symbol + Array, BlockDimension, Bundle, ComponentAccess, ConditionalDimension, CriticalRegion, + CustomDimension, Jump, Scalar, StencilDimension, Symbol ) @@ -677,6 +677,64 @@ def test_subdimension_stencil_gap_function_order(self, orders): expected = S.ImaginaryUnit if order == 8 else S.Infinity assert a.distance(b) == b.distance(a) == (expected,) + def test_distance_cache(self): + grid = Grid(shape=(8, 8)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + g = Function(name='g', grid=grid) + ispace = IterationSpace([Interval(x), Interval(y)]) + a = TimedAccess(f[x + 1, y], 'W', 0, ispace) + b = TimedAccess(f[x, y], 'R', 1, ispace) + + cache = _cached_distance.cache + _cached_distance.clear() + try: + assert a.distance(b) == (1, 0) + assert len(cache) == 1 + + # Timestamps and access modes are not part of the geometry + assert a.rebuild(timestamp=5).distance(b.rebuild(mode='W')) == (1, 0) + assert len(cache) == 1 + + # Anything else is + rev = IterationSpace([Interval(x), Interval(y)], directions={x: Backward}) + shifted = IterationSpace([Interval(x, 0, 1), Interval(y)]) + c0 = a.rebuild(access=ComponentAccess(f[x, y], 0)) + c1 = a.rebuild(access=ComponentAccess(f[x, y], 1)) + pairs = [ + (b, a, False), + (a.rebuild(ispace=rev), b.rebuild(ispace=rev), False), + (a.rebuild(ispace=shifted), b, False), + (a.rebuild(access=g[x + 1, y]), b.rebuild(access=g[x, y]), False), + (a, b, True), + (c0, c1, False), + ] + for n, (i, j, logical) in enumerate(pairs, start=2): + i.distance(j, logical=logical) + assert len(cache) == n + + assert b.distance(a) == (-1, 0) + assert c0.distance(c1) == (S.ImaginaryUnit,) + + # Cleared at the end of each Operator build + Operator(Eq(f, 1)) + assert len(cache) == 0 + finally: + _cached_distance.clear() + + @pytest.mark.parametrize('offset,expected', [(4, True), (6, False), (8, False)]) + @pytest.mark.parametrize('swap', [False, True]) + def test_disjoint_test(self, offset, expected, swap): + grid = Grid(shape=(16,)) + x, = grid.dimensions + xl = SubDimension.left('xl', x, 4) + + # `12 - xl` spans [9, 12]; `xl + offset` spans [offset, offset + 3] + e0, e1 = 12 - xl, xl + offset + if swap: + e0, e1 = e1, e0 + assert disjoint_test(e0, e1, xl, Interval(xl)) is expected + class TestSpace: From c2e941f601a71b601a132d90c598ff2948248bbd Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 8 Oct 2026 17:23:37 +0100 Subject: [PATCH 2/4] compiler: Drop stamps from the disjoint_subdims_axis cache key --- devito/ir/support/basic.py | 4 +++- tests/test_ir.py | 36 ++++++++++++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index 78947c0b08..e9225f6059 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -1665,7 +1665,9 @@ def disjoint_subdims(a0, a1): f = a0.function.c0 space_order = f.space_order if isinstance(f, Function) else 0 - if disjoint_subdims_axis(e0, e1, d0, d1, it0, it1, space_order): + # The result doesn't depend on stamps; drop them from the cache key + if disjoint_subdims_axis(e0, e1, d0, d1, it0.reset(), it1.reset(), + space_order): return True return False diff --git a/tests/test_ir.py b/tests/test_ir.py index a4c8274e7e..1deee19f8f 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -16,7 +16,7 @@ from devito.ir.stree import stree_build from devito.ir.support.basic import ( AFFINE, IRREGULAR, REGULAR, IterationInstance, Relation, Scope, TimedAccess, Vector, - _cached_distance, disjoint_test, mocksym0, mocksym1 + _cached_distance, disjoint_subdims_axis, disjoint_test, mocksym0, mocksym1 ) from devito.ir.support.guards import GuardOverflow from devito.ir.support.space import ( @@ -25,7 +25,7 @@ ) from devito.ir.support.utils import detect_halo_writes from devito.symbolics import DefFunction, FieldFromPointer, uxreplace -from devito.tools import prod +from devito.tools import Stamp, prod from devito.tools.data_structures import frozendict from devito.types import ( Array, BlockDimension, Bundle, ComponentAccess, ConditionalDimension, CriticalRegion, @@ -722,6 +722,38 @@ def test_distance_cache(self): finally: _cached_distance.clear() + def test_disjoint_subdims_cache(self): + grid = Grid(shape=(24, 24)) + x, y = grid.dimensions + xl = SubDimension.left('xl', x, 4, separated=True) + xm = SubDimension.middle('xm', x, 4, 4, separated=True) + xm = xm._rebuild(thickness=(xl.ltkn, xm.rtkn)) + f = Function(name='f', grid=grid) + + def accesses(stamp): + il, im = Interval(xl, 0, 0, stamp), Interval(xm, 0, 0, stamp) + left = TimedAccess(f[xl, y], 'W', 0, IterationSpace([il, Interval(y)])) + middle = TimedAccess(f[xm, y], 'R', 1, IterationSpace([im, Interval(y)])) + return il, left, middle + + il0, left0, middle0 = accesses(Stamp()) + il1, left1, middle1 = accesses(Stamp()) + assert il0 != il1 and il0.reset() == il1.reset() + + _cached_distance.clear() + disjoint_subdims_axis.clear() + try: + assert left0.distance(middle0) == (S.ImaginaryUnit,) + assert len(disjoint_subdims_axis.cache) == 1 + + # Stamps aren't part of the per-axis cache key + assert left1.distance(middle1) == (S.ImaginaryUnit,) + assert len(_cached_distance.cache) == 2 + assert len(disjoint_subdims_axis.cache) == 1 + finally: + _cached_distance.clear() + disjoint_subdims_axis.clear() + @pytest.mark.parametrize('offset,expected', [(4, True), (6, False), (8, False)]) @pytest.mark.parametrize('swap', [False, True]) def test_disjoint_test(self, offset, expected, swap): From eb6e0d46e32c97a2a5f165b2dbb89936c7cb5543 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Fri, 9 Oct 2026 12:10:17 +0100 Subject: [PATCH 3/4] compiler: Cache read expansion in Scope.reads_smart_gen --- devito/ir/support/basic.py | 13 ++++++++++--- tests/test_ir.py | 38 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index e9225f6059..219134a96c 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -1197,9 +1197,7 @@ def reads_smart_gen(self, f): if isinstance(f, (Function, Temp, TempArray, TBArray)): for i in self.getreads(f): - expand = erange if uses_subdims(i) else extrema - - for j in expand(i.access): + for j in _expand_read(i.access, uses_subdims(i)): yield TimedAccess(j, i.mode, i.timestamp, i.ispace) else: @@ -1615,6 +1613,15 @@ def skippable_interval(d, ispace, it): return d is None or (d in ispace and not d._defines & it.dim._defines) +@memoized_func(scope='build') +def _expand_read(access, subdims): + """ + The accesses a read may perform once its StencilDimensions are resolved; + see `Scope.reads_smart_gen`. + """ + return erange(access) if subdims else extrema(access) + + @memoized_func(scope='build') def _cached_distance(access0, ispace0, access1, ispace1, logical): """ diff --git a/tests/test_ir.py b/tests/test_ir.py index 1deee19f8f..658e38e681 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -7,7 +7,9 @@ Constant, Dimension, Eq, Function, Grid, Inc, Operator, SubDimension, TimeDimension, TimeFunction, switchconfig ) -from devito.finite_differences.differentiable import IndexSum, LocalSum +from devito.finite_differences.differentiable import ( + IndexDerivative, IndexSum, LocalSum, Weights +) from devito.ir.cgen import ccode from devito.ir.clusters import Cluster, ClusterGroup from devito.ir.equations import LoweredEq @@ -16,7 +18,8 @@ from devito.ir.stree import stree_build from devito.ir.support.basic import ( AFFINE, IRREGULAR, REGULAR, IterationInstance, Relation, Scope, TimedAccess, Vector, - _cached_distance, disjoint_subdims_axis, disjoint_test, mocksym0, mocksym1 + _cached_distance, _expand_read, disjoint_subdims_axis, disjoint_test, mocksym0, + mocksym1 ) from devito.ir.support.guards import GuardOverflow from devito.ir.support.space import ( @@ -722,6 +725,37 @@ def test_distance_cache(self): finally: _cached_distance.clear() + def test_expand_read(self): + grid = Grid(shape=(8,)) + x, = grid.dimensions + f = Function(name='f', grid=grid) + h = StencilDimension('h', 0, 2) + w = Weights(name='w', dimensions=h, initvalue=[1., 2., 3.]) + + stencil = f[x + h] + plain = f[x] + bound = IndexDerivative(f[x + h]*w, {x: h}) + + cache = _expand_read.cache + _expand_read.clear() + try: + # Unbounded StencilDimension: every point, or just the extrema + assert _expand_read(stencil, True) == (f[x], f[x + 1], f[x + 2]) + assert tuple(_expand_read(stencil, False)) == (f[x], f[x + 2]) + assert len(cache) == 2 + + assert tuple(_expand_read(plain, True)) == (plain,) + assert tuple(_expand_read(plain, False)) == (plain, plain) + + # A StencilDimension bound by an IndexDerivative isn't expanded + assert tuple(_expand_read(bound, True)) == (bound,) + assert len(cache) == 5 + + _expand_read(stencil, True) + assert len(cache) == 5 + finally: + _expand_read.clear() + def test_disjoint_subdims_cache(self): grid = Grid(shape=(24, 24)) x, y = grid.dimensions From 9b4c1d34170d770c7d86ecf93f4ef975d054366e Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Mon, 12 Oct 2026 16:42:40 +0100 Subject: [PATCH 4/4] compiler: Build the fusion DAG as a transitive reduction when BREAK-free --- devito/passes/clusters/fusion.py | 68 ++++++++++++++++++++++++++++++++ tests/test_ir.py | 68 ++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+) diff --git a/devito/passes/clusters/fusion.py b/devito/passes/clusters/fusion.py index bbd4694614..e84d79cede 100644 --- a/devito/passes/clusters/fusion.py +++ b/devito/passes/clusters/fusion.py @@ -7,6 +7,7 @@ from devito.ir.support import ( InitArray, PrefetchUpdate, ReleaseLock, Scope, SyncArray, WaitLock, WithLock ) +from devito.ir.support.basic import Dependence from devito.symbolics import search from devito.tools import ( DAG, CacheInstances, as_tuple, flatten, frozendict, memoized_func, timed_pass @@ -51,6 +52,44 @@ def _fusion_hazards(scope0, scope1, prefix): return NO_HAZARD +def _maybe_break(cgroups, prefix): + """ + True if any two of `cgroups` may induce a `BREAK` hazard, False if that's + provably impossible. + + A `BREAK` requires an anti-dependence carried along `prefix`. All distinct + read/write access geometries that may occur in two different ClusterGroups + are checked, ignoring timestamps and `Scope.rules`, hence conservatively. + """ + if not prefix: + return False + + # Map each distinct access geometry to the ClusterGroups it occurs in + writes = defaultdict(dict) + for n, cg in enumerate(cgroups): + for f, v in cg.scope.writes.items(): + for w in v: + writes[f].setdefault((w.access, w.ispace), (w, set()))[1].add(n) + reads = defaultdict(dict) + for n, cg in enumerate(cgroups): + for f in writes: + for r in cg.scope.reads_smart_gen(f): + reads[f].setdefault((r.access, r.ispace), (r, set()))[1].add(n) + + for f, v in writes.items(): + for r, rn in reads[f].values(): + for w, wn in v.values(): + if len(rn | wn) == 1: + # Only ever within the same ClusterGroup + continue + + d = Dependence(r, w) + if not d.is_imaginary and d.cause & prefix: + return True + + return False + + class Keys(CacheInstances): """ @@ -297,10 +336,39 @@ def _build_dag(self, cgroups, prefix): """ A DAG representing the data dependences across the ClusterGroups within a given scope. + + If no `BREAK` hazard is possible, the DAG only encodes an ordering, so + just its transitive reduction is built. """ prefix = frozenset(i.dim for i in as_tuple(prefix)) dag = DAG(nodes=cgroups) + + if not _maybe_break(cgroups, prefix): + # Visit the ClusterGroups backwards, tracking reachability as + # bitsets, so that already ordered pairs needn't be analyzed. E.g., + # with edges `cg1 -> cg2` and `cg0 -> cg1`, `(cg0, cg2)` is skipped + reach = [0]*len(cgroups) + barrier = len(cgroups) + for n in reversed(range(len(cgroups))): + cg0 = cgroups[n] + if cg0.scope.has_barrier: + barrier = n + + r = 0 + for n1 in range(n + 1, len(cgroups)): + if r >> n1 & 1: + continue + + # Fences have the same effect as any hazard + cg1 = cgroups[n1] + if n1 >= barrier or _fusion_hazards(cg0.scope, cg1.scope, prefix): + dag.add_edge(cg0, cg1) + r |= 1 << n1 | reach[n1] + reach[n] = r + + return dag + for n, cg0 in enumerate(cgroups): # Track whether there is any fence between `cg0` and the current `cg1`. fenced = cg0.scope.has_barrier diff --git a/tests/test_ir.py b/tests/test_ir.py index 658e38e681..d4cbf0d0ac 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -27,6 +27,9 @@ NullInterval, null_ispace ) from devito.ir.support.utils import detect_halo_writes +from devito.passes.clusters.fusion import ( + BREAK, EDGE, Fusion, _fusion_hazards, _maybe_break +) from devito.symbolics import DefFunction, FieldFromPointer, uxreplace from devito.tools import Stamp, prod from devito.tools.data_structures import frozendict @@ -1838,6 +1841,71 @@ def test_local_sums(self): assert not cluster.rebuild(exprs=[Eq(f[x], 0)]).local_sums +class TestFusionDAG: + + @pytest.fixture + def setup(self): + d = Dimension(name='d') + x = Dimension(name='x') + ispace = IterationSpace([Interval(d), Interval(x)]) + f, g, h = [Function(name=n, dimensions=(d, x), shape=(4, 4)) for n in 'fgh'] + + def cgroups(*exprs): + return [ClusterGroup(Cluster(e, ispace), ispace) for e in exprs] + + return d, x, f, g, h, cgroups + + def test_no_break(self, setup): + d, x, f, g, h, cgroups = setup + + # Reads behind or at the write along `d` aren't carried anti-dependences + for write, hazard in [(f[d, x], EDGE), (f[d + 1, x], None)]: + cgs = cgroups(Eq(g[d, x], f[d, x]), Eq(write, 1)) + assert not _maybe_break(cgs, frozenset({d})) + assert _fusion_hazards(cgs[0].scope, cgs[1].scope, frozenset({d})) == hazard + + def test_break(self, setup): + d, x, f, g, h, cgroups = setup + + # Reading `f` ahead of the next write along `d` is + cgs = cgroups(Eq(g[d, x], f[d + 1, x]), Eq(f[d, x], 1), Eq(h[d, x], 1)) + assert _maybe_break(cgs, frozenset({d})) + assert _fusion_hazards(cgs[0].scope, cgs[1].scope, frozenset({d})) == BREAK + + # Hence the full DAG, `BREAK` edges included + dag = Fusion(toposort=True)._build_dag(cgs, [Interval(d)]) + assert set(dag.edges) == {(cgs[0], cgs[1]), (cgs[1], cgs[2])} + + def test_break_stencil_reads(self): + x = Dimension(name='x') + s = StencilDimension('s', 0, 1) + ispace = IterationSpace([Interval(x)]) + f, g = [Function(name=n, dimensions=(x,), shape=(4,)) for n in 'fg'] + + # The read `f[x + s]` expands to `f[x]` and `f[x + 1]`, the latter + # being ahead of the write along `x` + cgs = [ClusterGroup(Cluster(e, ispace), ispace) + for e in (Eq(g[x], f[x + s]), Eq(f[x + s], 1))] + assert _maybe_break(cgs, frozenset({x})) + + def test_transitive_reduction(self, setup): + d, x, f, g, h, cgroups = setup + + # `h` depends on both `f` and `g`, but `(cg0, cg2)` is implied by + # `cg0 -> cg1 -> cg2` + cgs = cgroups(Eq(f[d, x], 1), Eq(g[d, x], f[d, x]), + Eq(h[d, x], f[d, x] + g[d, x])) + dag = Fusion(toposort=True)._build_dag(cgs, [Interval(d)]) + assert set(dag.edges) == {(cgs[0], cgs[1]), (cgs[1], cgs[2])} + assert dag.all_downstreams(cgs[0]) == [cgs[1], cgs[2]] + + # A fence orders everything around it, even absent data dependences + cgs = cgroups(Eq(f[d, x], 1), Eq(Symbol(name='s'), CriticalRegion(True)), + Eq(g[d, x], 1)) + dag = Fusion(toposort=True)._build_dag(cgs, [Interval(d)]) + assert set(dag.edges) == {(cgs[0], cgs[1]), (cgs[1], cgs[2])} + + class TestGuards: def test_guard_overflow(self):