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

@@ -1153,7 +1153,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:

devito/passes/clusters/cse.py

@@ -98,7 +98,9 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     min_cost = options['cse-min-cost']
     mode = options['cse-algo']
     try:
-        dtype = np.promote_types(options['scalar-min-type'], cluster.dtype).type
+        dtype = cluster.dtype
+        if not np.issubdtype(dtype, np.integer):
+            dtype = np.promote_types(options['scalar-min-type'], dtype).type
     except TypeError:
         dtype = cluster.dtype
 

devito/passes/iet/languages/openacc.py

@@ -2,8 +2,8 @@
 
 from devito.arch import AMDGPUX, NVIDIAX
 from devito.ir import (
-    Block, BusyWaitCall, Call, DeviceCall, DummyExpr, EntryFunction, FindSymbols, List,
-    ParallelTree, Pragma, Return, make_callable
+    Block, BusyWaitCall, Call, Conditional, DeviceCall, DummyExpr, EntryFunction,
+    FindSymbols, List, ParallelTree, Pragma, Return, make_callable
 )
 from devito.passes import is_on_device, needs_transfer
 from devito.passes.iet.definitions import DeviceAwareDataManager
@@ -226,7 +226,8 @@ def place_devptr(self, iet, **kwargs):
             return iet, {}
 
         symbols = FindSymbols('basics').visit(iet)
-        functions = [f for f in symbols if needs_transfer(f, self.gpu_fit)]
+        functions = [f for f in FindSymbols().visit(iet)
+                     if needs_transfer(f, self.gpu_fit)]
 
         efuncs = []
         calls = []
@@ -248,7 +249,9 @@ def place_devptr(self, iet, **kwargs):
 
             ret = Return(ctdp)
 
-            body = List(body=[init, dpf, castf, ret])
+            # Empty allocations have no present device address to discover.
+            nbytes = FieldFromPointer(f._C_field_nbytes, f._C_symbol)
+            body = List(body=[init, Conditional(nbytes > 0, dpf), castf, ret])
 
             name = self.sregistry.make_name(prefix='map_device_ptr')
             efuncs.append(make_callable(name, body, retval=hp))

devito/passes/iet/linearization.py

@@ -391,6 +391,11 @@ def linearize_transfers(iet, sregistry=None, **kwargs):
 
                 imask = [(wildcard, size)]
 
+        # Preserve snapshot metadata even when the flattened mask is empty:
+        # `_rebuild()` without positional arguments restores the original items.
+        if isinstance(imask0, IMask):
+            imask = IMask(*imask, **imask0.__dict__)
+
         rebuilt = n._rebuild(imask=imask)
 
         if exprs:

tests/test_gpu_common.py

@@ -14,7 +14,7 @@
 from devito.arch import Cpu64, Device, get_cpu_info, get_gpu_info
 from devito.exceptions import InvalidArgument
 from devito.ir import (
-    Call, Conditional, Expression, FindNodes, FindSymbols, Section,
+    Call, Conditional, Expression, FindNodes, FindSymbols, Section, Transfer,
     retrieve_iteration_tree
 )
 from devito.passes.iet.languages.openmp import OmpIteration
@@ -285,7 +285,9 @@ def test_linearize_transfer_no_overflow(self):
 
         # The transfer bound is a product of the four `size[i]`, each cast to
         # `long`; the multiplication is thus carried out in 64-bit arithmetic
-        for transfer in op.body.maps + op.body.unmaps:
+        transfers = FindNodes(Transfer).visit(op.body.maps + op.body.unmaps)
+        assert len(transfers) == 3
+        for transfer in transfers:
             code = transfer.ccode.value
             for i in range(4):
                 assert f'(long)(u_vec->size[{i}])' in code
@@ -453,7 +455,10 @@ def test_tasking_fused(self):
         assert len(retrieve_iteration_tree(op)) == 3
         locks = [i for i in FindSymbols().visit(op) if isinstance(i, Lock)]
         assert len(locks) == 1  # Only 1 because it's only `tmp` that needs protection
-        assert len(op._func_table) == 5 + int(configuration['language'] == 'openacc')
+        if configuration['language'] == 'openacc':
+            assert len(op._func_table) == 7
+        else:
+            assert len(op._func_table) == 5
         exprs = FindNodes(Expression).visit(op._func_table['copy_to_host0'].root)
         b = 17 if configuration['language'] == 'openacc' else 16  # No `qid` w/ OMP
         assert str(exprs[b]) == 'const int deviceid = sdata0->deviceid;'
@@ -508,7 +513,10 @@ def test_tasking_unfused_two_locks(self):
         body = sections[2].body[0].body[0]
         assert str(body.body[0]) == 'release_lock0(lock0);'
         assert str(body.body[1]) == 'activate0(time,sdata0);'
-        assert len(op._func_table) == 5 + int(configuration['language'] == 'openacc')
+        if configuration['language'] == 'openacc':
+            assert len(op._func_table) == 7
+        else:
+            assert len(op._func_table) == 5
         calls = FindNodes(Call).visit(op._func_table['copy_to_host0'].root)
         stores = [str(i) for i in calls if i.name == '__atomic_store_n']
         assert stores[0] == '__atomic_store_n(&lock0[0],1,__ATOMIC_RELEASE);'
@@ -563,7 +571,10 @@ def test_tasking_forcefuse(self):
         assert str(body.body[1]) == 'sdata0[0].time = time;'
         assert str(body.body[2]) == \
             '__atomic_store_n(&sdata0[0].flag,2,__ATOMIC_RELEASE);'
-        assert len(op._func_table) == 5 + int(configuration['language'] == 'openacc')
+        if configuration['language'] == 'openacc':
+            assert len(op._func_table) == 7
+        else:
+            assert len(op._func_table) == 5
         exprs = FindNodes(Expression).visit(op._func_table['copy_to_host0'].root)
         b = 21 if configuration['language'] == 'openacc' else 20  # No `qid` w/ OMP
         assert exprs[b].write is u
@@ -636,7 +647,10 @@ def test_tasking_multi_output(self):
         body = op1._func_table['activate0'].root.body
         assert str(body.body[-1]) == \
             '__atomic_store_n(&sdata0[wi0].flag,2,__ATOMIC_RELEASE);'
-        assert len(op1._func_table) == 5 + int(configuration['language'] == 'openacc')
+        if configuration['language'] == 'openacc':
+            assert len(op1._func_table) == 7
+        else:
+            assert len(op1._func_table) == 5
         exprs = FindNodes(Expression).visit(op1._func_table['copy_to_host0'].root)
         b = 21 if configuration['language'] == 'openacc' else 20  # No `qid` w/ OMP
         assert exprs[b].write is usave
@@ -695,8 +709,9 @@ def test_streaming_basic(self, opt, ntmps):
         op = Operator(eqn, opt=opt)
 
         # Check generated code
+        # OpenACC maps both the Function and its staging array.
         if configuration['language'] == 'openacc':
-            assert len(op._func_table) == 10
+            assert len(op._func_table) == 11
         else:
             assert len(op._func_table) == 8
         assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == ntmps
@@ -771,7 +786,7 @@ def test_streaming_two_buffers(self, opt, ntmps):
         # Check generated code
         arrays = [i for i in FindSymbols().visit(op) if i.is_Array]
         if configuration['language'] == 'openacc':
-            assert len(op._func_table) == 10
+            assert len(op._func_table) == 11
             assert len(arrays) == ntmps
         else:
             assert len(op._func_table) == 8
@@ -982,7 +997,7 @@ def test_streaming_multi_input(self, opt, ntmps):
 
         # Check generated code
         if configuration['language'] == 'openacc':
-            assert len(op1._func_table) == 10
+            assert len(op1._func_table) == 11
         else:
             assert len(op1._func_table) == 8
         assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == ntmps
@@ -1074,7 +1089,7 @@ def test_streaming_postponed_deletion(self, opt, ntmps):
 
         # Check generated code
         if configuration['language'] == 'openacc':
-            assert len(op1._func_table) == 10
+            assert len(op1._func_table) == 11
         else:
             assert len(op1._func_table) == 8
         assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == ntmps
@@ -1160,7 +1175,7 @@ def test_composite_buffering_tasking_multi_output(self):
         # loop nests, but they eventually get mapped to the same pair of efuncs,
         # since devito attempts to maximize code reuse
         if configuration['language'] == 'openacc':
-            assert len(op1._func_table) == 9
+            assert len(op1._func_table) == 10
         else:
             assert len(op1._func_table) == 7
 
@@ -1383,7 +1398,7 @@ def test_save_multi_output(self):
         # are identical, so they get mapped to the same efuncs (init + copy)
         # There also are two extra functions to allocate and free arrays
         if configuration['language'] == 'openacc':
-            assert len(op._func_table) == 9
+            assert len(op._func_table) == 10
         else:
             assert len(op._func_table) == 7
 
@@ -1442,7 +1457,7 @@ def test_save_w_nonaffine_time(self):
         # We just check the generated code here
         assert len([i for i in FindSymbols().visit(op) if isinstance(i, Lock)]) == 1
         if configuration['language'] == 'openacc':
-            assert len(op._func_table) == 9
+            assert len(op._func_table) == 11
         else:
             assert len(op._func_table) == 7
 
@@ -1502,7 +1517,7 @@ def test_streaming_w_shifting(self, opt, ntmps):
 
         # Check generated code
         if configuration['language'] == 'openacc':
-            assert len(op._func_table) == 10
+            assert len(op._func_table) == 11
         else:
             assert len(op._func_table) == 8
         assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == ntmps
@@ -1560,16 +1575,21 @@ def test_streaming_complete(self):
                                   {'npthreads': 1}))
 
         # Check generated code
-        diff = int(configuration['language'] == 'openmp')
-        extra = int(configuration['language'] == 'openacc')
-        assert len(op1._func_table) == 14 - diff + extra
-        assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == 9 - diff
-        assert len(op2._func_table) == 14 - diff + extra
-        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) == 12 - diff + extra
-        assert len([i for i in FindSymbols().visit(op3) if i.is_Array]) == 9 - diff
+        if configuration['language'] == 'openacc':
+            expected_funcs = (16, 16, 14)
+            expected_arrays = 9
+        elif configuration['language'] == 'openmp':
+            expected_funcs = (13, 13, 11)
+            expected_arrays = 8
+        else:
+            expected_funcs = (14, 14, 12)
+            expected_arrays = 9
+        for op, nfuncs in zip((op1, op2, op3), expected_funcs, strict=True):
+            assert len(op._func_table) == nfuncs
+            arrays = [i for i in FindSymbols().visit(op) if i.is_Array]
+            assert len(arrays) == expected_arrays
 
         op0.apply(time_m=15, time_M=35, save_shift=0)
         op1.apply(time_m=15, time_M=36, save_shift=0, u=u1)

tests/test_gpu_openacc.py

@@ -189,6 +189,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))
 
@@ -393,7 +423,7 @@ def test_basic(self, mode):
 
         # Make sure we've indeed generated OpenACC+MPI code
         assert 'acc parallel' in str(op)
-        assert len(op._func_table) == 4
+        assert len(op._func_table) == 5
 
         op(time_M=1)