[Core] Modernize execution synchronization and cancellation

[GopalakrishnanN]

Description

Modernize graph execution synchronization and cancellation using C++20 primitives:

• Replace StreamExecutionContext::WaitAll() busy-spinning with std::atomic::wait/notify_all and acquire-release task-count publication.
• Make first-failure status publication race-free and stop sibling workers through a context-local cancellation source.
• Snapshot a cancellation token at Session::Run() entry and propagate it by value through execution steps, kernel contexts, partial execution, CPU control-flow subgraphs, and contrib generation/search subgraphs.
• Make RunOptionsSetTerminate safe during active runs. UnsetTerminate installs a fresh stop state for future runs while already-snapshotted tokens remain stopped.
• Make every RunSince() exit and exception path complete its task exactly once with a scope guard.
• Keep the completion state alive across the final notify_all(), and re-arm cross-stream barriers when a partial-execution context is reused.
• Cancellation uses a small portable in-house primitive (onnxruntime::CancellationToken/CancellationSource/CancellationCallback in core/framework/cancellation.h) rather than std::stop_token, because the <stop_token> family is not uniformly available across ORT's configured toolchains (see Portability). It provides std::stop_callback-equivalent registration/deregistration semantics.
• OrtRunOptions gains user-defined copy/move members so that a copied or moved-from instance keeps an independent, usable termination state. These allocate a small shared state and are therefore not noexcept.
• Preserve C ABI signatures and the existing cancellation error text (Exiting due to terminate flag being set to true.). No ORT_API_VERSION change is required; the only public C header change is a documentation clarification on UnsetTerminate.
• Add direct conformance tests for the cancellation primitive, plus countdown, competing-failure, cancellation fan-out/reset, context-reuse, control-flow, and completion-wake benchmark coverage.

Motivation and Context

The parallel executor previously kept the caller runnable until all worker tasks completed:

while (remain_tasks_.Get()) {
  SpinPause();
}

Under concurrent ORT_PARALLEL runs, these waiting callers compete with executor workers and can consume most available CPU capacity. The previous implementation also read and wrote a non-thread-safe Status from multiple workers and used a plain bool termination flag even though the API permits termination from another thread.

The task counter now forms an acquire-release chain: every worker publishes its writes through an acq_rel decrement, the final transition to zero notifies waiters, and the caller observes completion with an acquire load. Downstream work increments the count before scheduling and before the parent can decrement, and incrementing a zero count is rejected.

RunOptions now uses a mutex-protected stop source. Active runs hold token snapshots, so one terminate request cancels all currently active runs using the options object, while reset affects only future runs. Internal worker failures use a local linked cancellation source so they stop sibling workers without mutating user-owned options.

Performance

Completion primitive

On a 24-logical-CPU Xeon w5-2455X with 32 waiters, using an alternating preserved-binary A/B benchmark:

• Median process CPU/completion: 14.10 ms -> 4.95 ms (-64.9%)
• Effective cores consumed: 17.96 -> 4.68 (-73.9%)
• Median wake latency: 1.630 ms -> 1.633 ms (+0.21%)
• P95 wake latency: 2.052 ms -> 2.078 ms (+1.28%)

The sustained 32-waiter completion benchmark had a wall-time cost (795.8 us -> 1038.8 us, +30.5%) because parked threads must be rescheduled. The intended benefit is CPU efficiency and reduced worker starvation, not universally lower primitive wake latency.

End-to-end Session::Run

Clean parent/modified Release binaries were compared with real CPU inference, one shared session, ORT_PARALLEL, intra-op=1, inter-op=12, deterministic inputs, 10 alternating paired rounds, and paired-bootstrap 95% confidence intervals.

SqueezeNet:

Concurrent runs	P50	P95	P99	Throughput	CPU/request
24	-26.7%	-34.9%	-35.4%	+47.9% [+36.1%, +60.8%]	-64.3%
32	-46.1%	-29.5%	Inconclusive	+60.1% [+42.8%, +78.2%]	-68.0%
48	-50.9%	-42.0%	-52.7%	+100.5% [+83.3%, +116.4%]	-75.1%

For concurrency 24, 32, and 48, throughput and CPU/request improved in all 10 paired runs (two-sided exact sign test p=0.002). At concurrency 48, effective process occupancy fell from 18.87 to 9.31 cores.

For latency-sensitive MNIST at concurrency 48:

• P50: statistically neutral (1.009 ms -> 1.030 ms)
• P95: 58.73 ms -> 3.46 ms (-94.2%)
• P99: 86.48 ms -> 4.75 ms (-94.6%)
• Throughput: 5,469 -> 35,173 QPS (+538.3%)
• CPU/request: 4.118 ms -> 0.261 ms (-93.6%)

A 32-way ORT_SEQUENTIAL negative control, which does not wait in the parallel executor, was neutral for average latency, throughput, P95/P99, and effective CPU occupancy. Low-concurrency parallel latency and throughput were also statistically inconclusive. This is a workload- and concurrency-specific improvement, not a universal latency claim.

Measurement provenance: the end-to-end A/B compared 361184e619 (base) against 92e204b154 (this change's initial implementation), on a single Windows/x86-64 host. Subsequent commits on this branch were the portability swap to the in-house cancellation primitive, lifetime/reuse hardening, and added tests; none of them alter the park-instead-of-spin completion mechanism that the measurement isolates.

Validation

Validated on Windows/MSVC Release at the current head, which includes a merge of main:

• onnxruntime_test_all: 1,899 tests from 147 suites — 1,872 passed, 27 skipped, 0 failures.
• Focused synchronization/cancellation suites (CancellationTest.*, StreamExecutionContextTest.*, RunOptionsTerminationTest.*, ParallelExecutor.*): 15 tests from 4 suites, all passed.
• Race-sensitive subset repeated 100 times: 14 tests per iteration, 0 failures across all 100 iterations.
• CPU control-flow (Loop.*:Scan*.*:If.*:ControlFlow*.*): 64 tests from 4 suites, all passed.
• Loop.InfiniteLoopTermination: passed using the thread-safe RunOptions helper.
• Training partial-execution framework/session targets compiled.
• Python onnxruntime_pybind11_state compiled.
• onnxruntime_benchmark built and executed.
• Lintrunner/clang-format and git diff --check: clean.

TSAN was not run because this Windows/MSVC environment does not have a readily configured TSAN target. This is the main acknowledged gap: the change is fundamentally about memory ordering, and repeated runs exercise timing rather than the happens-before graph. Reviewers who can run the focused suites under TSAN on Linux are very welcome to.

Portability

ORT already requires C++20, and std::atomic::wait/notify_all is available across ORT's configured toolchains (MSVC, Android NDK 28 libc++, Emscripten 4.0.23), so the completion barrier uses it directly.

The <stop_token> family (std::stop_token/std::stop_source/std::stop_callback) is not uniformly available: Android NDK 28 libc++ does not provide it, and older Apple libc++ only exposes it behind -fexperimental-library, which in turn fails to link libc++experimental on iOS. Cancellation therefore uses a small portable in-house primitive built only on <atomic>, <mutex>, <condition_variable>, and <memory>, with std::stop_callback-equivalent semantics (the callback destructor blocks until a concurrent invocation completes, with a same-thread guard to avoid self-deadlock). No experimental-library flag is required, and the previously proposed AppleClang compile/link flag has been removed. The primitive has dedicated conformance tests in onnxruntime/test/framework/cancellation_test.cc covering post-cancellation registration, exactly-once invocation under competing stop requests, destructor/invocation races, and callback-driven release of the owning state.

The added synchronization/cancellation runtime slightly increases the minimal MinSizeRel arm64 libonnxruntime.so, so the Android AndroidBinarySizeCheckJob_MinimalBaseline size budget was raised. The current value (1,484,800 bytes) also absorbs an independent increase made on main.

On threaded Wasm workers, atomic waiting can park. Browser main threads cannot block in Atomics.wait, so Emscripten may retain spinning behavior there while preserving correctness.

[Copilot]

Pull request overview

This PR modernizes ONNX Runtime’s graph execution synchronization and cancellation path by replacing busy-waiting with C++20 blocking primitives, making cancellation propagation and failure publication thread-safe, and updating call sites (execution steps, control-flow subgraphs, training partial execution, and Python bindings) to pass a std::stop_token snapshot by value.

Changes:

• Replace StreamExecutionContext::WaitAll() spinning with std::atomic::wait/notify_all and introduce a thread-safe “first failure wins” status publication plus context-local stop fan-out.
• Snapshot and propagate termination via std::stop_token through InferenceSession::Run/PartialRun, executors, kernel contexts, and CPU control-flow / contrib subgraphs.
• Update tests and microbenchmarks to validate completion wake behavior, competing failures, cancellation fan-out/reset, and termination behavior.

Reviewed changes

Copilot reviewed 29 out of 29 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/providers/cpu/controlflow/loop_test.cc	Updates loop termination test to use the new RunOptions termination API.
onnxruntime/test/onnx/microbenchmark/tptest.cc	Adds a microbenchmark for StreamExecutionContext completion wake behavior.
onnxruntime/test/framework/stream_execution_context_test.cc	Adds unit tests for CountDownBarrier, FirstFailureStatus, and RunOptions termination snapshot/reset.
onnxruntime/test/framework/parallel_executor_test.cc	Adds concurrent Run cancellation fan-out/reset coverage using a custom op that waits on cancellation.
onnxruntime/python/onnxruntime_pybind_state.cc	Rebinds RunOptions.terminate as a property backed by stop-token state, preserving the Python surface API.
onnxruntime/core/session/inference_session.h	Extends internal RunImpl to accept a std::stop_token snapshot.
onnxruntime/core/session/inference_session.cc	Snapshots termination token at Run/PartialRun entry and propagates it down to graph execution utilities.
onnxruntime/core/providers/cpu/controlflow/scan_utils.cc	Propagates stop-token cancellation into Scan subgraph execution.
onnxruntime/core/providers/cpu/controlflow/loop.cc	Propagates stop-token cancellation into Loop subgraph execution.
onnxruntime/core/providers/cpu/controlflow/if.cc	Propagates stop-token cancellation into If subgraph execution.
onnxruntime/core/framework/utils.h	Updates execution utility signatures to accept std::stop_token and removes the RunOptions-based ExecuteGraph overload.
onnxruntime/core/framework/utils.cc	Implements the updated stop-token-based ExecuteGraph/ExecuteSubgraph/ExecutePartialGraph plumbing.
onnxruntime/core/framework/stream_execution_context.h	Introduces atomic wait/notify completion barrier, first-failure status wrapper, and stop-token integration into the stream execution context.
onnxruntime/core/framework/stream_execution_context.cc	Implements first-failure publication, stop fan-out, scope-guarded task completion, and the new blocking WaitAll.
onnxruntime/core/framework/sequential_executor.h	Updates executor APIs to take std::stop_token rather than a bool& terminate flag.
onnxruntime/core/framework/sequential_executor.cc	Threads cancellation token through scheduling, kernel execution contexts, and partial execution.
onnxruntime/core/framework/sequential_execution_plan.h	Updates execution-step interface to take a stop-token.
onnxruntime/core/framework/run_options.cc	Implements mutex-protected stop-source state for RunOptions terminate/set/unset with snapshot semantics.
onnxruntime/core/framework/partial_graph_execution_state.h	Adds stop-token propagation for training partial execution contexts.
onnxruntime/core/framework/partial_graph_execution_state.cc	Resets/reuses StreamExecutionContext with per-run stop-token and re-computes valid streams consistently.
onnxruntime/core/framework/op_kernel_context_internal.h	Stores and exposes a per-run cancellation stop-token to kernels via OpKernelContextInternal.
onnxruntime/core/framework/execution_steps.h	Updates step Execute signatures to use stop-token.
onnxruntime/core/framework/execution_steps.cc	Threads stop-token through kernel launch and downstream scheduling steps.
onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h	Propagates stop-token cancellation into contrib subgraph execution.
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h	Propagates stop-token cancellation into contrib subgraph execution.
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h	Propagates stop-token cancellation into contrib subgraph execution.
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h	Propagates stop-token cancellation into contrib subgraph execution.
include/onnxruntime/core/framework/run_options.h	Updates public RunOptions to provide stop-token based termination APIs and internal termination state storage.
cmake/CMakeLists.txt	Adds an AppleClang(<17) compile/link flag for libc++ stop-token support on legacy Xcode.

[Copilot]

Pull request overview

Copilot reviewed 30 out of 30 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.

[GopalakrishnanN]

Review synthesis — synchronization & cancellation modernization

I ran a multi-angle review over the whole diff (readability, correctness, adversarial/race, memory-model/spec, cross-module integration). Overall this is a solid, well-tested refactor and the hard core holds up: the cancellation.h register/deregister races, the stop_callback destructor two-halves, the completion-count fork-join invariant, the cancellation funneling through the internal stop_source_, and the CountDownBarrier happens-before chain (acq_rel-per-RMW + acquire-on-Wait) all check out under adversarial interleaving. No correctness issue in the primitive itself.

Two issues in the new CountDownBarrier are worth addressing before merge — both are new to the spin→wait/notify change, and neither is caught by the current tests.

Major

1. notify_all() can run on a destroyed barrier (use-after-free) — stream_execution_context.h (CountDownBarrier::Dec/Wait) + sequential_executor.cc (ExecuteThePlan → WaitAll())

Wait() returns immediately when its initial load reads 0, without ever calling atomic::wait. So this interleaving is reachable:

1. Worker does the final Dec(): compare_exchange sets the count 1→0, then the worker is descheduled before reaching v_.notify_all().
2. Main thread calls WaitAll() → Wait() loads 0 and returns without blocking.
3. Main thread finishes ExecuteThePlan and destroys the stack-local ctx (which owns remain_tasks_).
4. Worker resumes and calls v_.notify_all() on the now-destroyed atomic → UB.

(The path where Wait() actually blocks is safe — there the notify necessarily precedes the wake. The bug is specifically the "load reads 0, never blocks" path, where the notify becomes a dangling access.) The old while (remain_tasks_.Get()) SpinPause(); had no post-decrement notify, so the decrementing worker had no access to the barrier after it hit zero — hence this is new. Suggested fix: give each scheduled task shared ownership of the completion state (e.g. a task-group whose lease holds a shared_ptr<CompletionState>), so the barrier outlives the last notify; raw &ctx captures don't extend its lifetime.

2. count_down_barriers_ isn't reset on context reuse → ORT_ENFORCE throw on 2nd iteration — stream_execution_context.cc (ctor barrier loop vs ResetForExecution), partial_graph_execution_state.cc (reuse path)

The 2-party stream barriers are Set(2) only in the constructor. ResetForExecution re-Sets remain_tasks_ but not count_down_barriers_. On a reused context (PartialGraphExecutionState caches the ctx and the else-branch calls only ResetForExecution), a barrier that reached 0 in iteration N gets decremented again in N+1. The new Dec() added ORT_ENFORCE(value > 0), which turns the old silent fetch_sub underflow into a hard throw on the second iteration.

Reachability depends on whether reused partial-graph contexts run a plan with num_barriers > 0 (multi-stream) — you're best placed to confirm. If so, ORTModule/partial-graph training that reuses the state across steps would fail on step 2. Suggested fix: re-initialize the barriers in ResetForExecution (guarded by ORT_ENABLE_STREAM), mirroring the constructor.

3. No dedicated conformance test for the new primitive — the cancellation.h contracts (immediate invoke when registering after stop; exactly-once under concurrent request_stop; destructor blocking while the callback runs on another thread; no block for same-thread self-deregistration) are only exercised indirectly. A focused cancellation_test.cc would lock these in against future regressions.

4. terminate_token means two different things along the call stack — sequential_executor.cc

The parameter named terminate_token at the top of ExecuteThePlan/PartialExecuteThePlan is the raw caller token, but the identically-named parameter threaded into RunSince/ExecuteKernel/ExecutionStep::Execute actually carries ctx.GetCancellationToken() — the internal source that merges external-terminate and first-failure (so one stream's failure stops its siblings). The single crossover local is called cancellation_token only because the name was taken. A future contributor "fixing" the apparent inconsistency by feeding the outer token straight through would silently break cross-stream failure propagation. Recommend reserving terminate_token for the raw token, renaming the threaded param to cancellation_token, and adding a one-line comment on GetCancellationToken().

Minor

• CancellationSource::request_stop() doesn't pin state_ — if an invoked callback drops the last reference to the state, the post-invoke work touches freed memory. Not reachable in the executor (its callbacks keep another source alive), but cancellation.h is a public header. One-liner: auto state = state_; return state && state->RequestStop();.
• stop_possible() diverges from std::stop_token — returns true even after all sources are gone without a stop request (std returns false). Currently unused, so harmless; worth a comment or a live-source count.
• Moved-from OrtRunOptions leaves termination_state_ null, so a later RequestTerminate()/GetTerminateToken() null-derefs. No live move-then-use path today; consider re-seeding in a user-defined move or = deleteing moves.
• UnsetTerminate no longer un-cancels an in-flight run — the old bool was a live toggle; the new snapshot + fresh-source-on-Reset model means a run holding an earlier snapshot stays cancelled. This matches the stop_token monotonic-latch model and looks intended — worth a doc note on RunOptionsUnsetTerminate. (Terminate requested after Run() starts still propagates via shared state.)
• Pre-existing exception-safety gap: ScheduleDownstream does AddTask() before ThreadPool::Schedule, whose std::function copy can throw — leaking the reservation (hang) or unwinding ctx while a scheduled task holds &ctx. Not introduced here, but the modernization is a natural place to harden it.
• Include-graph ripple: cancellation.h pulls <mutex>/<condition_variable>/<thread>/<atomic> into widely-included headers (run_options.h, utils.h), which likely explains the android.yml MinSizeRel arm64 size baseline bump. A pimpl/out-of-line boundary would limit the spread.
• Test/doc polish: add a copy-while-stopped RunOptions case (existing copy test only covers the reset/non-stopped path); refresh the stale CountDownBarrier class comment (it predates Inc()/dynamic task addition); document that Invoke() is noexcept (throwing callback ⇒ std::terminate, like std::stop_callback); note why external_stop_callback_ is a std::optional (re-arming a non-movable registration) and why TerminationState::RequestStop releases the lock before requesting stop.

Note on the OrtRunOptions::terminate removal (checked — C API is safe)

I verified the removed bool terminate member does not break the EP boundary: the plugin-EP bridge forwards &run_options as an opaque const OrtRunOptions*, and out-of-tree EPs interact through the preserved C API (RunOptionsSetTerminate/getter) and the Python property shim, not the struct layout. In-tree EPs are co-compiled. The only real residual is a C++ source break for out-of-tree code that included the internal header and read .terminate directly — worth a line in the release notes, but not an ABI hazard for C API consumers.

Verification notes

Cross-module coverage looks complete — every in-tree GetTerminateFlag() consumer (control-flow + contrib transformers) and every run_options.terminate consumer is updated. Since this branch predates the current main plugin_ep/ tree, a rebase should re-confirm no newer .terminate/GetTerminateFlag consumers slipped in.

I couldn't build/run locally, so this is static review only — both Majors are exactly what a sanitizer run would surface. If CI can exercise the parallel executor under ASan/TSan (and a ≥2-iteration partial-graph/ORTModule training run on a multi-stream plan for #2), that would be the fastest confirmation.