JDBetteridge/compression

devito/passes/clusters/asynchrony.py

@@ -102,7 +102,14 @@ def callback(self, clusters, prefix):
             protected = self._schedule_waitlocks(c0, d, clusters, locks, syncs)
             self._schedule_withlocks(c0, d, protected, locks, syncs)
 
-        processed = [c.rebuild(syncs={**c.syncs, **syncs[c]}) for c in clusters]
+        # Tasking may run again after prefetching has attached input waits.
+        # Preserve them when adding locks for asynchronous output readers.
+        processed = []
+        for c in clusters:
+            waits = {d: [s for s in ops if isinstance(s, WaitLock)]
+                     for d, ops in c.syncs.items()}
+            ops = normalize_syncs(waits, {**c.syncs, **syncs[c]})
+            processed.append(c.rebuild(syncs=ops))
 
         return processed
 

devito/passes/clusters/blocking.py

@@ -101,11 +101,16 @@ def _process_fatd(self, clusters, level, prefix=None):
 
         return super()._process_fatd(clusters, level, prefix)
 
-    def _has_data_reuse(self, cluster):
+    def _has_data_reuse(self, cluster, dims=None):
         # A sufficient condition for the existence of data reuse in `cluster`
         # is that the same Function is accessed twice at the same memory location,
         # which translates into the existence of any Relation across Indexeds
-        if any(r.function.is_AbstractFunction for r in cluster.scope.r_gen()):
+        for r in cluster.scope.r_gen():
+            if not r.function.is_AbstractFunction:
+                continue
+            # E.g., `u.forward = u + 1` has no reuse along parallel Dimensions
+            if dims is not None and all(r.distance_mapper.get(d) == 0 for d in dims):
+                continue
             return True
         if search(cluster.exprs, IndexSum):
             return True
@@ -208,9 +213,11 @@ def callback(self, clusters, prefix):
                     return clusters
 
                 properties = c.properties.block(d)
+                dims = [i for i in c.ispace.itdims
+                        if c.properties.is_parallel_relaxed(i)]
 
                 if any(self._has_short_trip_count(i) for i in c.ispace.itdims) or \
-                   not self._has_data_reuse(c):
+                   not self._has_data_reuse(c, dims):
                     properties = properties.notune(d)
 
             elif self._has_data_reuse(c):

devito/passes/clusters/buffering.py

@@ -1056,7 +1056,8 @@ def _select_buffer(f, dim, guard, cgroup, xds, async_degree, sregistry, callback
         assert len(buffers) == 1, "Unexpected form of multi-level buffering"
         buffer, = buffers
         xd = buffer.indices[dim]
-        extra_kwargs = {'is_autopaddable': buffer.is_autopaddable}
+        extra_kwargs = {'is_autopaddable': buffer.is_autopaddable,
+                        'buffer': buffer}
     else:
         size = infer_buffer_size(f, dim, cgroup)
         if async_degree is not None:

tests/test_gpu_openacc.py

@@ -186,6 +186,36 @@ def test_short_multi_tile_keeps_outer_dim_on_device(self):
             assert len(iters) == len(v)
             assert all(i.step == j for i, j in zip(iters, v, strict=True))
 
+    @pytest.mark.parametrize('shifted', [False, True])
+    def test_short_multi_tile_time_update(self, shifted):
+        grid = Grid(shape=(8, 8, 8))
+        t = grid.stepping_dim
+        x, y, z = grid.dimensions
+        u = TimeFunction(name='u', grid=grid)
+
+        source = u[t, x + 1, y, z] if shifted else u
+        op = Operator(Eq(u.forward, source + 1), name='short_tile_time_update',
+                      platform='nvidiaX', language='openacc',
+                      opt=('advanced', {'par-tile': ((16, 4), (8, 8)),
+                                        'blocklevels': 1, 'blockinner': True,
+                                        'blockrelax': 'device-aware'}))
+
+        root = 'y0_blk0' if shifted else 'x0_blk0'
+        bns, _ = assert_blocking(op, {root})
+        iters = FindNodes(Iteration).visit(bns[root])
+        steps = [i.step for i in iters if i.dim.is_Block and i.dim._depth == 1]
+        assert steps == ([4, 16] if shifted else [4, 4, 16])
+
+        tree, = retrieve_iteration_tree(op)
+        assert tree[0].dim is grid.time_dim
+        if shifted:
+            assert tree[1].dim is x
+            assert tree[1].limits == (x.symbolic_min, x.symbolic_max, 1)
+        else:
+            assert all(i.dim.is_Block for i in tree[1:])
+        assert tree[1].pragmas[0].ccode.value ==\
+            'acc parallel loop tile(16,4,4) present(u)'
+
     def test_multi_tile_blocking_structure(self):
         grid = Grid(shape=(8, 8, 8))