From 12166a8ceb0ac744776676105b3087a7f7013b3e Mon Sep 17 00:00:00 2001 From: mloubout Date: Thu, 20 Aug 2026 12:18:06 -0400 Subject: [PATCH] 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. --- devito/ir/iet/nodes.py | 29 ++++++++++++++++++++++++++++- devito/passes/iet/asynchrony.py | 16 ++++++++++++++-- devito/passes/iet/orchestration.py | 6 +++++- tests/test_gpu_common.py | 21 +++++++++++++++------ tests/test_iet.py | 19 ++++++++++++++++++- 5 files changed, 80 insertions(+), 11 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 4e63438dad7..aeaa18477b2 100644 --- a/devito/ir/iet/nodes.py +++ b/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,32 @@ def size(self): return self.ispace.size +class ThreadFence(Call): + + """ + A memory fence, ordering a thread's accesses either side of it. + + Threads that hand work to each other through a shared flag need one on both + sides of the handshake: a release before the flag that publishes a request + or a completion, so everything it stands for is visible first, and an + acquire before reading what it stands for. Without them the two stores can + be observed out of order and an update is 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): """ diff --git a/devito/passes/iet/asynchrony.py b/devito/passes/iet/asynchrony.py index aa205e818cf..5bfe56e87f2 100644 --- a/devito/passes/iet/asynchrony.py +++ b/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,13 @@ 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]) + # The flag is what publishes the request, so everything the thread will + # read on the back of it -- the arguments here, and anything the caller + # wrote before, such as a lock released back to the thread -- has to be + # visible first. `volatile` does not order stores, and on a weakly + # ordered target this one can be seen before them: the thread then acts + # on stale state, and a lock handed over that way is lost. + activation.append(ThreadFence('release')) activation.append( DummyExpr(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 2) ) @@ -138,12 +145,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) ])] diff --git a/devito/passes/iet/orchestration.py b/devito/passes/iet/orchestration.py index ae2d55766cd..f9969801617 100644 --- a/devito/passes/iet/orchestration.py +++ b/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') diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index 9b6c5fdfe11..c0e4d00adec 100644 --- a/tests/test_gpu_common.py +++ b/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 diff --git a/tests/test_iet.py b/tests/test_iet.py index 6d2b3126c67..a66d5aa62bf 100644 --- a/tests/test_iet.py +++ b/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"]')