[WebGPU] Expose safe graph capture I/O in Python

[Ananya Anand (4n4ny4)]

Description

Expose a safe Python workflow for ONNX Runtime’s existing WebGPU graph capture and replay functionality.

This change:

• Adds session-owned WebGPU OrtValue allocation.
• Supports in-place CPU-to-GPU and same-session GPU-to-GPU updates.
• Requires fixed WebGPU inputs and outputs through IOBinding.
• Adds explicit captured-graph release.
• Preserves session, allocator, OrtValue, and IOBinding lifetimes.
• Rejects unsafe raw-pointer, DLPack, cross-session, and buffer-rebinding operations.
• Adds focused Python coverage and a reusable graph-capture benchmark.

Motivation and Context

WebGPU graph replay retains the buffer bindings and addresses used during capture. Python previously lacked a safe way to allocate session-owned WebGPU buffers, update their contents between runs, retrieve outputs explicitly, and release captured resources.

Using ordinary transient inputs or rebinding outputs after capture can cause replay to continue using the original buffers, resulting in stale inputs, unexpected outputs, or unsafe resource lifetimes.

This change provides a supported Python workflow based on fixed session-owned buffers:

1. Allocate and bind WebGPU input/output OrtValues.
2. Capture the graph during the initial run.
3. Update input contents in place.
4. Replay using run_with_iobinding.
5. Copy outputs explicitly to CPU when needed.
6. Release the captured graph before rebinding or destroying buffers.

On a TITAN V development benchmark, WebGPU graph capture reduced steady-state p50 latency from 5.506 ms to 5.014 ms, an improvement of approximately 8.9%.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR exposes a safe Python workflow for WebGPU graph capture/replay by introducing session-owned device OrtValue allocation, enforcing session provenance for session-scoped buffers, and restricting unsafe I/O patterns (raw pointers, DLPack export, cross-session use, and output rebinding) that can break captured replay semantics.

Changes:

• Add session-owned OrtValue creation/update APIs and a Python-facing release_captured_graph() workflow to manage captured resources safely.
• Enforce provenance/lifetime rules for session-scoped and WebGPU OrtValues across run, run_async, run_with_ort_values, IOBinding, and OrtValueVector usage patterns.
• Add focused Python tests plus a reusable WebGPU graph-capture benchmark script demonstrating fixed device I/O and capture/replay.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/test/python/webgpu_graph_capture_benchmark.py	Adds a standalone benchmark demonstrating fixed WebGPU I/O binding, capture/replay, in-place updates, and explicit output readback.
onnxruntime/test/python/onnxruntime_test_python.py	Adds unit tests covering session-scoped OrtValue semantics and WebGPU graph-capture safety restrictions.
onnxruntime/python/onnxruntime_pybind_state.cc	Adds pybind methods for session-owned OrtValue allocation/update plus WebGPU graph-capture status/release plumbing.
onnxruntime/python/onnxruntime_pybind_ortvalue.cc	Blocks unsafe WebGPU operations (data_ptr, numpy, DLPack) and prevents putting WebGPU OrtValues into raw OrtValueVector.
onnxruntime/python/onnxruntime_pybind_iobinding.cc	Ensures WebGPU binding uses the correct allocator name and keeps the session alive for SessionIOBinding.
onnxruntime/python/onnxruntime_inference_collection.py	Adds Python APIs for session-owned OrtValues, graph-capture release, and provenance validation/enforcement in Session, IOBinding, and OrtValue.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Jiajia Qin (qjia7)]

Review frame

• Review contract: Reviewed PR 32074 at exact head a30855282ecb8da56d8d6cdd4cb3b74165f0575a against its recorded base 3d3cfa6c004551ee31d3acff6d01fb572712ca91, including all six changed files, issue discussion, reviews, inline discussion, and head CI. GitHub reports only the passing license/cla check.
• Problem/feature validity: Validated. WebGPU replay retains the bind groups/buffer handles recorded during capture, while the pre-PR Python allocation path does not obtain the allocator from a particular session. Python therefore needs a session-scoped way to allocate and update fixed WebGPU inputs/outputs and to release captured resources.
• Risk/scope: Deep. This adds public Python APIs and changes allocator, data-transfer, I/O-binding, graph-ID, and object-lifetime behavior across Python and C++.
• Direction gate: Pass, with incomplete enforcement. Session allocator lookup plus Python-side session provenance is the right ownership direction and is consistent with InferenceSession::GetAllocator, SessionIOBinding, and the WebGPU EP's per-session context. The owner solution must additionally associate each captured graph ID with immutable I/O buffer identities until release; C1 is the missing part of that safety boundary. The implementation also has the compatibility and transfer-routing gaps below.

Confirmed findings

C1 [P2]: Reject a different or mutated I/O binding while a graph ID is captured

Suggested inline location: onnxruntime/python/onnxruntime_inference_collection.py:534

Suggested comment:

[P2] Please associate each captured graph ID with the IOBinding and fixed OrtValues used for capture, and reject a different/rebound/cleared binding until release_captured_graph(id) succeeds. This check currently establishes only that the binding belongs to the same session. Once graph 0 is captured, InferenceSession::RunImpl replays graph 0 by ID and WebGPU's ReplayGraph submits the previously captured bind groups; it never consults the buffers in the newly supplied binding. The new test demonstrates the resulting silent failure at lines 2069-2079: running alternate_io_binding updates the original output and leaves the alternate output zero. Clearing the original binding before release can additionally return those buffers to WebGPU's cache while the captured bind group still identifies them. This contradicts the PR's claim that unsafe buffer-rebinding operations are rejected. Please enforce the lifecycle in the API (or expose a capture object that owns the fixed binding), and change this test to expect rejection rather than stale output.

Evidence and attribution: The PR exposes the previously unavailable session-owned buffers and advertises a safe Python workflow. Session.run_with_iobinding adds only a session-identity check. The PR's own test proves that another same-session binding is accepted and silently ignored by replay.

S1 [P2]: Preserve the documented per-run gpu_graph_id=-1 escape hatch

Suggested inline location: onnxruntime/python/onnxruntime_inference_collection.py:323

Suggested comment:

[P2] Please make this validation depend on the effective gpu_graph_id instead of rejecting every run whenever the session option enables capture. Core defines kGraphAnnotationSkip = -1, and CachedExecutionProviderForGraphReplay::AllowGraphCaptureOnRun() returns false for that ID; WebGPU's OnRunStart likewise does not create a per-graph manager or begin capture for -1. Such a run is therefore the supported non-capturing path and transient feeds are safe. Before this PR, Python callers could use it, but run, run_async, and run_with_ort_values now reject it before C++ sees the run options, and the bind-time IOBinding checks also make a skipped run with ordinary bindings impossible. Please allow gpu_graph_id=-1, defer fixed-buffer validation until run_with_iobinding where the run options are known, and add regression coverage for both convenience-run and IOBinding skip paths.

Evidence and attribution: InferenceSession::CachedExecutionProviderForGraphReplay explicitly reserves -1 as kGraphAnnotationSkip, and the recursive capture/replay path is guarded by AllowGraphCaptureOnRun. The unconditional Python guard is added by this PR and removes that existing behavior.

C2 [P2]: Do not let the session-wide transfer search route CUDA buffers through WebGPU

Suggested inline location: onnxruntime/python/onnxruntime_pybind_state.cc:144

Suggested comment:

[P2] Please ensure this session-scoped copy selects the transfer implementation for the tensors' actual provider. These new creation/update APIs are generic (the CPU path is tested and device_type="cuda" is accepted), but DataTransferManager::CopyTensor chooses the first registered transfer whose CanCopy returns true. The built-in WebGPU DataTransfer::CanCopy accepts every CPU/GPU or GPU/GPU pair without checking vendor ID. In a session registered as [WebGpuExecutionProvider, CUDAExecutionProvider], updating a session-created CUDA OrtValue therefore selects WebGPU first; DataTransferImpl::CopyTensor casts the CUDA allocation to WGPUBuffer, and BufferManager::Upload passes it to wgpuBufferGetMappedRange. The plugin WebGPU transfer already avoids this by requiring vendor ID 0. Please mirror that vendor/device check in the built-in transfer (or restrict this API to locations it can select unambiguously) and add a mixed-provider regression test.

Evidence and attribution: The PR's CopySessionOrtValue rejects only GPU-to-GPU pairs where exactly one memory-info name is WEBGPU_BUFFER; CPU-to-CUDA and CUDA-to-CUDA take the unrestricted session-wide search. The unsafe WebGPU cast is a direct consequence. The PR newly exposes this path for arbitrary session allocators and routes all session-scoped updates through it.

Clarifications

Q1: The reported 8.9% speedup is not reproducible from the added benchmark

webgpu_graph_capture_benchmark.py unconditionally enables graph capture at line 53 and has no baseline mode, so the same script cannot produce the stated 5.506 ms non-capture result. Please provide the exact baseline/capture commands and raw result JSON (including backend and input-copy mode), or add a capture toggle so the comparison uses identical fixed I/O, readback, warmup, and timing boundaries.

Test coverage

• C1: The WebGPU test exercises the wrong-binding case, but asserts the silent stale-output behavior instead of enforcing the advertised safety invariant.
• S1: No test covers gpu_graph_id=-1 on a graph-enabled WebGPU session.
• C2: No test covers a session with WebGPU followed by another GPU provider.
• The positive WebGPU test is skipped when the EP or adapter is unavailable. At review time GitHub reports only license/cla; there is no WebGPU-enabled build/test result attached to the PR.
• Local static verification: both changed Python modules compile with Python's compile() and git diff --check is clean. I did not claim runtime or C++ build success because the available local build does not contain a PR-compatible Python WebGPU extension.

Verdict

The feature is valid and the session-provenance design direction is appropriate. C1 and C2 block the advertised safe workflow because accepted calls can silently target old buffers or pass a foreign GPU handle to Dawn. S1 is an existing graph-control compatibility regression and should also be fixed before merge. Q1 remains a performance-evidence clarification. I found no cleanup-only items worth posting.

[Jiajia Qin (qjia7)]

WebGPU OrtValues can also be created before a session through an environment-registered shared allocator. Please distinguish session-owned WebGPU buffers from shared WebGPU buffers instead of requiring every WebGPU OrtValue to belong to this session.