Skip to content

compiler: order the async task handshake - #3009

Open
mloubout wants to merge 1 commit into
mainfrom
fix-async-memory-ordering
Open

compiler: order the async task handshake#3009
mloubout wants to merge 1 commit into
mainfrom
fix-async-memory-ordering

Conversation

@mloubout

@mloubout mloubout commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

The lock and flag an asynchronous task synchronises on are written by both
threads and declared volatile. That re-issues the loads and orders nothing, so
the handshake can lose 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. That is the state a stalled run sits in.

What's here

  • ir/iet/nodes.pyThreadFence, a Call subclass in the same vein as
    Prodder, so it renders itself and the passes emit one agreed thing rather
    than each hand-rolling a call.
  • passes/iet/asynchrony.py — a release before the flag that publishes a
    request, an acquire before the task reads what it stands for, a release
    before the flag that reports completion.
  • passes/iet/orchestration.py — in release_lock, an acquire after the wait
    and a release after the lock is handed back, before the request that refills
    it. _make_waitlock wants an acquire too, so the body's reads cannot be
    hoisted above the wait, but adding a second node to that List changes the
    IET's shape — a one-node List is denested, a two-node one is not — and
    test_streaming_fused matches on that shape. Left out with a note; the
    ordering the deadlock turns on is the store side, which the other five carry.
  • tests/test_iet.pytest_thread_fence_cgen, that the node renders for
    both orders and rejects any other. Runs anywhere.
  • tests/test_gpu_common.py — the existing lock tests index those two
    callables positionally, so they move with the fences; they also now assert the
    fences are where they should be.

__atomic_thread_fence is a compiler builtin valid in C and C++ alike, so the
device targets that render through CXXPrinter are unaffected.

Two alternatives considered and rejected:

  • declaring the two objects _Atomic is the standard-clean fix and satisfies
    ThreadSanitizer, where a fence does not — but _Atomic is not valid C++ under
    g++ ('_Atomic' does not name a type), and the qualifier mapper is shared with
    CXXPrinter, so it would break the CUDA/HIP/SYCL paths. std::atomic<int> is
    a type rather than a qualifier and cannot come through that mapper.
  • the Fence hierarchy in types/parallel.py (ThreadCommit, ThreadArrive,
    ThreadWait) reads like the right vocabulary, but those classify Clusters so
    passes cannot reorder across them, and emit nothing.

Reproducer

Drop this in debug-scripts/ and run python async_handshake_deadlock.py 6 900:
a checkpointed wavefield streamed to disk, written by a forward operator and
read back by an adjoint one, six workers for contention. It needs several
workers — one on its own rarely hits the window.

async_handshake_deadlock.py
"""The async task handshake loses an update and both threads wait forever.

A checkpointed wavefield streamed to disk is written by a forward operator and
read back by an adjoint one.  Each pair hands off through the generated
`lock0`/`flag` protocol roughly once per time step.  Declared `volatile int`
those two objects are neither atomic nor ordered, so a delivery can be lost --
the lock handed back after the value that replaced it -- and then the compute
thread waits for data it has already spent the request for while the task waits
to be asked.  Nothing moves, both threads spin at 100% CPU.

    python debug-scripts/async_handshake_deadlock.py            # 6 workers, 15 min
    python debug-scripts/async_handshake_deadlock.py 4 300      # 4 workers, 5 min

Needs contention: one worker on its own rarely hits the window.  A worker that
goes quiet for over 90 s while still burning CPU is the deadlock; the script
says so and prints what the generated operator declared.
"""

import os
import sys
import time
from multiprocessing import Process, Value

from devitopro import TimeFunction
from devitopro.types.enriched import Disk

from devito import ConditionalDimension, Eq, Function, Grid, Inc, Operator

SHAPE = (240, 200)
SO = 8
NT = 1200          # ~1200 handshakes per operator call, as in a real gradient
FACTOR = 5


def build():
    grid = Grid(shape=SHAPE, extent=(12000., 10000.))
    t_sub = ConditionalDimension(name='t_sub', parent=grid.time_dim,
                                 factor=FACTOR)
    b = Function(name='b', grid=grid, space_order=SO)
    b.data[:] = 0.5
    u = TimeFunction(name='u', grid=grid, space_order=SO, time_order=2)
    v = TimeFunction(name='v', grid=grid, space_order=SO, time_order=2)
    g = Function(name='g', grid=grid)
    usave = TimeFunction(name='usave', grid=grid, space_order=SO, time_order=2,
                         time_dim=t_sub, save=NT // FACTOR + 2, layers=Disk,
                         compression='cvxcompress')

    fwd = Operator([Eq(u.forward, (b * u).laplace + 2 * u - u.backward),
                    Eq(usave, u.forward)])
    adj = Operator([Eq(v.backward, (b * v).laplace + 2 * v - v.forward),
                    Inc(g, usave * v)])
    return fwd, adj, u, v, usave


def declaration(op):
    """How the generated operator guards the handshake."""
    code = str(op.ccode)
    decls = [line.strip() for line in code.splitlines()
             if ('lock0[1]' in line or 'int flag' in line)]
    return decls + [f'{code.count("__atomic_thread_fence")} fence(s)']


def worker(beat, cycles, index):
    fwd, adj, u, v, usave = build()
    if index == 0:                          # one worker reports the code
        print(f'  generated: {"; ".join(declaration(adj))}', flush=True)
    for _ in range(cycles):
        u.data[:] = 0.
        fwd.apply(time_M=NT, dt=1.)
        v.data[:] = 0.
        adj.apply(time_M=NT - FACTOR - 2, dt=1.)
        for f in usave.values() if hasattr(usave, 'values') else [usave]:
            f._reset()
        beat.value = int(time.time())       # heartbeat


def main():
    nworkers = int(sys.argv[1]) if len(sys.argv) > 1 else 6
    budget = int(sys.argv[2]) if len(sys.argv) > 2 else 900
    os.environ.setdefault('OMP_NUM_THREADS', '4')

    beats = [Value('l', 0) for _ in range(nworkers)]
    procs = [Process(target=worker, args=(b, 10 ** 6, i))
             for i, b in enumerate(beats)]
    for p in procs:
        p.start()
    print(f'{nworkers} workers, {budget}s budget', flush=True)

    start = time.time()
    stalled = []
    while time.time() - start < budget and len(stalled) == 0:
        time.sleep(15)
        now = time.time()
        for i, b in enumerate(beats):
            if b.value and now - b.value > 90 and procs[i].is_alive():
                stalled.append(i)
    for p in procs:
        p.terminate()

    if stalled:
        print(f'DEADLOCK: worker(s) {stalled} stopped making progress while '
              f'still running', flush=True)
        return 1
    print(f'no deadlock in {int(time.time() - start)}s', flush=True)
    return 0


if __name__ == '__main__':
    sys.exit(main())

Same script, same budget, only the compiler differing:

generated outcome
before volatile int flag; volatile int lock0[1], 0 fences DEADLOCK: worker(s) [2] stopped making progress while still running
after same declarations, 6 fences no deadlock in 900s

The failing run also trips CvxCompress's own assertion —
Decompress: nx=216, ny=256, nz=1, nx_check=0 — the compute decompressing a
buffer the task never filled. Same lost update wearing its other face: this does
not only hang, it can hand garbage to the decompressor.

Notes

Found through a streaming-checkpoint FWI gradient on arm64, where it stalled one
worker in six within minutes. x86's store ordering hides it.

The test_gpu_common.py assertions need a device; devitopro carries an
equivalent one for the streaming path, which runs anywhere
(devitocodespro/devitopro#954).

@mloubout
mloubout force-pushed the fix-async-memory-ordering branch from b54188d to 9bf31a1 Compare August 20, 2026 16:15
@codecov

codecov Bot commented Aug 20, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.68%. Comparing base (436199c) to head (12166a8).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3009      +/-   ##
==========================================
+ Coverage   83.67%   83.68%   +0.01%     
==========================================
  Files         257      257              
  Lines       54663    54687      +24     
  Branches     4683     4683              
==========================================
+ Hits        45739    45765      +26     
+ Misses       8114     8113       -1     
+ Partials      810      809       -1     
Flag Coverage Δ
pytest-gpu-aomp-amdgpuX 68.47% <100.00%> (-0.01%) ⬇️
pytest-gpu-gcc- 78.29% <40.62%> (-0.02%) ⬇️
pytest-gpu-icx- 78.21% <40.62%> (-0.02%) ⬇️
pytest-gpu-nvc-nvidiaX 69.13% <100.00%> (+0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@mloubout
mloubout force-pushed the fix-async-memory-ordering branch from 9bf31a1 to 991305b Compare August 20, 2026 16:35
@mloubout mloubout changed the title compiler: make the async handshake's lock and flag atomic compiler: order the async task handshake Aug 20, 2026
@mloubout
mloubout force-pushed the fix-async-memory-ordering branch 2 times, most recently from efc7bb8 to 7e6ac61 Compare August 20, 2026 18:23
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.
@mloubout
mloubout force-pushed the fix-async-memory-ordering branch from 7e6ac61 to 12166a8 Compare August 20, 2026 18:34
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

probably blank lines

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants