[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 std::stop_source.
• Snapshot std::stop_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.
• Preserve C ABI signatures and the existing cancellation error text. No ORT_API_VERSION change is required.
• Add countdown, competing-failure, cancellation fan-out/reset, 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 stop 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.

Validation

• Release onnxruntime_test_all: 1,852 tests from 142 suites, 0 failures (1,825 passed, 27 skipped).
• Focused synchronization/cancellation tests: 9 tests from 3 suites, all passed; repeated 50 times without failure.
• CPU control-flow tests: 66 tests from 5 suites, all passed.
• Loop.InfiniteLoopTermination: passed with the thread-safe RunOptions helper.
• Debug focused synchronization/cancellation tests: 9 tests, all passed.
• 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.

Portability

ORT already requires C++20. MSVC supports both atomic wait and stop tokens. Android CI uses NDK 28 and Wasm CI uses Emscripten 4.0.23. Legacy Xcode 15/16 libc++ gates stop-token support behind -fexperimental-library, so a narrowly scoped AppleClang <17 compile/link flag is included; Xcode 26+ is unaffected.

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.