compiler: order the async task handshake

The lock and flag an asynchronous task synchronises on are plain volatile ints,
written by both threads with no ordering imposed.  `volatile` guarantees the
loads are re-issued; it does not order stores, so on a weakly ordered target
they can be observed out of order and the handshake loses an update:

    compute: lock0[0] = 0;      (release_lock0)  -- observed late
             sdata0->flag = 2;  (activate0)      -- observed first
    task:    sees the request, delivers lock0[0] = 2, sets flag = 1
    compute: the late lock0[0] = 0 lands, wiping the delivery
             the next release_lock0 waits for a 2 nobody will write again

Both threads then spin forever: the compute waiting for data whose request it
has already spent, the task waiting to be asked.  `lock == 0 && flag == 1` is
reachable only this way -- the task sets the lock before the flag, so a
completed cycle must leave the lock at 2 -- and that is the state a stalled run
sits in.

Fenced on both sides: a release before the flag that publishes a request or a
completion, an acquire before reading what either stands for.  `__atomic_thread_fence`
is a builtin, valid in C and C++ alike, so the device targets that render
through CXXPrinter are unaffected; declaring the two objects `_Atomic` would
have been the standard-clean alternative but that is not valid C++ under g++.

Found through a streaming-checkpoint FWI gradient on arm64, where it stalled one
worker in six within minutes and, in the same runs, had CvxCompress trip its own
assertion on a buffer the task never filled.  x86's store ordering hides it.

The tests index the two callables positionally, so they move with the fences.

Author: mloubout

devito/ir/iet/nodes.py

@@ -18,7 +18,7 @@
     AFFINE, INBOUND, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, SEQUENTIAL,
     VECTORIZED, Forward, PrefetchUpdate, Property, WithLock
 )
-from devito.symbolics import CallFromPointer, ListInitializer
+from devito.symbolics import CallFromPointer, ListInitializer, Macro
 from devito.tools import (
     Signer, as_tuple, ctypes_to_cstr, filter_ordered, filter_sorted, flatten
 )
@@ -61,6 +61,7 @@
     'Section',
     'Switch',
     'SyncSpot',
+    'ThreadFence',
     'TimedList',
     'Transfer',
     'Using',
@@ -1302,6 +1303,31 @@ def size(self):
         return self.ispace.size
 
 
+class ThreadFence(Call):
+
+    """
+    A memory fence, ordering a thread's accesses either side of it.
+
+    Threads exchanging work through a shared flag need memory ordering on both
+    sides: a release before publishing the flag, and an acquire after observing it.
+    This ensures the state associated with the flag is visible before it is used.
+    Without this ordering, updates may be observed in the wrong order and lost.
+
+    `__atomic_thread_fence` is a compiler builtin, so this renders the same in C
+    and in C++ and needs no header.
+    """
+
+    def __init__(self, order):
+        assert order in ('acquire', 'release')
+        super().__init__('__atomic_thread_fence',
+                         [Macro(f'__ATOMIC_{order.upper()}')])
+        self._order = order
+
+    @property
+    def order(self):
+        return self._order
+
+
 class Prodder(Call):
 
     """

devito/passes/iet/asynchrony.py

@@ -7,7 +7,7 @@
 from devito.ir import (
     AsyncCall, AsyncCallable, BlankLine, Call, Callable, Conditional, DummyEq, DummyExpr,
     EntryFunction, FindNodes, FindSymbols, Increment, Iteration, List, PointerCast,
-    Return, ThreadCallable, Transformer, While, make_callable, maybe_alias
+    Return, ThreadCallable, ThreadFence, Transformer, While, make_callable, maybe_alias
 )
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.engine import iet_pass
@@ -86,6 +86,11 @@ def _lower_async_objs(iet, tracker=None, sregistry=None, **kwargs):
                 arguments.append(i)
         activation.extend([DummyExpr(FieldFromComposite(i.base, sdata[d]), i)
                            for i in arguments])
+
+        # Publishing the request must happen last. Everything the thread depends on
+        # must be visible before this flag is set. `volatile` does not provide that
+        # ordering, so weakly ordered targets could otherwise observe stale state.
+        activation.append(ThreadFence('release'))
         activation.append(
             DummyExpr(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 2)
         )
@@ -138,12 +143,17 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
     tbase = threads.indexed
 
     # Prepend the SharedData fields available upon thread activation
-    preactions = [DummyExpr(i, FieldFromPointer(i.base, sbase)) for i in ncfields]
+    preactions = [ThreadFence('acquire')]
+    preactions.extend(DummyExpr(i, FieldFromPointer(i.base, sbase))
+                      for i in ncfields)
     preactions.append(BlankLine)
 
     # Append the flag reset
     postactions = [List(body=[
         BlankLine,
+        # Whatever the task produced -- data in a buffer, a lock handed back --
+        # has to be visible before the flag says it is done.
+        ThreadFence('release'),
         DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 1)
     ])]
 

devito/passes/iet/orchestration.py

@@ -7,7 +7,7 @@
 from devito.exceptions import CompilationError
 from devito.ir.iet import (
     AsyncCall, AsyncCallable, BlankLine, Block, BusyWait, Call, Callable, Conditional,
-    DummyExpr, List, SyncSpot, Transformer, derive_parameters, make_callable
+    DummyExpr, List, SyncSpot, ThreadFence, Transformer, derive_parameters, make_callable
 )
 from devito.ir.iet.visitors import Visitor
 from devito.ir.support import (
@@ -53,7 +53,11 @@ def _make_waitlock(self, iet, sync_ops, *args):
     def _make_releaselock(self, iet, sync_ops, *args):
         pre = []
         pre.append(BusyWait(Or(*[CondNe(s.handle, 2) for s in sync_ops])))
+        pre.append(ThreadFence('acquire'))
         pre.extend(DummyExpr(s.handle, 0) for s in sync_ops)
+        # Hand the lock back before anything that follows can be observed --
+        # in particular before the request that asks the thread to refill it.
+        pre.append(ThreadFence('release'))
 
         name = self.sregistry.make_name(prefix="release_lock")
         parameters = derive_parameters(pre, ordering='canonical')

tests/test_gpu_common.py

@@ -413,11 +413,16 @@ def test_tasking_in_isolation(self, opt):
         assert str(sections[0].body[0].body[0].body[0].body[0]) == 'while(lock0[0] == 0);'
         body = op._func_table['release_lock0'].root.body
         assert str(body.body[0].condition) == 'Ne(lock0[0], 2)'
-        assert str(body.body[1]) == 'lock0[0] = 0;'
+        # An acquire fence pairs with the thread's release before the lock is
+        # read back, and a release fence publishes it before the next request
+        assert 'atomic_thread_fence' in str(body.body[1])
+        assert str(body.body[2]) == 'lock0[0] = 0;'
+        assert 'atomic_thread_fence' in str(body.body[3])
         body = op._func_table['activate0'].root.body
         assert str(body.body[0].condition) == 'Ne(sdata0[0].flag, 1)'
         assert str(body.body[1]) == 'sdata0[0].time = time;'
-        assert str(body.body[2]) == 'sdata0[0].flag = 2;'
+        assert 'atomic_thread_fence' in str(body.body[2])
+        assert str(body.body[3]) == 'sdata0[0].flag = 2;'
 
         op.apply(time_M=nt-2)
 
@@ -529,12 +534,15 @@ def test_tasking_forcefuse(self):
                 'while(lock0[0] == 0 || lock1[0] == 0);')  # Wait-lock
         body = op._func_table['release_lock0'].root.body
         assert str(body.body[0].condition) == 'Ne(lock0[0], 2) | Ne(lock1[0], 2)'
-        assert str(body.body[1]) == 'lock0[0] = 0;'  # Set-lock
-        assert str(body.body[2]) == 'lock1[0] = 0;'  # Set-lock
+        assert 'atomic_thread_fence' in str(body.body[1])
+        assert str(body.body[2]) == 'lock0[0] = 0;'  # Set-lock
+        assert str(body.body[3]) == 'lock1[0] = 0;'  # Set-lock
+        assert 'atomic_thread_fence' in str(body.body[4])
         body = op._func_table['activate0'].root.body
         assert str(body.body[0].condition) == 'Ne(sdata0[0].flag, 1)'  # Wait-thread
         assert str(body.body[1]) == 'sdata0[0].time = time;'
-        assert str(body.body[2]) == 'sdata0[0].flag = 2;'
+        assert 'atomic_thread_fence' in str(body.body[2])
+        assert str(body.body[3]) == 'sdata0[0].flag = 2;'
         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
@@ -597,8 +605,9 @@ def test_tasking_multi_output(self):
         assert str(sections[0].body[0].body[0].body[0].body[0]) ==\
             'while(lock0[t2] == 0);'
         body = op1._func_table['release_lock0'].root.body
+        assert 'atomic_thread_fence' in str(body.body[1])
         for i in range(3):
-            assert 'lock0[t' in str(body.body[1 + i])  # Set-lock
+            assert 'lock0[t' in str(body.body[2 + i])  # Set-lock
         body = op1._func_table['activate0'].root.body
         assert str(body.body[-1]) == 'sdata0[wi0].flag = 2;'
         assert len(op1._func_table) == 5

tests/test_iet.py

@@ -12,7 +12,7 @@
 from devito.ir.iet import (
     Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr,
     ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List,
-    Switch, Transformer, filter_iterations, make_callable, make_efunc,
+    Switch, ThreadFence, Transformer, filter_iterations, make_callable, make_efunc,
     retrieve_iteration_tree
 )
 from devito.ir.iet.visitors import sorted_efuncs
@@ -118,6 +118,23 @@ def test_nested_calls_cgen():
     assert str(code) == 'foo(bar());'
 
 
+def test_thread_fence_cgen():
+    """
+    A fence renders as the builtin, which is valid in C and C++ alike.
+
+    Threads that hand work to each other through a shared flag need one either
+    side of the handshake; without them the stores can be observed out of order
+    and a hand-off is lost.
+    """
+    assert str(CGen().visit(ThreadFence('acquire'))) == \
+        '__atomic_thread_fence(__ATOMIC_ACQUIRE);'
+    assert str(CGen().visit(ThreadFence('release'))) == \
+        '__atomic_thread_fence(__ATOMIC_RELEASE);'
+
+    with pytest.raises(AssertionError):
+        ThreadFence('seq_cst')
+
+
 @pytest.mark.parametrize('mode,expected', [
     ('basics', '["x"]'),
     ('symbolics', '["f"]')