[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Description

Modernizes executor synchronization and cancellation with C++20 atomic wait/notify, race-free first-failure publication, linked context cancellation, exactly-once task completion, and reusable barriers.

Snapshots termination state at Run entry and propagates tokens through steps, kernel contexts, partial execution, CPU control-flow and generation/search subgraphs. Reset affects future runs without reviving already-snapshotted cancelled work. C ABI signatures/API version remain unchanged; the C header clarifies UnsetTerminate.

Uses the existing portable in-house cancellation primitive instead of assuming uniform std::stop_token support across configured toolchains.

Refreshed on main 9f6baac7319a789e729ccee1dba9b8e3ac9c2e48, preserving main's session serialization, EPContext and node-test-loader fixes:

• Snapshot capture remains before acquiring the session serialization lock, so queued runs cannot accidentally capture a new reset epoch.
• Reconciles T5/Whisper status handling and session capture/replay behavior.
• Full compilation exposed three stale bool-based MoE test-context calls; both shared uncancelled declarations now use CancellationToken, preserving their intent rather than introducing implicit bool conversion.

Motivation and Context

Busy-waiting callers consume worker capacity, unsynchronized status publication races, and a plain terminate flag is unsuitable for the public cross-thread termination contract. Token snapshots and acquire/release completion provide explicit lifetime and reset semantics.

This is independent of #29817's planner change; no hard prerequisite is asserted. Their shared header and behavior were integrated and exercised together.

Fresh validation

Windows Release CPU/DirectML MSVC/Ninja build of the actual runtime, provider tests and Python extension passed. This is broader than the earlier syntax/extracted-method harnesses.

.\onnxruntime_test_all.exe --gtest_filter=CancellationTest.*:RunOptionsTerminationTest.*:StreamExecutionContextTest.*:ParallelExecutor.TestStatusPropagation:ParallelExecutor.ConcurrentRunCancellationFanoutAndReset:ExecutionProviderCapabilitiesTest.*:MoeExpertCountingTest.*:PlannerTest.*:ThreadPoolTest.WorkerThreadCountExcludesCallerAndHybridGranularity

51 passed, zero skips: primitive 4, RunOptions 2, stream context 7, real parallel executor 2, capability 7, migrated MoE contexts 11, planner 17, worker-count 1. ConcurrentRunCancellationFanoutAndReset executed against real initialized sessions.

.\onnxruntime_provider_test.exe --gtest_filter=Loop.*:If.*:Scan8.*:Scan9.*

69 passed, zero skips; 1 existing disabled case. Includes Loop.InfiniteLoopTermination.

The freshly built native Python extension, not the installed older wheel, also passed terminate getter/setter reset checks and five cancellation/reset/reuse cycles through an initialized parallel CPU session, verifying full output shape and every element against an independent reference.

Earlier focused evidence remains valid in its scope: 4 primitive cases repeated 100 times; inference/training+ATEN syntax checks. Those were not a substitute for the full integration results above.

Performance and remaining gates

The earlier body's completion-primitive and SqueezeNet/SSD measurements are historical author-reported results; they have not been repeated for this main reconciliation. No fresh CPU-efficiency/SLA claim is made here.

Fresh Linux training passed on this head: training job, CTest 15/15 passed. ConcurrentRunCancellationFanoutAndReset and SnapshotFanoutAndReset explicitly ran and passed. Windows native first-failure publication/status-propagation stress also completed 1,000 repetitions per case, 2,000 executions, zero failures.

Current-head checks are settled without unsuccessful checks. Targeted cached PartialRun reuse and exact device-stream dispatch remain narrower coverage questions; a green general GPU job is not silently treated as proof of every path. Architecture-owner approval is still required.

The final main refresh includes #33168 rather than retrying its deterministic dtype failures; #33183 was closed only after the complete 1,650-model gate passed.

[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.

[Gopalakrishnan Nallasamy (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.

[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Additional related baseline evidence: #33187 records an ARM64 lost-status failure in unchanged main executor code on #32194, while this PR current ARM64 gate passed (job113040928202). Local actual TestStatusPropagation + competing FirstFailureStatus tests are now running1000 repetitions. A green control is not proof of eliminating an intermittent race; preserving the failure report rather than rerunning it away.

[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Completed actual native Windows Release stress: ParallelExecutor.TestStatusPropagation and StreamExecutionContextTest.FirstFailureStatusPublishesOneCompetingFailure ran1000 repetitions each (2000 executions), with1000 passed pairs, zero failed cases and clean process exit0. The real error-return/throw paths are exercised. This strengthens the first-failure publication evidence but is not ARM64 stress or proof of eliminating every race in #33187.