[js/webgpu] allow setting env.webgpu.adapter (#19940)

### Description
Allow user to set `env.webgpu.adapter` before creating the first
inference session.

Feature request:
https://github.com/microsoft/onnxruntime/pull/19857#issuecomment-1999984753

@xenova
This commit is contained in:
Yulong Wang 2024-03-19 12:55:00 -07:00 committed by GitHub
parent 8293aa1564
commit 01c7aaf6aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 35 additions and 16 deletions

View file

@ -166,16 +166,20 @@ export declare namespace Env {
*/ */
forceFallbackAdapter?: boolean; forceFallbackAdapter?: boolean;
/** /**
* Get the adapter for WebGPU. * Set or get the adapter for WebGPU.
* *
* This property is only available after the first WebGPU inference session is created. * Setting this property only has effect before the first WebGPU inference session is created. The value will be
* used as the GPU adapter for the underlying WebGPU backend to create GPU device.
*
* If this property is not set, it will be available to get after the first WebGPU inference session is created. The
* value will be the GPU adapter that created by the underlying WebGPU backend.
* *
* When use with TypeScript, the type of this property is `GPUAdapter` defined in "@webgpu/types". * When use with TypeScript, the type of this property is `GPUAdapter` defined in "@webgpu/types".
* Use `const adapter = env.webgpu.adapter as GPUAdapter;` in TypeScript to access this property with correct type. * Use `const adapter = env.webgpu.adapter as GPUAdapter;` in TypeScript to access this property with correct type.
* *
* see comments on {@link Tensor.GpuBufferType} * see comments on {@link Tensor.GpuBufferType}
*/ */
readonly adapter: unknown; adapter: unknown;
/** /**
* Get the device for WebGPU. * Get the device for WebGPU.
* *

View file

@ -252,8 +252,10 @@ export class WebGpuBackend {
} }
}; };
Object.defineProperty(this.env.webgpu, 'device', {value: this.device}); Object.defineProperty(
Object.defineProperty(this.env.webgpu, 'adapter', {value: adapter}); this.env.webgpu, 'device', {value: this.device, writable: false, enumerable: true, configurable: false});
Object.defineProperty(
this.env.webgpu, 'adapter', {value: adapter, writable: false, enumerable: true, configurable: false});
// init queryType, which is necessary for InferenceSession.create // init queryType, which is necessary for InferenceSession.create
this.setQueryType(); this.setQueryType();

View file

@ -93,18 +93,31 @@ export const initEp = async(env: Env, epName: string): Promise<void> => {
if (typeof navigator === 'undefined' || !navigator.gpu) { if (typeof navigator === 'undefined' || !navigator.gpu) {
throw new Error('WebGPU is not supported in current environment'); throw new Error('WebGPU is not supported in current environment');
} }
const powerPreference = env.webgpu?.powerPreference;
if (powerPreference !== undefined && powerPreference !== 'low-power' && powerPreference !== 'high-performance') { let adapter = env.webgpu.adapter as GPUAdapter | null;
throw new Error(`Invalid powerPreference setting: "${powerPreference}"`);
}
const forceFallbackAdapter = env.webgpu?.forceFallbackAdapter;
if (forceFallbackAdapter !== undefined && typeof forceFallbackAdapter !== 'boolean') {
throw new Error(`Invalid forceFallbackAdapter setting: "${forceFallbackAdapter}"`);
}
const adapter = await navigator.gpu.requestAdapter({powerPreference, forceFallbackAdapter});
if (!adapter) { if (!adapter) {
throw new Error( // if adapter is not set, request a new adapter.
'Failed to get GPU adapter. You may need to enable flag "--enable-unsafe-webgpu" if you are using Chrome.'); const powerPreference = env.webgpu.powerPreference;
if (powerPreference !== undefined && powerPreference !== 'low-power' &&
powerPreference !== 'high-performance') {
throw new Error(`Invalid powerPreference setting: "${powerPreference}"`);
}
const forceFallbackAdapter = env.webgpu.forceFallbackAdapter;
if (forceFallbackAdapter !== undefined && typeof forceFallbackAdapter !== 'boolean') {
throw new Error(`Invalid forceFallbackAdapter setting: "${forceFallbackAdapter}"`);
}
adapter = await navigator.gpu.requestAdapter({powerPreference, forceFallbackAdapter});
if (!adapter) {
throw new Error(
'Failed to get GPU adapter. ' +
'You may need to enable flag "--enable-unsafe-webgpu" if you are using Chrome.');
}
} else {
// if adapter is set, validate it.
if (typeof adapter.limits !== 'object' || typeof adapter.features !== 'object' ||
typeof adapter.requestDevice !== 'function') {
throw new Error('Invalid GPU adapter set in `env.webgpu.adapter`. It must be a GPUAdapter object.');
}
} }
if (!env.wasm.simd) { if (!env.wasm.simd) {