[Plugin EP] Port graph capture/replay APIs by adrianlizarraga · Pull Request #27958 · microsoft/onnxruntime

Adrian Lizarraga (adrianlizarraga) · 2026-04-03T00:08:35Z

NOTE: This PR cannot be merged until the ORT version is updated to 1.26.0 in the main branch

Description

Ports graph capture/replay APIs (e.g., CUDA Graph) to the Plugin EP (OrtEp) C API so that plugin-based execution providers can participate in ORT-managed graph capture and replay.

What changed

New Plugin EP C API functions (onnxruntime_ep_c_api.h):

OrtEp::IsGraphCaptureEnabled — indicates whether the EP has graph capture enabled.
OrtEp::IsGraphCaptured — indicates whether a graph has been captured for a given annotation ID.
OrtEp::ReplayGraph — replays a previously captured graph.
OrtEp::GetGraphCaptureNodeAssignmentPolicy — returns the node assignment validation policy for graph capture.

All four are optional (NULL defaults to safe behavior) and version-gated (ort_version_supported >= 26).
If IsGraphCaptureEnabled returns true, IsGraphCaptured and ReplayGraph must also be implemented;
otherwise PluginExecutionProvider logs a warning and disables graph capture for that EP.

New OrtGraphCaptureNodeAssignmentPolicy enum (onnxruntime_ep_c_api.h):
Replaces the hardcoded EP-name checks in InferenceSession::Initialize() with a policy-based approach:

ALL_NODES_ON_EP — all nodes must be on the target EP (e.g., TensorRT).
ALLOW_CPU_FOR_SHAPES — CPU nodes allowed for shape computation if no memcpy nodes exist (e.g., CUDA, WebGPU, DML).

Refactored InferenceSession graph capture selection (inference_session.cc):

Removed the hardcoded graph_support_ep_list and per-EP strcmp checks.
Now iterates over all registered EPs and uses IsGraphCaptureEnabled() + GetGraphCaptureNodeAssignmentPolicy() to select and validate the graph-capturing EP.
AreAllComputeNodesAssignedToCudaOrJsOrDmlEpWebGpuEp() → generalized to AreAllComputeNodesAssignedToEpOrCpu(), which also requires at least one node on the target EP.
IExecutionProvider::GetGraphCaptureNodeAssignmentPolicy() added to the base class (defaults to ALL_NODES_ON_EP).

Bounded graph capture recursion (inference_session.cc/h):

Run() now delegates to RunImpl() with a graph_capture_depth parameter.
Caps internal run attempts at kMaxGraphCaptureRunAttempts = 8, returning a clear error if the EP never reports IsGraphCaptured() == true.

EP implementations:

WebGPU plugin EP: Fully implements all four graph capture APIs by forwarding to the underlying IExecutionProvider.
CUDA plugin EP: Stubs with TODOs (returns disabled/not-implemented).
NvTensorRTRTX EP: IsGraphCaptureEnabled() now returns false since this EP manages graph capture internally (not via ORT).

C++ wrapper (onnxruntime_cxx_api.h / onnxruntime_cxx_inline.h):

Added Ort::Env::CopyTensor() convenience overload for copying a single tensor (wraps CopyTensors with num_tensors=1).

Tests

ep_plugin_provider_test.cc: Unit tests for each new PluginExecutionProvider graph capture method, including NULL function pointer defaults, version < 26 backward compatibility, and validation that IsGraphCaptureEnabled() returns false when IsGraphCaptured or ReplayGraph are NULL.
test_graph_capture.cc: End-to-end test for WebGPU plugin EP graph capture/replay using IO binding (warm-up + capture run, then replay with different inputs).

Motivation and Context

Previously, graph capture support was limited to a hardcoded list of EPs (kCudaExecutionProvider, kTensorrtExecutionProvider, kJsExecutionProvider, kWebGpuExecutionProvider, kDmlExecutionProvider) with EP-specific validation logic in InferenceSession. This made it impossible for plugin EPs to participate in ORT-managed graph capture/replay without modifying the core session code.

This PR makes graph capture/replay extensible to any EP, including out-of-tree plugin EPs, by exposing it through the OrtEp C API.

…ve the hardcoded list of EPs that support graph capture from inference_session.cc

…ation

Copilot

Pull request overview

This PR (draft) ports graph capture/replay support to the Plugin EP pathway by extending the OrtEp C API, wiring those callbacks through the plugin EP provider wrapper, and updating session initialization logic to validate/capture based on an EP-provided node-assignment policy.

Changes:

Bump ORT version/API version to 1.26.0 / ORT_API_VERSION=26 and add new OrtEp graph capture/replay callbacks plus OrtGraphCaptureNodeAssignmentPolicy.
Update InferenceSession::Initialize() to select any EP with graph capture enabled, validate graph assignment via EP policy, and cache a single EP for replay.
Add/extend tests for plugin EP graph capture APIs and add an end-to-end autoep WebGPU plugin EP graph capture/replay test.

Reviewed changes

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

Show a summary per file

File	Description
VERSION_NUMBER	Bumps runtime version to 1.26.0.
include/onnxruntime/core/session/onnxruntime_c_api.h	Bumps `ORT_API_VERSION` to 26.
onnxruntime/core/session/onnxruntime_c_api.cc	Updates version string static assert to 1.26.0.
include/onnxruntime/core/session/onnxruntime_ep_c_api.h	Adds graph capture/replay callbacks to `OrtEp` and introduces `OrtGraphCaptureNodeAssignmentPolicy`.
include/onnxruntime/core/framework/execution_provider.h	Adds `IExecutionProvider::GetGraphCaptureNodeAssignmentPolicy()` with a strict default.
onnxruntime/core/session/inference_session.cc	Generalizes graph-capture EP selection/validation and uses EP-specified node assignment policy.
onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h	Exposes graph capture/replay APIs on `PluginExecutionProvider`.
onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc	Implements plugin-side forwarding for graph capture/replay and policy query with version gating.
onnxruntime/core/providers/webgpu/ep/ep.h	Declares plugin adapter entrypoints for graph capture/replay and assignment policy.
onnxruntime/core/providers/webgpu/ep/ep.cc	Wires WebGPU plugin EP adapter function pointers and forwards to EP impl.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h	Returns `ALLOW_CPU_FOR_SHAPES` policy for WebGPU EP.
onnxruntime/core/providers/js/js_execution_provider.h	Returns `ALLOW_CPU_FOR_SHAPES` policy for JS EP.
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h	Returns `ALLOW_CPU_FOR_SHAPES` policy for DML EP wrapper.
onnxruntime/core/providers/cuda/cuda_execution_provider.h	Returns `ALLOW_CPU_FOR_SHAPES` policy for CUDA EP.
onnxruntime/core/providers/cuda/plugin/cuda_ep.h	Declares plugin CUDA EP adapter entrypoints for graph capture/replay and policy.
onnxruntime/core/providers/cuda/plugin/cuda_ep.cc	Wires plugin CUDA EP adapter function pointers (currently stubbed).
onnxruntime/test/framework/ep_plugin_provider_test.cc	Adds unit tests for plugin EP graph capture/replay function-pointer behavior and version gating.
onnxruntime/test/autoep/test_graph_capture.cc	Adds end-to-end test exercising WebGPU plugin EP graph capture + replay via public APIs.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…ionProvider::IsGraphCaptureEnabled() as ORT previously never managed graph capture for this EP; Update use of webgpu name with constant

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…ted when necessary

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Tianlei Wu (tianleiwu)

Review Summary

Well-structured PR that generalizes the hardcoded EP-specific graph capture logic in InferenceSession into a policy-based, extensible design driven by IExecutionProvider virtual methods. The C API additions are thoroughly documented, version-gated, and defensively validated. The recursion-depth guard in Run() is a valuable safety net.

Highlights

AreAllComputeNodesAssignedToEpOrCpu() now requires has_node_on_provider, which is a correctness improvement — avoids false positives where no nodes are on the target EP.
PluginExecutionProvider::IsGraphCaptureEnabled() validation that checks IsGraphCaptured/ReplayGraph are non-null is a valuable defensive measure — without it, ORT would silently hang in the recursive warm-up loop.
NvTensorRTRTX change to return false from IsGraphCaptureEnabled() is correct: it already returned false from the public IsGraphCaptured() and manages capture internally.
Excellent API documentation on the four new OrtEp members — exactly the level of detail an out-of-tree EP author needs.

Minor Observations (not blocking)

CopyTensor API consistency: The new single-tensor CopyTensor(const OrtValue*, OrtValue*, OrtSyncStream*) takes raw pointers while the existing CopyTensors takes const std::vector<Value>& wrappers. Works correctly but inconsistent style.
CUDA plugin stubs: Since IsGraphCaptureEnabledImpl returns false, the other three stub registrations are dead code. Fine as scaffolding for upcoming work.

Overall: clean, well-tested, ready to merge. Two inline suggestions below.

This reverts commit 7afe4c2.

… an empty graph (#29457) ### Description  InferenceSession::Initialize() validates, for any EP that has graph capture enabled and uses the ALLOW_CPU_FOR_SHAPES node-assignment policy (DirectML, WebGPU), that the partitioned graph is eligible for capture via AreAllComputeNodesAssignedToEpOrCpu(). That helper currently requires at least one node to be assigned to the EP: return has_node_on_provider && !HasMemcpyNodes(graph); For a graph with no nodes, like a model that is fully consumed by the EP's runtime graph fusion, or that trivially folds away (like a lone Constant node), has_node_on_provider is false, so the session throws: This session cannot use the graph capture feature as requested by the user as all compute graph nodes have not been partitioned to the DmlExecutionProvider An empty graph contains nothing that violates the "all compute on the EP or CPU, no Memcpy" requirement, so it should be allowed. This change permits the empty-graph case: return (has_node_on_provider || graph.NumberOfNodes() == 0) && !HasMemcpyNodes(graph); Behavior is unchanged for all non-empty graphs (the expression collapses to the original has_node_on_provider). Adds CApiTest.DmlGraphCaptureEmptyGraph : creates a DML session with ep.dml.enable_graph_capture=1 on a single Constant model (folds to an empty graph) and asserts session creation does not throw. The test fails without this fix and passes with it. ### Motivation and Context  #27958 refactored the graph-capture selection logic and replaced AreAllComputeNodesAssignedToCudaOrJsOrDmlEpWebGpuEp() (which returned true for an empty graph) with AreAllComputeNodesAssignedToEpOrCpu() and its new has_node_on_provider requirement. The behavior verified with a DLL-swap: onnxruntime-directml 1.24.x/1.25.2 create the session successfully; 1.27 throws. The failure surfaces through ONNX Runtime GenAI on DirectML: GenAI creates a per-device "allocator-initialization" session on a Constant-only model with DML graph capture enabled to obtain a device allocator (model.cpp). Under ORT 1.27 that session throws at construction, which breaks every GroupQueryAttention-based GenAI LLM on the DirectML EP (observed as WinML EP certification failures across DeepSeek/Llama/Qwen/Phi models --------- Co-authored-by: Aditya Rastogi <adityar@ntdev.microsoft.com>

Adrian Lizarraga (adrianlizarraga) added 10 commits April 2, 2026 13:54

First incorrect draft impl

54fe58c

Add IExecutionProvider::GetGraphCaptureNodeAssignmentPolicy() to remo…

b25bb09

…ve the hardcoded list of EPs that support graph capture from inference_session.cc

Minor update to docs. Stills need improvement

3146b85

Add onnxruntime:: namespace qualifier for DML use of enum

023c0c0

Merge branch 'main' into adrianl/PluginEp_CudaGraphCaptureReplay

bd71d7d

Implement new graph capture/replay C APIs in webgpu's OrtEp implement…

857f0b4

…ation

Add stubs with todos to cuda's plugin EP

0741b28

Use only one enum for the graph capture node assignment policy

f08d05b

Update C doc comments

211bb7e

First attempt at a webgpu plugin EP test that uses graph capture

91fb965

Adrian Lizarraga (adrianlizarraga) commented Apr 4, 2026

View reviewed changes

Comment thread onnxruntime/test/autoep/test_graph_capture.cc Outdated

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 4, 2026 00:15

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 4, 2026 00:16 View session

Copilot AI reviewed Apr 4, 2026

View reviewed changes

Comment thread onnxruntime/core/session/inference_session.cc

Comment thread onnxruntime/test/autoep/test_graph_capture.cc Outdated

Adrian Lizarraga (adrianlizarraga) added 2 commits April 6, 2026 10:59

Add ifdef guards for webgpu plugin EP test

c38ba6c

Review comments: Update NvTensorRTRTX EP to return false from IExecut…

61b8ff6

…ionProvider::IsGraphCaptureEnabled() as ORT previously never managed graph capture for this EP; Update use of webgpu name with constant

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 6, 2026 19:05

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 6, 2026 19:06 View session

Copilot AI reviewed Apr 6, 2026

View reviewed changes

Comment thread onnxruntime/core/session/inference_session.cc

Adrian Lizarraga (adrianlizarraga) added 4 commits April 6, 2026 12:34

Add InferenceSession::RunImpl() that takes in the recursive run depth

977dea1

Ensure at least one compute node is assigned to the graph capturing EP

8645b8d

Add ORT_ENFORCE

b158299

Merge branch 'main' into adrianl/PluginEp_CudaGraphCaptureReplay

0bc9cdc

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 6, 2026 20:28

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 6, 2026 20:29 View session

Copilot AI reviewed Apr 6, 2026

View reviewed changes

Comment thread onnxruntime/test/autoep/test_graph_capture.cc Outdated

Comment thread include/onnxruntime/core/session/onnxruntime_cxx_inline.h Outdated

Adrian Lizarraga (adrianlizarraga) added 2 commits April 6, 2026 16:05

Fix graph capture/replay test for the webgpu plugin EP (runs locally)

e289ae2

CXX API: add Ort::Env::CopyTensor() to copy only one tensor

ff09625

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 6, 2026 23:42

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 6, 2026 23:43 View session

Adrian Lizarraga (adrianlizarraga) added 2 commits April 9, 2026 13:22

Address doc comments

9205f22

Check that OrtEp::IsGraphCaptured and OrtEp::ReplayGraph are implemen…

5694f6c

…ted when necessary

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 9, 2026 21:14

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 9, 2026 21:15 View session

Adrian Lizarraga (adrianlizarraga) changed the title ~~[DRAFT] [Plugin EP] Port graph capture/replay APIs~~ [Plugin EP] Port graph capture/replay APIs Apr 9, 2026

Adrian Lizarraga (adrianlizarraga) marked this pull request as ready for review April 9, 2026 21:18

Adrian Lizarraga (adrianlizarraga) requested review from Edward Chen (edgchen1), Tianlei Wu (tianleiwu) and Dmitri Smirnov (yuslepukhin) April 9, 2026 21:19

Copilot AI reviewed Apr 9, 2026

View reviewed changes

Comment thread include/onnxruntime/core/session/onnxruntime_ep_c_api.h

Comment thread include/onnxruntime/core/session/onnxruntime_ep_c_api.h Outdated

Address comments

cb66889

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 9, 2026 21:25

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 9, 2026 21:26 View session

Copilot AI reviewed Apr 9, 2026

View reviewed changes

Comment thread onnxruntime/test/autoep/test_graph_capture.cc

Adrian Lizarraga (adrianlizarraga) added 2 commits April 9, 2026 14:35

Sync IO binding inputs in test

91851ac

Improved documentation for graph annotation ID

ff929f3

Adrian Lizarraga (adrianlizarraga) requested a review from Copilot April 9, 2026 21:42

Copilot started reviewing on behalf of Adrian Lizarraga (adrianlizarraga) April 9, 2026 21:43 View session

Copilot AI reviewed Apr 9, 2026

View reviewed changes

Comment thread onnxruntime/core/session/inference_session.cc

Adrian Lizarraga (adrianlizarraga) commented Apr 9, 2026

View reviewed changes

Comment thread include/onnxruntime/core/session/onnxruntime_ep_c_api.h

Adrian Lizarraga (adrianlizarraga) commented Apr 9, 2026

View reviewed changes

Comment thread include/onnxruntime/core/session/onnxruntime_ep_c_api.h

Tianlei Wu (tianleiwu) reviewed Apr 9, 2026

View reviewed changes

Comment thread onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc

Comment thread onnxruntime/test/autoep/test_selection.cc

Tianlei Wu (tianleiwu) approved these changes Apr 9, 2026

View reviewed changes

Tianlei Wu (tianleiwu) merged commit 7afe4c2 into main Apr 10, 2026
103 of 106 checks passed

Tianlei Wu (tianleiwu) deleted the adrianl/PluginEp_CudaGraphCaptureReplay branch April 10, 2026 00:16

Adrian Lizarraga (adrianlizarraga) mentioned this pull request Apr 10, 2026

FF rel-1.25.0 to latest main for cut #28000

Closed

Sanaa Hamel (sanaa-hamel-microsoft) added a commit that referenced this pull request Apr 10, 2026

Revert "[Plugin EP] Port graph capture/replay APIs (#27958)"

c29e914

This reverts commit 7afe4c2.

BrewTestBot mentioned this pull request May 8, 2026

onnxruntime 1.26.0 Homebrew/homebrew-core#281672

Merged

adrastogi mentioned this pull request Jul 1, 2026

Fix a regression in graph-capture session initialization that rejects an empty graph #29457

Merged

Uh oh!

Conversation

Adrian Lizarraga (adrianlizarraga) commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

What changed

Tests

Motivation and Context

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Review Summary

Highlights

Minor Observations (not blocking)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Adrian Lizarraga (adrianlizarraga) commented Apr 3, 2026 •

edited

Loading