From d3a6117e4254c6b5fc39c760ae696426d3a23cd0 Mon Sep 17 00:00:00 2001
From: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
Date: Tue, 12 Mar 2024 19:50:51 -0700
Subject: [PATCH] [js/webgpu] expose a few properties in WebGPU API (#19857)

### Description
This change exposes a few properties in `ort.env.webgpu` to resolve
feature requirement mentioned in properties in
https://github.com/microsoft/onnxruntime/pull/14579#discussion_r1519612619.

- Add `powerPreference` and `forceFallbackAdapter` in `ort.env.webgpu`,
to allow users to set the value of the properties before the first
inference session is created.
- Add readonly property `adapter` in `ort.env.webgpu` to allow users to
get the adapter instance. Now users can access `ort.env.webgpu.device`
and `ort.env.webgpu.adapter`.

@xenova @beaufortfrancois
---
 js/common/lib/env.ts                   | 35 ++++++++++++++++++++++++++
 js/web/lib/wasm/jsep/backend-webgpu.ts |  1 +
 js/web/lib/wasm/wasm-core-impl.ts      | 10 +++++++-
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/js/common/lib/env.ts b/js/common/lib/env.ts
index 73a47d1a4f937..dd8bde2b596f4 100644
--- a/js/common/lib/env.ts
+++ b/js/common/lib/env.ts
@@ -143,9 +143,44 @@ export declare namespace Env {
        */
       ondata?: (data: WebGpuProfilingData) => void;
     };
+    /**
+     * Set or get the power preference.
+     *
+     * Setting this property only has effect before the first WebGPU inference session is created. The value will be
+     * used as options for `navigator.gpu.requestAdapter()`.
+     *
+     * See {@link https://gpuweb.github.io/gpuweb/#dictdef-gpurequestadapteroptions} for more details.
+     *
+     * @defaultValue `undefined`
+     */
+    powerPreference?: 'low-power'|'high-performance';
+    /**
+     * Set or get the force fallback adapter flag.
+     *
+     * Setting this property only has effect before the first WebGPU inference session is created. The value will be
+     * used as options for `navigator.gpu.requestAdapter()`.
+     *
+     * See {@link https://gpuweb.github.io/gpuweb/#dictdef-gpurequestadapteroptions} for more details.
+     *
+     * @defaultValue `undefined`
+     */
+    forceFallbackAdapter?: boolean;
+    /**
+     * Get the adapter for WebGPU.
+     *
+     * This property is only available after the first WebGPU inference session is created.
+     *
+     * 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.
+     *
+     * see comments on {@link GpuBufferType}
+     */
+    readonly adapter: unknown;
     /**
      * Get the device for WebGPU.
      *
+     * This property is only available after the first WebGPU inference session is created.
+     *
      * When use with TypeScript, the type of this property is `GPUDevice` defined in "@webgpu/types".
      * Use `const device = env.webgpu.device as GPUDevice;` in TypeScript to access this property with correct type.
      *
diff --git a/js/web/lib/wasm/jsep/backend-webgpu.ts b/js/web/lib/wasm/jsep/backend-webgpu.ts
index 27c5566ab9fed..182c1cd351c9d 100644
--- a/js/web/lib/wasm/jsep/backend-webgpu.ts
+++ b/js/web/lib/wasm/jsep/backend-webgpu.ts
@@ -231,6 +231,7 @@ export class WebGpuBackend {
     };
 
     Object.defineProperty(this.env.webgpu, 'device', {value: this.device});
+    Object.defineProperty(this.env.webgpu, 'adapter', {value: adapter});
 
     // init queryType, which is necessary for InferenceSession.create
     this.setQueryType();
diff --git a/js/web/lib/wasm/wasm-core-impl.ts b/js/web/lib/wasm/wasm-core-impl.ts
index 37b9ed6a1002f..afab9ba00b0c4 100644
--- a/js/web/lib/wasm/wasm-core-impl.ts
+++ b/js/web/lib/wasm/wasm-core-impl.ts
@@ -89,7 +89,15 @@ export const initEp = async(env: Env, epName: string): Promise<void> => {
     if (typeof navigator === 'undefined' || !navigator.gpu) {
       throw new Error('WebGPU is not supported in current environment');
     }
-    const adapter = await navigator.gpu.requestAdapter();
+    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}"`);
+    }
+    const 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.');