[js/webgpu] Support capture and replay for jsep

js/common/lib/inference-session.ts

@@ -111,7 +111,7 @@ export declare namespace InferenceSession {
     optimizedModelFilePath?: string;
 
     /**
-     * Wether enable profiling.
+     * Whether enable profiling.
      *
      * This setting is a placeholder for a future use.
      */
@@ -154,6 +154,12 @@ export declare namespace InferenceSession {
      */
     preferredOutputLocation?: OnnxValueDataLocation|{readonly [outputName: string]: OnnxValueDataLocation};
 
+    /**
+     * Whether enable graph capture.
+     * This setting is available only in ONNXRuntime Web for WebGPU EP.
+     */
+    enableGraphCapture?: boolean;
+
     /**
      * Store configurations for a session. See
      * https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/
@@ -238,7 +244,6 @@ export declare namespace InferenceSession {
   export interface WebGpuExecutionProviderOption extends ExecutionProviderOption {
     readonly name: 'webgpu';
     preferredLayout?: 'NCHW'|'NHWC';
-    graphCaptureEnabled?: boolean;
   }
   export interface WebNNExecutionProviderOption extends ExecutionProviderOption {
     readonly name: 'webnn';

js/web/lib/wasm/binding/ort-wasm.d.ts

@@ -13,9 +13,9 @@ export declare namespace JSEP {
   type ReleaseKernelFunction = (kernel: number) => void;
   type RunFunction =
       (kernel: number, contextDataOffset: number, sessionHandle: number, errors: Array<Promise<string|null>>) => number;
-  type CaptureBeginFunction = (sessionHandle: number) => void;
-  type CaptureEndFunction = (sessionHandle: number) => void;
-  type ReplayFunction = (sessionHandle: number) => void;
+  type CaptureBeginFunction = () => void;
+  type CaptureEndFunction = () => void;
+  type ReplayFunction = () => void;
 }
 
 export interface OrtWasmModule extends EmscriptenModule {
@@ -181,11 +181,14 @@ export interface OrtWasmModule extends EmscriptenModule {
       (gpuBuffer: GPUBuffer, size: number,
        type: Tensor.GpuBufferDataTypes) => () => Promise<Tensor.DataTypeMap[Tensor.GpuBufferDataTypes]>;
   /**
-   *  [exported from js_internal_api.js] Called when InferenceSession.run started.
+   *  [exported from js_internal_api.js] Called when InferenceSession.run started. This function will be called before
+   * _OrtRun[WithBinding]() is called.
+   * @param sessionId - specify the session ID.
    */
-  jsepOnRunStart: () => void;
+  jsepOnRunStart: (sessionId: number) => void;
   /**
-   * [exported from js_internal_api.js] Release a session.
+   * [exported from js_internal_api.js] Release a session. This function will be called before _OrtReleaseSession() is
+   * called.
    * @param sessionId - specify the session ID.
    * @returns
    */

js/web/lib/wasm/jsep/backend-webgpu.ts

@@ -10,7 +10,7 @@ import {createView, TensorView} from './tensor-view';
 import {createGpuDataManager, downloadGpuData, GpuDataManager} from './webgpu/gpu-data-manager';
 import {RunFunction, WEBGPU_OP_RESOLVE_RULES} from './webgpu/op-resolve-rules';
 import {ProgramManager} from './webgpu/program-manager';
-import {ComputeContext, GpuData, ProgramInfo, ProgramInputTensorInfoDependency, StatusType, TimestampQuery} from './webgpu/types';
+import {ComputeContext, GpuData, ProgramInfo, ProgramInputTensorInfoDependency, SessionState, TimestampQuery} from './webgpu/types';
 
 interface CommandInfo {
   readonly kernelId: number;
@@ -111,9 +111,9 @@ export class WebGpuBackend {
   programManager: ProgramManager;
 
   /**
-   * representing the session ID of which is currently being captured/replay.
-   * `null` means no session is being captured.
-   * only valid when captureGraphEnabled = true.
+   * representing the session ID of which is currently being run.
+   * `null` means no session is being run.
+   * only valid when session.run is executed.
    */
   currentSessionId: number|null = null;
 
@@ -169,9 +169,9 @@ export class WebGpuBackend {
   queryType: TimestampQuery;
 
   env: Env;
-  status: StatusType = StatusType.default;
+  sessionStatus: SessionState = 'default';
   /**
-   * a SessionID -> CommandInfo[] mapping.
+   * a SessionID -> CommandInfo[] mapping. It's used to record all GPU commands for corresponding session.
    */
   capturedCommandList: Map<number, CommandInfo[]> = new Map();
 
@@ -519,7 +519,7 @@ export class WebGpuBackend {
         () => `[ProgramManager] run "${program.name}" (key=${key}) with ${normalizedDispatchGroup[0]}x${
             normalizedDispatchGroup[1]}x${normalizedDispatchGroup[2]}`);
 
-    if (this.queryType !== 'none' || this.status === StatusType.capture) {
+    if (this.queryType !== 'none' || this.sessionStatus === 'capturing') {
       const pendingKernelInfo: PendingKernelInfo = {
         kernelId: this.currentKernelId!,
         programName: artifact.programInfo.name,
@@ -701,32 +701,31 @@ export class WebGpuBackend {
     }
   }
 
-  captureBegin(sessionHandle: number): void {
-    LOG_DEBUG('info', () => `captureBegin ${sessionHandle}`);
-    this.currentSessionId = sessionHandle;
-    let sessionCommandList = this.capturedCommandList.get(sessionHandle);
-    let sessionPendingKernels = this.capturedPendingKernels.get(sessionHandle);
+  captureBegin(): void {
+    LOG_DEBUG('info', () => 'captureBegin');
+    let sessionCommandList = this.capturedCommandList.get(this.currentSessionId!);
+    let sessionPendingKernels = this.capturedPendingKernels.get(this.currentSessionId!);
     if (!sessionCommandList) {
       sessionCommandList = [];
-      this.capturedCommandList.set(sessionHandle, sessionCommandList);
+      this.capturedCommandList.set(this.currentSessionId!, sessionCommandList);
       sessionPendingKernels = [];
-      this.capturedPendingKernels.set(sessionHandle, sessionPendingKernels);
+      this.capturedPendingKernels.set(this.currentSessionId!, sessionPendingKernels);
     }
-    this.status = StatusType.capture;
+    // flush the left commands before we change the status.
+    this.flush();
+    this.sessionStatus = 'capturing';
   }
-  captureEnd(sessionHandle: number): void {
-    LOG_DEBUG('info', () => `captureEnd ${sessionHandle}`);
+  captureEnd(): void {
+    LOG_DEBUG('info', () => 'captureEnd');
     // flush the left commands before we change the status.
     this.flush();
-    this.currentSessionId = null;
-    this.status = StatusType.default;
+    this.sessionStatus = 'default';
   }
-  replay(sessionHandle: number): void {
-    LOG_DEBUG('info', () => `replay ${sessionHandle}`);
-    this.currentSessionId = sessionHandle;
-    this.status = StatusType.replay;
-    const sessionCommandList = this.capturedCommandList.get(sessionHandle);
-    const sessionPendingKernels = this.capturedPendingKernels.get(sessionHandle);
+  replay(): void {
+    LOG_DEBUG('info', () => 'replay');
+    this.sessionStatus = 'replaying';
+    const sessionCommandList = this.capturedCommandList.get(this.currentSessionId!);
+    const sessionPendingKernels = this.capturedPendingKernels.get(this.currentSessionId!);
     const length = sessionCommandList!.length;
     this.pendingKernels = [];
     for (let i = 0; i < length; i++) {
@@ -750,7 +749,7 @@ export class WebGpuBackend {
     }
     // flush the left commands before we change the status.
     this.flush();
-    this.status = StatusType.default;
+    this.sessionStatus = 'default';
   }
 
   onReleaseSession(sessionId: number): void {
@@ -764,7 +763,8 @@ export class WebGpuBackend {
     this.gpuDataManager.onReleaseSession(sessionId);
   }
 
-  onRunStart(): void {
+  onRunStart(sessionId: number): void {
+    this.currentSessionId = sessionId;
     this.setQueryType();
   }
 }

js/web/lib/wasm/jsep/init.ts

@@ -203,9 +203,9 @@ export const init = async(module: OrtWasmModule, env: Env, gpuAdapter: GPUAdapte
         return backend.computeKernel(kernel, context, errors);
       },
       // jsepCaptureBegin
-      (sessionHandle: number) => backend.captureBegin(sessionHandle),
+      () => backend.captureBegin(),
       // jsepCaptureEnd
-      (sessionHandle: number) => backend.captureEnd(sessionHandle),
+      () => backend.captureEnd(),
       // jsepReplay
-      (sessionHandle: number) => backend.replay(sessionHandle));
+      () => backend.replay());
 };

js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts

@@ -4,7 +4,7 @@
 import {WebGpuBackend} from '../backend-webgpu';
 import {LOG_DEBUG} from '../log';
 
-import {GpuData, GpuDataId, GpuDataType, StatusType} from './types';
+import {GpuData, GpuDataId, GpuDataType} from './types';
 
 /**
  * manages GpuDataId -> GpuBuffer
@@ -331,7 +331,7 @@ class GpuDataManagerImpl implements GpuDataManager {
       return;
     }
 
-    if (this.backend.status === StatusType.default) {
+    if (this.backend.sessionStatus === 'default') {
       for (const buffer of this.buffersPending) {
         // eslint-disable-next-line no-bitwise
         if ((buffer.usage & GPUBufferUsage.STORAGE) === GPUBufferUsage.STORAGE) {

js/web/lib/wasm/jsep/webgpu/program-manager.ts

@@ -7,7 +7,7 @@ import {WebGpuBackend} from '../backend-webgpu';
 import {LOG_DEBUG} from '../log';
 
 import {createShaderHelper} from './ops/common';
-import {Artifact, GpuData, ProgramInfo, StatusType} from './types';
+import {Artifact, GpuData, ProgramInfo} from './types';
 
 /**
  * ProgramManager is the main class behind running computations
@@ -51,7 +51,7 @@ export class ProgramManager {
     const bindGroup = device.createBindGroup(
         {layout: buildArtifact.computePipeline.getBindGroupLayout(0), entries, label: buildArtifact.programInfo.name});
 
-    if (this.backend.status === StatusType.capture) {
+    if (this.backend.sessionStatus === 'capturing') {
       const commandInfo = {
         kernelId: this.backend.currentKernelId!,
         computePipeline: buildArtifact.computePipeline,

js/web/lib/wasm/jsep/webgpu/types.ts

@@ -5,11 +5,7 @@ import {TensorView} from '../tensor-view';
 
 import {ShaderHelper} from './ops/common';
 
-export enum StatusType {
-  default = 0,
-  capture = 1,
-  replay = 2
-}
+export type SessionState = 'default'|'capturing'|'replaying';
 
 export enum GpuDataType {
   default = 0,

js/web/lib/wasm/session-options.ts

@@ -112,18 +112,6 @@ const setExecutionProviders =
                       `Can't set a session config entry: 'preferredLayout' - ${webgpuOptions.preferredLayout}.`);
                 }
               }
-              if (webgpuOptions?.graphCaptureEnabled !== undefined) {
-                if (typeof webgpuOptions.graphCaptureEnabled !== 'boolean') {
-                  throw new Error(`graphCaptureEnabled must be a boolean value: ${webgpuOptions.graphCaptureEnabled}`);
-                }
-                const keyDataOffset = allocWasmString('graphCaptureEnabled', allocs);
-                const valueDataOffset = allocWasmString(webgpuOptions.graphCaptureEnabled.toString(), allocs);
-                if (getInstance()._OrtAddSessionConfigEntry(sessionOptionsHandle, keyDataOffset, valueDataOffset) !==
-                    0) {
-                  checkLastError(`Can't set a session config entry: 'graphCaptureEnabled' - ${
-                      webgpuOptions.graphCaptureEnabled}.`);
-                }
-              }
             }
             break;
           case 'wasm':
@@ -180,6 +168,18 @@ export const setSessionOptions = (options?: InferenceSession.SessionOptions): [n
       setExecutionProviders(sessionOptionsHandle, sessionOptions.executionProviders, allocs);
     }
 
+    if (sessionOptions.enableGraphCapture !== undefined) {
+      if (typeof sessionOptions.enableGraphCapture !== 'boolean') {
+        throw new Error(`enableGraphCapture must be a boolean value: ${sessionOptions.enableGraphCapture}`);
+      }
+      const keyDataOffset = allocWasmString('enableGraphCapture', allocs);
+      const valueDataOffset = allocWasmString(sessionOptions.enableGraphCapture.toString(), allocs);
+      if (wasm._OrtAddSessionConfigEntry(sessionOptionsHandle, keyDataOffset, valueDataOffset) !== 0) {
+        checkLastError(
+            `Can't set a session config entry: 'enableGraphCapture' - ${sessionOptions.enableGraphCapture}.`);
+      }
+    }
+
     if (sessionOptions.freeDimensionOverrides) {
       for (const [name, value] of Object.entries(sessionOptions.freeDimensionOverrides)) {
         if (typeof name !== 'string') {