Clears GPU Cache when there are no more active sessions

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

@@ -902,6 +902,10 @@ export class WebGpuBackend {
     this.sessionStatus = 'default';
   }
 
+  onCreateSession(): void {
+    this.gpuDataManager.onCreateSession();
+  }
+
   onReleaseSession(sessionId: number): void {
     this.unregisterBuffers(sessionId);
     if (this.capturedCommandList.has(sessionId)) {

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

@@ -64,11 +64,18 @@ export interface GpuDataManager {
    */
   dispose(): void;
 
+  /**
+   * create session related data.
+   */
+  onCreateSession(): void;
+
   /**
    * release session related data.
    * @param sessionId - specify the session ID.
    */
   onReleaseSession(sessionId: number): void;
+
+  sessionCount: number;
 }
 
 interface StorageCacheValue {
@@ -214,6 +221,7 @@ class GpuDataManagerImpl implements GpuDataManager {
       this.freeUniformBuffers.set(key, []);
     }
   }
+  sessionCount = 0;
 
   upload(id: GpuDataId, data: Uint8Array): void {
     const srcArrayBuffer = data.buffer;
@@ -460,6 +468,11 @@ class GpuDataManagerImpl implements GpuDataManager {
     this.capturedPendingBuffers = new Map();
   }
 
+  onCreateSession() {
+    LOG_DEBUG('verbose', () => '[WebGPU] Incrementing session count');
+    this.sessionCount += 1;
+  }
+
   onReleaseSession(sessionId: number) {
     // release the captured pending buffers.
     const pendingBuffers = this.capturedPendingBuffers.get(sessionId);
@@ -469,6 +482,17 @@ class GpuDataManagerImpl implements GpuDataManager {
       });
       this.capturedPendingBuffers.delete(sessionId);
     }
+
+    LOG_DEBUG('verbose', () => '[WebGPU] Decrementing session count');
+    // release the storage cache if no active sessions.
+    this.sessionCount -= 1;
+    if (this.sessionCount === 0) {
+      LOG_DEBUG('verbose', () => '[WebGPU] Clearing cache');
+      this.storageCache.forEach((storage) => {
+        storage.gpuData.buffer.destroy();
+      });
+      this.storageCache = new Map();
+    }
   }
 }
 

js/web/lib/wasm/wasm-core-impl.ts

@@ -317,6 +317,8 @@ export const createSession = async (
       checkLastError("Can't create a session.");
     }
 
+    wasm.jsepOnCreateSession?.();
+
     // clear current MLContext after session creation
     if (wasm.currentContext) {
       wasm.jsepRegisterMLContext!(sessionHandle, wasm.currentContext);

js/web/lib/wasm/wasm-types.ts

@@ -141,6 +141,12 @@ export declare namespace JSEP {
      * @param sessionId - specify the session ID.
      */
     jsepOnRunStart: (sessionId: number) => void;
+    /**
+     * [exported from pre-jsep.js] Create a session. This function will be called after _OrtCreateSession() is
+     * called.
+     * @returns
+     */
+    jsepOnCreateSession: () => void;
     /**
      * [exported from pre-jsep.js] Release a session. This function will be called before _OrtReleaseSession() is
      * called.

onnxruntime/wasm/pre-jsep.js

@@ -192,6 +192,9 @@ Module['jsepInit'] = (name, params) => {
     Module['jsepCreateDownloader'] = (gpuBuffer, size, type) => {
       return backend['createDownloader'](gpuBuffer, size, type);
     };
+    Module['jsepOnCreateSession'] = sessionId => {
+      backend['onCreateSession'](sessionId);
+    };
     Module['jsepOnReleaseSession'] = sessionId => {
       backend['onReleaseSession'](sessionId);
     };