compiler: fix buffering corner case such as interp

Author: mloubout

.github/workflows/lint.yaml

@@ -88,6 +88,6 @@ jobs:
           for DOCKERFILE in docker/Dockerfile.*; \
             do \
             echo "    Linting $DOCKERFILE"; \
-            hadolint "$DOCKERFILE" \
+            hadolint --ignore DL3066 "$DOCKERFILE" \
             || exit 1; \
           done

devito/ir/equations/algorithms.py

@@ -365,6 +365,8 @@ def generate_conditionals(expr, input_expr, ordering):
                 cond = d.relation(cond, GuardFactor(d))
             conditionals[d] = cond
 
+    if not conditionals and not input_expr.implicit_dims:
+        return expr, conditionals
     # Merge conditionals when possible. E.g., if an implicit_dim shares
     # its parent Dimension with another ConditionalDimension, the two
     # conditions can be merged into a single guard.

devito/ir/stree/algorithms.py

@@ -169,7 +169,15 @@ def preprocess(clusters, options=None, **kwargs):
     for c in clusters:
         if c.is_halo_touch:
             hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs)
-            queue.append(c.rebuild(exprs=[], halo_scheme=hs))
+            # Peel syncs (e.g. a WaitLock) off, they must survive an unbound halo.
+            # Restrict the IterationSpace to the syncs' prefix so we don't carry
+            # inner Dimensions (e.g. a layer's VirtualDimension) that have no
+            # content left, which would yield empty guarded/iteration nodes
+            if c.syncs:
+                key = lambda d: any(d in s._defines for s in c.syncs)  # noqa: B023
+                ispace = c.ispace.prefix(key)
+                processed.append(c.rebuild(exprs=[], ispace=ispace))
+            queue.append(c.rebuild(exprs=[], syncs={}, halo_scheme=hs))
 
         elif c.is_dist_reduce:
             processed.append(c)

devito/ir/support/guards.py

@@ -65,7 +65,14 @@ class GuardFactor(Guard, CondEq, Pickable):
 
     __rargs__ = ('d',)
 
-    def __new__(cls, d, **kwargs):
+    def __new__(cls, *args, **kwargs):
+        if len(args) != 1:
+            # Reconstruction with relational args (e.g. via sympy `_subs`): the
+            # factor semantics no longer hold, so degrade to a plain relational
+            base = CondNe if issubclass(cls, CondNe) else CondEq
+            return base(*args, **kwargs)
+
+        d, = args
         assert d.is_Conditional
 
         obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
@@ -169,6 +176,11 @@ def __new__(cls, d, index, direction,
         except TypeError:
             pass
 
+        return cls._new(p0, p1, d, index, direction, d_min=d_min, d_max=d_max, **kwargs)
+
+    @classmethod
+    def _new(cls, p0, p1, d, index, direction, d_min=None, d_max=None):
+
         obj = super().__new__(cls, p0, p1, evaluate=False)
 
         obj.d = d
@@ -572,7 +584,9 @@ def pairwise_or(*guards):
 
 @_uxreplace_handle.register(BaseGuardBoundNext)
 def _(expr, args, kwargs):
-    return expr.func(expr.d, expr.index, expr.direction, **kwargs)
+    p0, p1 = args
+    return expr._new(p0, p1, expr.d, expr.index, expr.direction,
+                     **kwargs)
 
 
 @singledispatch

devito/passes/clusters/asynchrony.py

@@ -8,7 +8,7 @@
     WaitLock, WithLock, normalize_syncs
 )
 from devito.passes.clusters.utils import in_critical_region, is_memcpy
-from devito.symbolics import IntDiv, retrieve_terminals, uxreplace
+from devito.symbolics import IntDiv, retrieve_dimensions, retrieve_terminals, uxreplace
 from devito.tools import OrderedSet, is_integer, timed_pass
 from devito.types import CustomDimension, Lock, VirtualDimension
 
@@ -339,27 +339,29 @@ def _actions_from_update_memcpy(c, d, bounds, clusters, actions, sregistry):
 
     syncs = {d: [
         ReleaseLock(handle, target),
-        PrefetchUpdate(handle, target, tindex, function, findex, d,
+        PrefetchUpdate(handle, target, tindex, f, findex, d,
                        origin=e.rhs, gid=gid)
     ]}
     syncs = {**c.syncs, **syncs}
 
     pc = c.rebuild(exprs=expr, ispace=ispace, guards=guards, syncs=syncs)
 
-    # Since we're turning `e` into a prefetch, we need to:
-    # 1) attach a WaitLock SyncOp to the first Cluster accessing `target`
-    # 2) insert the prefetch Cluster right after the last Cluster accessing `target`
-    # 3) drop the original Cluster performing a memcpy-like fetch
-    n = clusters.index(c)
-    first = None
-    last = None
-    for c1 in clusters[n+1:]:
-        if target in c1.scope.reads:
-            if first is None:
-                first = c1
-            last = c1
-    assert first is not None
-    assert last is not None
+    # Wait before the first, prefetch after the last access to `target`, then
+    # drop the memcpy `c`. `c` may be toposorted amid the readers, so scan them
+    # all; count only reads over the streamed `pd`, not the buffer-init loop.
+    reads = [c1 for c1 in clusters
+             if c1 is not c and target in c1.scope.reads and pd in c1.ispace.itdims]
+    assert reads
+    first = reads[0]
+    last = reads[-1]
+
+    # Advance `last` past its loop nest so the prefetch follows it rather than
+    # splitting it, e.g. severing an interpolation's store from its point loop
+    nest_dims = set(last.ispace.itdims) - set(d._defines)
+    for c1 in clusters[clusters.index(last)+1:]:
+        if not nest_dims & set(c1.ispace.itdims):
+            break
+        last = c1
 
     actions[first].syncs[d].append(WaitLock(handle, target))
     actions[last].insert.append(pc)
@@ -433,7 +435,10 @@ def _(expr, dim, dir):
 
         3//2 - 1 + 0 = 0
     """
-    if expr.lhs._defines & dim._defines:
+    dims = retrieve_dimensions(expr.lhs)
+    assert len(dims) == 1
+    ldim = dims.pop()
+    if ldim._defines & dim._defines:
         if dir == 1:
             return expr + dir
         else:

devito/passes/clusters/buffering.py

@@ -311,7 +311,11 @@ def _optimize(self, clusters, descriptors):
                         processed.append(c)
                         continue
 
-                    key1 = lambda d: not d._defines & v.dim._defines  # noqa: B023
+                    # Lift only along the buffer's own spatial Dimensions, not
+                    # unrelated ones (e.g. sparse points of an indirect read)
+                    key1 = lambda d: (not d._defines & v.dim._defines and  # noqa: B023
+                                      any(d._defines & bd._defines  # noqa: B023
+                                          for bd in v.bdims))  # noqa: B023
                     dims = c.ispace.project(key1).itdims
                     ispace = c.ispace.lift(dims, key0())
                     processed.append(c.rebuild(ispace=ispace))
@@ -331,8 +335,8 @@ def _make_task_groups(self, descriptors):
             return stamps
 
         task_sets = (
-            [b for b, v in descriptors.items() if v.is_writeonly],
-            [b for b, v in descriptors.items() if v.is_readonly],
+            [b for b, v in descriptors.items() if any(vb.is_writeonly for vb in v)],
+            [b for b, v in descriptors.items() if any(vb.is_readonly for vb in v)],
         )
 
         for tasks in task_sets:
@@ -519,30 +523,37 @@ def itdims(self):
     def ispace(self):
         # The IterationSpace within which the buffer will be accessed
 
-        # NOTE: The `key` is to avoid Clusters including `f` but not directly
-        # using it in an expression, such as HaloTouch Clusters
-        def key(c):
-            bufferdim = any(i in c.ispace.dimensions for i in self.bdims)
-            xd_only = all(d._defines & self.xd._defines for d in c.ispace.dimensions)
-            return bufferdim or xd_only
-
         ispaces = set()
+        indirect = set()
         for c in self.clusters:
-            if not key(c):
-                continue
-
-            # Skip wild clusters (e.g. HaloTouch Clusters)
+            # Skip wild clusters (e.g. HaloTouch Clusters), which include `f`
+            # but do not directly use it in an expression
             if c.is_wild:
                 continue
 
-            # Iterations space and buffering dims
-            edims = [d for d in self.bdims if d not in c.ispace.dimensions]
-            if not edims:
-                ispaces.add(c.ispace)
-            else:
-                # Add all missing buffering dimensions and reorder to
-                # avoid duplicates with different ordering
-                ispaces.add(c.ispace.insert(self.dim, edims).reorder())
+            bufferdim = any(i in c.ispace.dimensions for i in self.bdims)
+            xd_only = all(d._defines & self.xd._defines for d in c.ispace.dimensions)
+
+            if bufferdim or xd_only:
+                # `c` iterates (at least some of) the buffer's spatial dims
+                edims = [d for d in self.bdims if d not in c.ispace.dimensions]
+                if not edims:
+                    ispaces.add(c.ispace)
+                else:
+                    # Add all missing buffering dimensions and reorder to
+                    # avoid duplicates with different ordering
+                    ispaces.add(c.ispace.insert(self.dim, edims).reorder())
+            elif ((self.f in c.scope.reads or self.f in c.scope.writes) and
+                  self.dim.root in c.ispace.dimensions):
+                # `c` accesses `f` indirectly (e.g. interpolation), so it doesn't
+                # iterate `bdims`; span the buffer's own Dimensions instead
+                tispace = c.ispace.project(lambda i: i._defines & self.dim.root._defines)
+                indirect.add(tispace.insert(self.dim, list(self.bdims)))
+
+        # Indirect accessors define the ispace only for a read-only streamed
+        # buffer, where nothing iterates the buffer's own Dimensions directly
+        if not ispaces:
+            ispaces = indirect
 
         if len(ispaces) > 1:
             # Best effort to make buffering work in the presence of multiple
@@ -663,7 +674,11 @@ def last_idx(self):
         mapper = {}
         func = vmax if self.is_forward_buffering else vmin
         for c in self.lastwrite + self.firstread:
-            indices = extract_indices(self.f, self.dim, [c])
+            # Consider all Clusters sharing `c`'s guards, so the leading edge is
+            # found even when `f` is accessed at different offsets across them
+            # (e.g. `f` and `f.forward` in separate Eqs)
+            group = [c1 for c1 in self.clusters if c1.guards == c.guards]
+            indices = extract_indices(self.f, self.dim, group)
             idx = func(*[Vector(i) for i in indices])[0]
             mapper[c] = idx
 

tests/test_gpu_common.py

@@ -1419,7 +1419,7 @@ def test_streaming_complete(self):
         assert len([i for i in FindSymbols().visit(op2) if i.is_Array]) == 9 - diff
         # The streaming/buffering cannot be fused since the guards of
         # vas and vasb/vasf are not compatible.
-        assert len(op3._func_table) == 14 - diff
+        assert len(op3._func_table) == 12 - diff
         assert len([i for i in FindSymbols().visit(op3) if i.is_Array]) == 9 - diff
 
         op0.apply(time_m=15, time_M=35, save_shift=0)