microsoft / microsoft/onnxruntime
[Web] The `device` specified in `ort.env.webgpu` will not be used at runtime
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
Based on the [type definition](https://github.com/microsoft/onnxruntime/blob/main/js/common/lib/env.ts#L228), I attempted to set a custom WebGPU device via `ort.env.webgpu`.
I wanted to set the options while requesting the device:
```js
adapter.requestDevice({ requiredFeatures: ['bgra8unorm-storage'] })
```
However, when using a [pre-allocated tensor](https://onnxruntime.ai/docs/tutorials/web/ep-webgpu.html#create-input-tensor-from-a-gpu-buffer) (i.e. a WebGPU buffer), the following error is thrown:
`backend-webgpu.ts:273 An uncaught WebGPU validation error was raised: [Buffer "data A"] is associated with [Device "my custom device"], and cannot be used with [Device].`
This happens because the `device` specified in `ort.env.webgpu` is not actually used by the ONNX WebGPU backend.
Instead, a new device is created internally, see [this line in the source code](https://github.com/microsoft/onnxruntime/blob/2cc77384c3ae536fe286cb537c1442d85dd1c5cd/js/web/lib/wasm/jsep/backend-webgpu.ts#L257C25-L257C32).
I've considered a few possible solutions to address this:
1. **Allow passing a custom `GPUDeviceDescriptor`** via `ort.env.webgpu`.
This would work, but it's somewhat restrictive and may not offer enough flexibility.
2. **Provide a helper function to acquire the device manually with custom options.**
This would let users supply their own `GPUDeviceDescriptor`, while you could still inject any required internal options.
3. **Provide a helper to construct a `GPUDeviceDescriptor`** by merging user-defined options with the defaults required by ONNX Runtime.
This approach offers a good balance between flexibility and control, while also allowing users to see and debug which options the runtime is applying.
### To reproduce
I reproduced the bug using [your example](https://github.com/microsoft/onnxruntime-inference-examples/tree/main/js/quick-start_onnxruntime-web-script-tag), with some modifications to run it on the WebGPU backend using pre-allocated buffers.
[Here is a a codepen](https://codepen.io/dsffff/pen/emJLGOQ)
### Urgency
**Note: This isn't an urgent issue**
I managed to work around it for now by **monkey-patching** (not ideal, but effective in the short term) the `navigator.gpu.requestAdapter` and `adapter.requestDevice` methods.
This allowed me to require gpu features I needed. More on the patch below.
---
### Temporary Solution (for those who need it until a proper fix is implemented)
Avoid calling `requestDevice()` yourself. Instead, let `ort.InferenceSession.create()` handle device creation internally. However, before calling it, you can **monkey-patch the WebGPU API** to inject your desired `requiredFeatures`.
Place the following code **before** creating the inference session:
```ts
const originalRequestAdapter = navigator.gpu.requestAdapter;
const patchNavigatorGpu = (gpuDeviceDescriptor: GPUDeviceDescriptor = {}) => {
const requestAdapterPath = async (
options?: GPURequestAdapterOptions
) => {
const adapter = await originalRequestAdapter.call(navigator.gpu, options)!;
if (!adapter) {
return null;
}
const requestDevice = adapter.requestDevice;
adapter.requestDevice = (descriptor: GPUDeviceDescriptor = {}) => {
return requestDevice.call(adapter, {
...descriptor,
...gpuDeviceDescriptor,
requiredFeatures: [
...descriptor.requiredFeatures || [],
...gpuDeviceDescriptor.requiredFeatures || []
]
});
};
return adapter;
}
navigator.gpu.requestAdapter = requestAdapterPath;
}
```
This allows you to override the default device creation behavior and specify any required features while still letting ONNX Runtime manage the lifecycle.
Usage
```ts
patchNavigatorGpu({ requiredFeatures: ['bgra8unorm-storage'] });
// this will acquire the device internally and set it in the ort.env.webgpu.device property
const session = await ort.InferenceSession.create(model, {
executionProviders: ['webgpu']
})
// get the device
const device = ort.env.webgpu.device;
```
### ONNX Runtime Installation
Released Package
### ONNX Runtime Version or Commit ID
1.22.0
### Execution Provider
'webgpu' (WebGPU)
Contributor guide
Research direction
Start with js/common/lib/env.ts around the webgpu type definition and js/web/lib/wasm/jsep/backend-webgpu.ts around the device creation line. Reproduce the issue using the modified quick-start_onnxruntime-web-script-tag example or the linked CodePen with a pre-allocated buffer. Done means the device configured through ort.env.webgpu is used at runtime without the device-mismatch validation error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- javascript, typescript
- Domain
- machine-learning, web-dev
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100