Repository navigation
Fix WebGPU context memory identity and test host-owned caches - #33110
Conversation
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Copilot review overview
Review effort: Lite
Findings: 1
Open (3)
What changed in this PR
This PR fixes WebGPU stream validation for custom external-Dawn contexts and expands the external host test coverage to validate host-owned devices and shader-cache callbacks across multiple devices/isolation keys.
Changes:
- Validate WebGPU session streams against the session allocator’s actual memory-device ID (rather than the ORT context/deviceId).
- Extend the external Dawn host test to create and inject host-owned Dawn instances/devices, exercise cache callbacks, and verify device usability after ORT teardown.
- Update CI and build wiring to link Dawn’s proc table library and run the new host-device/cache test modes.
| File | Description |
|---|---|
| onnxruntime/test/webgpu/external_dawn/main.cc | Adds host-owned device + shader cache callback testing (cold/warm/isolated) and a required-feature rejection case. |
| onnxruntime/core/providers/webgpu/ep/ep.cc | Fixes stream validation to compare against allocator memory-device ID instead of the custom context ID. |
| docs/WebGPU_External_Dawn.md | Documents host-owned device handles, deviceId semantics, callback lifetime/ownership, and new test modes. |
| cmake/onnxruntime_unittests.cmake | Links the external Dawn host test against dawn::dawn_proc to support direct proc-table calls. |
| .github/workflows/windows_webgpu.yml | Runs the new --host_device and --no_implicit_sync host-owned device/cache CI checks (source, package, plugin). |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
Native Dawn ownership, plugin teardown, and cache callback behavior warrant final human review despite comprehensive Windows coverage.
Review effort: Balanced
Findings: None
Resolved since last review (3)
|
Hi Mustapha Jaber (@mustjab), I have two questions about the intended deployment scenario:
|
|
Jiajia Qin (@qjia7) I wonder whether this is more closely tied to multi-context/multi-GPU support. The allocator/context device-ID mismatch appears to have existed for a long time, and the new check seems to be what exposed it. Would it make sense to remove or relax this check as a minimal unblock, and address distinct allocator/tensor memory-device IDs together with multi-context or multi-GPU support? |
Mustapha Jaber (mustjab)
changed the title
|
Xiaofei Han (@xiaofeihan1), replying to your comment: Our current deployment uses external Dawn with the built-in WebGPU EP and a |
|
Xiaofei Han (@xiaofeihan1), replying to your comment: The stream check exposed an existing allocator/context identity mismatch. Following Jiajia's feedback, this PR now propagates the context ID through the I would keep that complete fix and the concurrent default/external-session |
|
Cool, I'm OK with this PR. I'm just curious how you identified this issue. It seems to be caused by the external Dawn + WebGPU Plugin EP setup, but you mentioned that you're using the built-in EP. |
## Summary Follow-up to the review of #33110. The shared WebGPU data-transfer wrapper does not need to own a context. A session stream already supplies its owning execution provider's transfer state, while a streamless copy can select its registered context from the GPU tensor's memory-device ID. - Remove the cached context, initialization mutex, and context-ID argument from the internal transfer creator. - Route session copies directly through their stream. Retain the tensor-selected registered context only for each streamless copy, including context zero, and flush successful streamless GPU-to-GPU copies. - Document that memory-device identity does not register Dawn objects, the current Env/global allocator context-zero restriction, and the registration requirements for external-context copies. Allocator/session context IDs and the public C API are unchanged. This does not add pre-session external-context registration. ## Validation - Scoped C++ lintrunner and staged whitespace checks passed. - Fresh Windows x64 Release external-Dawn plugin build on upstream main `864d87126864c64651f58a9b51d9006ad71b2c78` passed. - All five native-host modes passed against the fresh ORT/plugin DLLs: baseline, host-device, missing-sync, callback-failure, and missing-proc-table. The host executable was reused after verifying its source was unchanged. - Host-device coverage includes coexisting default/external context allocation and transfers, shader-cache reuse/isolation, and device usability after ORT teardown. Negative modes rejected invalid inputs or contained callback failures as expected. - The same follow-up also passed the five host modes in both built-in and plugin configurations before rebasing onto current main.
Summary
Native external-Dawn hosts can supply a Dawn instance and device with a positive
ORT context ID. WebGPU allocators previously advertised memory-device ID zero
regardless of that context, leaving allocator lookup, streams, and transfer
routing unable to distinguish separate Dawn memory domains. Plugin stream
validation exposed this mismatch during custom-context session initialization.
This PR fixes the identity mismatch rather than weakening the stream check:
default memory device, with range validation before context creation.
owning-context identity check.
copies between distinct contexts or foreign GPU vendors.
cache callbacks, cache isolation keys, lifetime checks, and failure cases.
and GPU-output IDs, CPU/GPU and same-context GPU/GPU transfers, and rejection
of cross-context copies.
built-in, and plugin CI paths, and document device/callback ownership.
Our current deployment uses built-in WebGPU with external Dawn and a host-owned
device. Plugin and concurrent-context tests are additional regression coverage,
not a claim that the deployment has migrated to plugins or requires a default
session alongside the host-owned session. Context IDs identify Dawn memory
domains, not physical GPU indices. No new public API is introduced.
Validation
Windows x64 Release with native Dawn:
plugin builds passed all five host modes, with CPU fallback disabled.
same-context transfers, and cross-context rejection.
different isolation key, and verified host-device use after ORT teardown.
rejection passed with their expected diagnostics.
forwarding regression.
Other native platforms remain unvalidated locally. Cache reuse is tested across
devices within one process; persistence across process restarts is not covered.