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

## 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 compatibilities, 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.

Author: adrianlizarraga

include/onnxruntime/core/framework/execution_provider.h

@@ -292,6 +292,15 @@ class IExecutionProvider {
     return Status::OK();
   }
 
+  /**
+     Get the node assignment validation policy for graph capture.
+     When graph capture is enabled, ORT validates that nodes are assigned to EPs
+     in a way compatible with graph capture. This tells ORT which policy to apply.
+   */
+  virtual OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP;
+  }
+
   /**
      Called when session creation is complete
      This provides an opportunity for execution providers to optionally synchronize and

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1389,6 +1389,10 @@ struct Env : detail::Base<OrtEnv> {
                      const std::vector<Value>& dst_tensors,
                      OrtSyncStream* stream) const;  ///< Wraps OrtApi::CopyTensors
 
+  /// Wraps OrtApi::CopyTensors
+  /// Copies only one src tensor to another dst tensor.
+  Status CopyTensor(const OrtValue* src_tensor, OrtValue* dst_tensor, OrtSyncStream* stream) const;
+
   /// \brief Wraps OrtApi::SetPerSessionThreadPoolCallbacks
   /// Stores work callbacks on the Env for per-session thread pools.
   /// Only affects sessions created after this call. Does not affect global thread pools.

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -1055,6 +1055,11 @@ inline Status Env::CopyTensors(const std::vector<Value>& src_tensors,
   return Status(status);
 }
 
+inline Status Env::CopyTensor(const OrtValue* src_tensor, OrtValue* dst_tensor, OrtSyncStream* stream) const {
+  OrtStatus* status = GetApi().CopyTensors(p_, &src_tensor, &dst_tensor, stream, 1);
+  return Status(status);
+}
+
 inline UnownedAllocator Env::CreateSharedAllocator(const OrtEpDevice* ep_device, OrtDeviceMemoryType mem_type,
                                                    OrtAllocatorType allocator_type,
                                                    const OrtKeyValuePairs* allocator_options) {

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -2027,6 +2027,23 @@ typedef enum OrtEpDataLayout {
   OrtEpDataLayout_Default = OrtEpDataLayout_NCHW,
 } OrtEpDataLayout;
 
+/**
+ * \brief Node assignment policies for graph capture validation.
+ *
+ * When graph capture is enabled, ORT validates that nodes are assigned to EPs in a way that is
+ * compatible with graph capture. An EP can specify which validation policy ORT should apply.
+ *
+ * \since Version 1.26.
+ */
+typedef enum OrtGraphCaptureNodeAssignmentPolicy {
+  /** All nodes in the main graph must be assigned to this EP. No CPU fallback is allowed. */
+  OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP = 0,
+
+  /** Compute nodes must be on this EP. CPU nodes are allowed for shape computation as long as
+   *  no memory copy nodes exist. */
+  OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES = 1,
+} OrtGraphCaptureNodeAssignmentPolicy;
+
 /**
  * \brief The OrtEp struct provides functions to implement for an execution provider.
  * \since Version 1.22.
@@ -2346,6 +2363,101 @@ struct OrtEp {
    */
   ORT_API2_STATUS(CreateProfiler, _In_ OrtEp* this_ptr,
                   _Outptr_result_maybenull_ OrtEpProfilerImpl** profiler);
+
+  /** \brief Indicate whether the graph capturing mode (e.g., CUDA graph) is enabled for the provider.
+   *
+   * Graph capture allows an EP to record a sequence of device (e.g., GPU) operations during an initial run and replay
+   * them on subsequent runs, bypassing per-kernel CPU launch overhead.
+   *
+   * Applications enable graph capture via EP-specific provider options (e.g., `enable_cuda_graph=1`
+   * for the CUDA EP). An EP should return true from this function if it has been configured to enable
+   * graph capture/replay.
+   *
+   * **ORT graph capture/replay summary:**
+   * During OrtSession initialization, ORT calls OrtEp::IsGraphCaptureEnabled() on each EP in the order specified during
+   * provider registration with the session. If an EP returns true, ORT validates that the graph is suitable for
+   * graph capture, and if so, caches the EP for graph capture during the next run. The graph validation ensures
+   * that there are no control flow nodes and that node-to-EP assignments are compatible with the policy specified
+   * by the EP via OrtEp::GetGraphCaptureNodeAssignmentPolicy().
+   * Note that an OrtSession only supports graph capture for one EP (i.e., the first EP to claim support).
+   *
+   * During the first call to OrtApi::Run() for the OrtSession, ORT performs multiple internal runs of the model
+   * until the EP indicates that the graph has been captured by returning `true` from `OrtEp::IsGraphCaptured()`.
+   * If the EP is unable to capture the graph within 8 runs, the call to OrtApi::Run() returns an error OrtStatus.
+   * Each internal run invokes `OrtEp::OnRunStart()`, normal execution, and `OrtEp::OnRunEnd()`. EPs should use
+   * these run callbacks to track the number of necessary warm-up runs and begin/end graph capture when ready.
+   *
+   * After successful graph capture, subsequent calls to OrtApi::Run() skip normal execution and ORT instead calls
+   * `OrtEp::ReplayGraph()` directly.
+   *
+   * Applications can capture and replay multiple graphs (e.g., one per distinct input shape) by setting the
+   * `"gpu_graph_id"` run config entry via `OrtApi::AddRunConfigEntry()` to different integer values. ORT passes
+   * the value as the `graph_annotation_id` parameter to `OrtEp::IsGraphCaptured()` and `OrtEp::ReplayGraph()`.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \return true if graph capture mode is enabled, false otherwise.
+   *
+   * \note Implementation of this function is optional. If set to NULL, ORT assumes graph capture is not enabled.
+   * \note If this function returns true, `OrtEp::IsGraphCaptured` and `OrtEp::ReplayGraph` must also be implemented.
+   *       If either is NULL, ORT will log a warning and ignore this EP for graph capture.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(bool, IsGraphCaptureEnabled, _In_ const OrtEp* this_ptr);
+
+  /** \brief Indicate whether a graph has been captured and instantiated.
+   *
+   * ORT calls this before each `Session::Run()`. If true, ORT calls `ReplayGraph()` instead of
+   * normal execution. After a run where this returns false, ORT automatically retries until it
+   * returns true (handling warm-up runs transparently).
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \param[in] graph_annotation_id Identifies which captured graph to query.
+   *            Applications can set this value via `OrtApi::AddRunConfigEntry()` with the key `"gpu_graph_id"`.
+   *            The default value is 0 when the run config entry is not set.
+   *            Setting different IDs allows the EP to capture and manage multiple graphs (e.g., one per
+   *            distinct input shape). A value of -1 means graph capture/replay should be skipped for this run.
+   * \return true if the graph has been captured, false otherwise.
+   *
+   * \note This function must be implemented if `OrtEp::IsGraphCaptureEnabled` is implemented and may return true.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(bool, IsGraphCaptured, _In_ const OrtEp* this_ptr, _In_ int graph_annotation_id);
+
+  /** \brief Run the instantiated (captured) graph.
+   *
+   * Called by ORT instead of normal execution when `IsGraphCaptured()` returns true.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \param[in] graph_annotation_id Identifies which captured graph to replay.
+   *            Applications can set this value via `OrtApi::AddRunConfigEntry()` with the key `"gpu_graph_id"`.
+   *            The default value is 0 when the run config entry is not set.
+   *            A value of -1 means graph replay should be skipped for this run.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \note This function must be implemented if `OrtEp::IsGraphCaptureEnabled` is implemented and may return true.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API2_STATUS(ReplayGraph, _In_ OrtEp* this_ptr, _In_ int graph_annotation_id);
+
+  /** \brief Get the node assignment validation policy for graph capture.
+   *
+   * When graph capture is enabled, ORT validates that nodes are assigned to EPs in a way that is
+   * compatible with graph capture. This function tells ORT which validation policy to apply.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \return The node assignment policy for graph capture.
+   *
+   * \note Implementation of this function is optional. If set to NULL, ORT uses
+   *       OrtGraphCaptureNodeAssignmentPolicy_ALL_NODES_ON_EP (strictest validation).
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(OrtGraphCaptureNodeAssignmentPolicy, GetGraphCaptureNodeAssignmentPolicy,
+            _In_ const OrtEp* this_ptr);
 };
 
 /** \brief The function signature that ORT will call to create OrtEpFactory instances.

onnxruntime/core/providers/cuda/cuda_execution_provider.h

@@ -124,6 +124,9 @@ class CUDAExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(CudaGraphAnnotation_t graph_annotation_id) const override;
   Status ReplayGraph(CudaGraphAnnotation_t graph_annotation_id) override;
+  OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+  }
   void RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const override;
   OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override;
   std::vector<AllocatorPtr> CreatePreferredAllocators() override;

onnxruntime/core/providers/cuda/plugin/cuda_ep.cc

@@ -58,6 +58,12 @@ CudaEp::CudaEp(CudaEpFactory& factory, const Config& config, const OrtLogger& lo
   Compile = nullptr;
   ReleaseNodeComputeInfos = nullptr;
 
+  // Graph capture/replay
+  IsGraphCaptureEnabled = IsGraphCaptureEnabledImpl;
+  IsGraphCaptured = IsGraphCapturedImpl;
+  ReplayGraph = ReplayGraphImpl;
+  GetGraphCaptureNodeAssignmentPolicy = GetGraphCaptureNodeAssignmentPolicyImpl;
+
   const OrtApi& ort_api = factory_.GetOrtApi();
   Ort::Status log_status(ort_api.Logger_LogMessage(&logger_, ORT_LOGGING_LEVEL_INFO,
                                                    "CUDA Plugin EP created",
@@ -304,5 +310,30 @@ OrtStatus* ORT_API_CALL CudaEp::SyncImpl(OrtEp* this_ptr) noexcept {
   EXCEPTION_TO_STATUS_END
 }
 
+bool ORT_API_CALL CudaEp::IsGraphCaptureEnabledImpl(const OrtEp* /*this_ptr*/) noexcept {
+  // TODO: forward to EpImpl()->IsGraphCaptureEnabled()
+  return false;
+}
+
+/*static*/
+bool ORT_API_CALL CudaEp::IsGraphCapturedImpl(const OrtEp* /*this_ptr*/, int /*graph_annotation_id*/) noexcept {
+  // TODO: forward to EpImpl()->IsGraphCaptured(graph_annotation_id)
+  return false;
+}
+
+/*static*/
+OrtStatus* ORT_API_CALL CudaEp::ReplayGraphImpl(OrtEp* /*this_ptr*/, int /*graph_annotation_id*/) noexcept {
+  // TODO: forward to EpImpl()->ReplayGraph(graph_annotation_id)
+  return Ort::GetApi().CreateStatus(ORT_NOT_IMPLEMENTED,
+                                    "Graph capture replay is not yet supported in the CUDA plugin EP.");
+}
+
+/*static*/
+OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL CudaEp::GetGraphCaptureNodeAssignmentPolicyImpl(
+    const OrtEp* /*this_ptr*/) noexcept {
+  // TODO: forward to EpImpl()->GetGraphCaptureNodeAssignmentPolicy()
+  return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+}
+
 }  // namespace cuda_plugin
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/plugin/cuda_ep.h

@@ -61,6 +61,17 @@ class CudaEp : public onnxruntime::ep::adapter::Ep {
 
   static OrtStatus* ORT_API_CALL SyncImpl(OrtEp* this_ptr) noexcept;
 
+  static bool ORT_API_CALL IsGraphCaptureEnabledImpl(const OrtEp* this_ptr) noexcept;
+
+  static bool ORT_API_CALL IsGraphCapturedImpl(const OrtEp* this_ptr,
+                                               int graph_annotation_id) noexcept;
+
+  static OrtStatus* ORT_API_CALL ReplayGraphImpl(OrtEp* this_ptr,
+                                                 int graph_annotation_id) noexcept;
+
+  static OrtGraphCaptureNodeAssignmentPolicy ORT_API_CALL GetGraphCaptureNodeAssignmentPolicyImpl(
+      const OrtEp* this_ptr) noexcept;
+
   CudaEpFactory& factory_;
   std::string name_;
   Config config_;

onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h

@@ -356,6 +356,11 @@ namespace Dml
             return m_impl->ReplayGraph(graph_annotation_id);
         }
 
+        OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override
+        {
+            return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+        }
+
     private:
         ComPtr<ExecutionProviderImpl> m_impl;
     };

onnxruntime/core/providers/js/js_execution_provider.h

@@ -72,6 +72,9 @@ class JsExecutionProvider : public IExecutionProvider {
   bool IsGraphCaptureEnabled() const override;
   bool IsGraphCaptured(int graph_annotation_id) const override;
   Status ReplayGraph(int graph_annotation_id) override;
+  OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override {
+    return OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES;
+  }
 
  private:
   bool IsGraphCaptureAllowed() const;

onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc

@@ -1280,7 +1280,9 @@ void NvExecutionProvider::HandleCudaGraphStart(cudaStream_t stream, bool require
 }
 
 bool NvExecutionProvider::IsGraphCaptureEnabled() const {
-  return cuda_graph_enable_;
+  // Return false so that ORT's framework does not cache this EP for ORT-managed graph capture/replay.
+  // NvTensorRTRTX manages CUDA graph capture/replay internally.
+  return false;
 }
 
 bool NvExecutionProvider::IsGraphCaptured(int graph_annotation_id) const {