compiler: fix buffering corner case such as interp

Author: mloubout

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)

devito/passes/clusters/asynchrony.py

@@ -330,20 +330,22 @@ def _actions_from_update_memcpy(c, d, bounds, clusters, actions, sregistry):
 
     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)

devito/passes/clusters/buffering.py

@@ -303,7 +303,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))
@@ -465,30 +469,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
@@ -609,7 +620,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