compiler: Reject host-evaluated guards reading device-written Functions

devito/core/gpu.py

@@ -3,9 +3,9 @@
 import numpy as np
 
 from devito.core.operator import CoreOperator, CustomOperator, ParTile
-from devito.exceptions import InvalidOperator
+from devito.exceptions import CompilationError, InvalidOperator
 from devito.operator.operator import rcompile
-from devito.passes import is_on_device, stream_dimensions
+from devito.passes import is_on_device, needs_transfer, stream_dimensions
 from devito.passes.clusters import (
     Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse,
     memcpy_prefetch, optimize_pows, tasking
@@ -15,6 +15,7 @@
     DeviceAccTarget, DeviceCXXOmpTarget, DeviceOmpTarget, check_stability, hoist_prodders,
     linearize, mpiize, pthreadify, relax_incr_dimensions
 )
+from devito.symbolics import retrieve_functions
 from devito.tools import as_tuple, timed_pass
 
 __all__ = [
@@ -142,6 +143,37 @@ def _normalize_gpu_fit(cls, oo, **kwargs):
             else:
                 return as_tuple(cls.GPU_FIT)
 
+    @classmethod
+    def _lower_clusters(cls, expressions, **kwargs):
+        clusters = super()._lower_clusters(expressions, **kwargs)
+
+        # The host copy of a Function written on the device is stale, so a guard
+        # evaluated on the host must not read it. A guard is evaluated on the host
+        # if it is outside the device kernels, e.g. `if (T[0][2][0] < 1.5F)` around
+        # a kernel launch, or if its Cluster runs on the host, e.g. a copy-back to
+        # a host Function
+        gpu_fit = kwargs['options']['gpu-fit']
+        offloaded = {c for c in clusters
+                     if c.properties.is_parallel_relaxed(c.ispace.itdims) and
+                     all(is_on_device(f, gpu_fit) for f in c.scope.writes)}
+        writes = {f for c in offloaded
+                  for f in c.scope.writes if needs_transfer(f, gpu_fit)}
+        for c in clusters:
+            for d, guard in c.guards.items():
+                dims = c.ispace.prefix(lambda i: d in i._defines).itdims  # noqa: B023
+                if c in offloaded and c.properties.is_parallel_relaxed(dims):
+                    continue
+
+                stale = writes.intersection(retrieve_functions(guard, deep=True))
+                if stale:
+                    names = ", ".join(sorted(f.name for f in stale))
+                    raise CompilationError(
+                        f"Unsupported guard `{guard}`: it is evaluated on the host, "
+                        f"but it reads `{names}`, updated on the device"
+                    )
+
+        return clusters
+
     @classmethod
     def _rcompile_wrapper(cls, **kwargs0):
         options0 = kwargs0.pop('options')

tests/test_gpu_common.py

@@ -12,7 +12,7 @@
     switchconfig, switchenv
 )
 from devito.arch import Cpu64, Device, get_cpu_info, get_gpu_info
-from devito.exceptions import InvalidArgument
+from devito.exceptions import CompilationError, InvalidArgument
 from devito.ir import (
     Conditional, Expression, FindNodes, FindSymbols, Section, retrieve_iteration_tree
 )
@@ -371,6 +371,29 @@ def test_reduction_many_dims(self):
         tree, = retrieve_iteration_tree(op1)
         assert 'collapse(3) reduction(+:s)' in str(tree[1].pragmas[0])
 
+    @pytest.mark.parametrize('opt', ['noop', 'advanced'])
+    def test_guard_reads_device_data(self, opt):
+        grid = Grid(shape=(4, 4))
+        x, y = grid.dimensions
+        time = grid.time_dim
+
+        T = TimeFunction(name='T', grid=grid, space_order=0, time_order=0)
+        f0 = Function(name='f0', grid=grid, space_order=0)
+        f1 = TimeFunction(name='f1', grid=grid, space_order=0, time_order=0, save=4)
+
+        # The guard is evaluated within the kernel, hence it sees the updated `T`
+        cy = ConditionalDimension(name='cy', parent=y, condition=T[0, x, y] < 1.5)
+        op = Operator([Eq(T, time), Inc(f0, T + 1, implicit_dims=cy)],
+                      opt=opt, name='GuardDevice')
+        op.apply(time_M=3)
+        assert np.all(f0.data == 3)
+
+        # The guard would be evaluated on the host, where `T` is stale
+        ct = ConditionalDimension(name='ct', parent=time, condition=T[0, 2, 0] < 1.5)
+        with pytest.raises(CompilationError, match='evaluated on the host'):
+            Operator([Eq(T, time), Eq(f1, T + 1, implicit_dims=ct)],
+                     opt=opt, name='GuardHost')
+
 
 class Bundle(SubDomain):
     """