w

cmake/onnxruntime_webassembly.cmake

@@ -434,7 +434,7 @@ jsepDownload:_pp_")
       "SHELL:-s ASSERTIONS=0"
       "SHELL:-s SAFE_HEAP=0"
       "SHELL:-s STACK_OVERFLOW_CHECK=0"
-      --closure 1
+      # --closure 1
     )
   endif()
 

cmake/patches/dawn/dawn.patch

@@ -34,8 +34,85 @@ index efd6491cd6..8ebc5d28b6 100644
  emcmake cmake -GNinja -DDAWN_EMSCRIPTEN_TOOLCHAIN="path/to/emscripten" ../..
 
  ninja
+diff --git a/third_party/emdawnwebgpu/library_webgpu.js b/third_party/emdawnwebgpu/library_webgpu.js
+index 5862ce4045..5c59dbd223 100644
+--- a/third_party/emdawnwebgpu/library_webgpu.js
++++ b/third_party/emdawnwebgpu/library_webgpu.js
+@@ -584,6 +584,10 @@ var LibraryWebGPU = {
+
+   emwgpuDelete__sig: 'vp',
+   emwgpuDelete: (ptr) => {
++    // const b = WebGPU.Internals.jsObjects[ptr];
++    // if (b instanceof GPUBuffer) {
++    //   console.log("deleting buffer", ptr);
++    // }
+     delete WebGPU.Internals.jsObjects[ptr];
+   },
+
+@@ -811,6 +815,61 @@ var LibraryWebGPU = {
+     {{{ runtimeKeepalivePush() }}}
+     WebGPU.Internals.futureInsert(futureId, adapter.requestDevice(desc).then((device) => {
+       {{{ runtimeKeepalivePop() }}}
++
++      // a set that caches all active buffers
++      const buffers = WebGPU.Internals.buffers ??= new Set();
++      // key is buffer usage, value is total size of buffers with that usage
++      const buffersTotalSize = WebGPU.Internals.buffersTotalSize ??= new Map();
++
++      WebGPU.Internals.buffersCreated ??= 0;
++      WebGPU.Internals.buffersDestroyed ??= 0;
++      WebGPU.Internals.buffersUploads ??= 0;
++      WebGPU.Internals.buffersExternalUploads ??= 0;
++      WebGPU.Internals.buffersDownloads ??= 0;
++      WebGPU.Internals.buffersExternalDownloads ??= 0;
++
++      if ('WEBGPU_STAT' in globalThis) {
++        // create a proxy so that we can monitor buffer usages
++        device = new Proxy(device, {
++          // when call device.createBuffer(), the returned buffer should be added into buffers
++          get: (target, prop, _receiver) => {
++            if (prop === 'createBuffer') {
++              return (desc) => {
++                const buffer = target.createBuffer(desc);
++                const originalDestroy = buffer.destroy.bind(buffer);
++                buffer.destroy = () => {
++                  const previousTotal = buffersTotalSize.get(buffer.usage);
++                  buffersTotalSize.set(buffer.usage, previousTotal - buffer.size);
++                  buffers.delete(buffer);
++                  WebGPU.Internals.buffersDestroyed++;
++                  originalDestroy();
++                };
++
++                if (buffer.usage === (GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC)) {
++                  WebGPU.Internals.buffersUploads++;
++                }
++                if (buffer.usage === (GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ)) {
++                  WebGPU.Internals.buffersDownloads++;
++                }
++
++                buffers.add(buffer);
++                const previousTotal = buffersTotalSize.get(buffer.usage) ?? 0;
++                buffersTotalSize.set(buffer.usage, previousTotal + buffer.size);
++                WebGPU.Internals.buffersCreated++;
++                return buffer;
++              };
++            }
++            const propertyValue = Reflect.get(target, prop);
++            if (typeof propertyValue === 'function') {
++              return propertyValue.bind(target);
++            } else {
++              return propertyValue;
++            }
++          },
++          set: (target, prop, value, _receiver) => Reflect.set(target, prop, value),
++        });
++      }
++
+       WebGPU.Internals.jsObjectInsert(queuePtr, device.queue);
+       WebGPU.Internals.jsObjectInsert(devicePtr, device);
+
 diff --git a/third_party/emdawnwebgpu/webgpu.cpp b/third_party/emdawnwebgpu/webgpu.cpp
-index ca52b1237b..b11462fb87 100644
+index ca52b1237b..a30ca583c3 100644
 --- a/third_party/emdawnwebgpu/webgpu.cpp
 +++ b/third_party/emdawnwebgpu/webgpu.cpp
 @@ -131,7 +131,6 @@ class RefCounted : NonMovable {
@@ -62,7 +139,14 @@ index ca52b1237b..b11462fb87 100644
 
    void Destroy();
    const void* GetConstMappedRange(size_t offset, size_t size);
-@@ -1168,7 +1169,11 @@ WGPUBuffer emwgpuCreateBuffer(const EventSource* source,
+@@ -1164,11 +1165,17 @@ WGPUAdapter emwgpuCreateAdapter(const EventSource* source) {
+
+ WGPUBuffer emwgpuCreateBuffer(const EventSource* source,
+                               bool mappedAtCreation = false) {
+-  return new WGPUBufferImpl(source, mappedAtCreation);
++  auto x = new WGPUBufferImpl(source, mappedAtCreation);
++  // printf(" #C++: emwgpuCreateBuffer %p\n", x);
++  return x;
  }
 
  WGPUDevice emwgpuCreateDevice(const EventSource* source, WGPUQueue queue) {
@@ -75,7 +159,7 @@ index ca52b1237b..b11462fb87 100644
  }
 
  WGPUQueue emwgpuCreateQueue(const EventSource* source) {
-@@ -1284,6 +1289,10 @@ WGPUBufferImpl::WGPUBufferImpl(const EventSource* source, bool mappedAtCreation)
+@@ -1284,6 +1291,10 @@ WGPUBufferImpl::WGPUBufferImpl(const EventSource* source, bool mappedAtCreation)
    }
  }
 
@@ -86,11 +170,27 @@ index ca52b1237b..b11462fb87 100644
  void WGPUBufferImpl::Destroy() {
    emwgpuBufferDestroy(this);
    AbortPendingMap("Buffer was destroyed before mapping was resolved.");
-@@ -1504,6 +1513,7 @@ WGPUFuture WGPUShaderModuleImpl::GetCompilationInfo(
+@@ -1504,6 +1515,7 @@ WGPUFuture WGPUShaderModuleImpl::GetCompilationInfo(
    void wgpu##Name##Release(WGPU##Name o) {       \
      if (o->Release()) {                          \
        delete o;                                  \
 +      emwgpuDelete(o);                           \
      }                                            \
    }
  WGPU_OBJECTS(DEFINE_WGPU_DEFAULT_ADDREF_RELEASE)
+@@ -1587,6 +1599,7 @@ WGPUFuture wgpuAdapterRequestDevice(
+ // ----------------------------------------------------------------------------
+
+ void wgpuBufferDestroy(WGPUBuffer buffer) {
++  // printf(" #C++: wgpuBufferDestroy %p\n", buffer);
+   buffer->Destroy();
+ }
+
+@@ -1639,6 +1652,7 @@ void wgpuBufferUnmap(WGPUBuffer buffer) {
+ WGPUBuffer wgpuDeviceCreateBuffer(WGPUDevice device,
+                                   const WGPUBufferDescriptor* descriptor) {
+   WGPUBuffer buffer = new WGPUBufferImpl(device, descriptor->mappedAtCreation);
++  // printf(" #C++: wgpuDeviceCreateBuffer %p\n", buffer);
+   emwgpuDeviceCreateBuffer(device, descriptor, buffer);
+   return buffer;
+ }

js/web/lib/wasm/ep-webgpu/README.md

@@ -0,0 +1,11 @@
+This folder "ep-webgpu" contains required TypeScript implementation for WebGPU EP support.
+
+"ep-webgpu" here contains "ep" in the name, to distinguish it from other the WebGPU implementation in the JSEP folder:
+
+- WebGPU EP is a C++ implementation. It uses the WebGPU C/C++ API provided by Dawn in the code. Emscripten will compile
+  the C++ code to WebAssembly, and add internal JavaScript code to make it work in the browser.
+- JSEP (JavaScript Execution Provider) is a hybrid implementation. It contains both JavaScript and C++ code, including
+  the interop between them. It uses the WebGPU JavaScript API provided by the browser.
+
+For WebGPU backend, when build definition `BUILD_DEFS.USE_WEBGPU_EP` is `true`, it is considered the WebGPU EP will be
+used. Otherwise, the JSEP will be used.

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

@@ -231,6 +231,16 @@ export class WebGpuBackend {
   private queryTimeBase?: bigint;
   queryType: TimestampQuery;
 
+  buffers = new Set();
+  buffersTotalSize = new Map();
+
+  buffersCreated = 0;
+  buffersDestroyed = 0;
+  buffersUploads = 0;
+  buffersExternalUploads = 0;
+  buffersDownloads = 0;
+  buffersExternalDownloads = 0;
+
   env: Env;
   sessionStatus: SessionState = 'default';
   /**
@@ -279,7 +289,67 @@ export class WebGpuBackend {
       requireFeatureIfAvailable('subgroups-f16' as GPUFeatureName);
     }
 
-    this.device = await adapter.requestDevice(deviceDescriptor);
+    const buffers = this.buffers;
+    const buffersTotalSize = this.buffersTotalSize;
+
+    const buffersUploadsIncrement = () => {
+      this.buffersUploads++;
+    };
+    const buffersDownloadsIncrement = () => {
+      this.buffersDownloads++;
+    };
+    const buffersCreatedIncrement = () => {
+      this.buffersCreated++;
+    };
+    const buffersDestroyedIncrement = () => {
+      this.buffersDestroyed++;
+    };
+
+    if ('WEBGPU_STAT' in globalThis) {
+      this.device = new Proxy(await adapter.requestDevice(deviceDescriptor), {
+        // when call device.createBuffer(), the returned buffer should be added into buffers
+        get: (target, prop, _receiver) => {
+          if (prop === 'createBuffer') {
+            return (desc: GPUBufferDescriptor) => {
+              const buffer = target.createBuffer(desc);
+              const originalDestroy = buffer.destroy.bind(buffer);
+              buffer.destroy = () => {
+                const previousTotal = buffersTotalSize.get(buffer.usage);
+                buffersTotalSize.set(buffer.usage, previousTotal - buffer.size);
+                buffers.delete(buffer);
+                buffersDestroyedIncrement();
+                originalDestroy();
+              };
+
+              // eslint-disable-next-line no-bitwise
+              if (buffer.usage === (GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC)) {
+                buffersUploadsIncrement();
+              }
+              // eslint-disable-next-line no-bitwise
+              if (buffer.usage === (GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ)) {
+                buffersDownloadsIncrement();
+              }
+
+              buffers.add(buffer);
+              const previousTotal = buffersTotalSize.get(buffer.usage) ?? 0;
+              buffersTotalSize.set(buffer.usage, previousTotal + buffer.size);
+              buffersCreatedIncrement();
+              return buffer;
+            };
+          }
+          const propertyValue = Reflect.get(target, prop);
+          if (typeof propertyValue === 'function') {
+            return propertyValue.bind(target);
+          } else {
+            return propertyValue;
+          }
+        },
+        set: (target, prop, value, _receiver) => Reflect.set(target, prop, value),
+      });
+    } else {
+      this.device = await adapter.requestDevice(deviceDescriptor);
+    }
+
     this.deviceInfo = new DeviceInfoImpl(this.device);
     this.adapterInfo = new AdapterInfoImpl(adapter.info || (await adapter.requestAdapterInfo()));
     this.gpuDataManager = createGpuDataManager(this);
@@ -844,6 +914,7 @@ export class WebGpuBackend {
   ): () => Promise<Tensor.DataType> {
     return async () => {
       const data = await downloadGpuData(this, gpuBuffer, size);
+      this.buffersExternalDownloads++;
       return createView(data.buffer, type);
     };
   }

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

@@ -260,6 +260,8 @@ export const createSession = async (
   let modelDataOffset: number, modelDataLength: number;
   const wasm = getInstance();
 
+  wasm.webgpuStat?.('createSession_start');
+
   if (Array.isArray(modelData)) {
     // if model data is an array, it must be a 2-elements tuple containing the pointer and size of the model data
     [modelDataOffset, modelDataLength] = modelData;
@@ -327,6 +329,7 @@ export const createSession = async (
     }
 
     wasm.jsepOnCreateSession?.();
+    wasm.webgpuStat?.('createSession_end');
 
     // clear current MLContext after session creation
     if (wasm.currentContext) {
@@ -436,6 +439,7 @@ export const createSession = async (
 
 export const releaseSession = (sessionId: number): void => {
   const wasm = getInstance();
+  wasm.webgpuStat?.('releaseSession_start');
   const session = activeSessions.get(sessionId);
   if (!session) {
     throw new Error(`cannot release session. invalid session id: ${sessionId}`);
@@ -462,6 +466,8 @@ export const releaseSession = (sessionId: number): void => {
     checkLastError("Can't release session.");
   }
   activeSessions.delete(sessionId);
+
+  wasm.webgpuStat?.('releaseSession_end');
 };
 
 export const prepareInputOutputTensor = async (
@@ -633,6 +639,8 @@ export const run = async (
   const outputValuesOffset = wasm.stackAlloc(outputCount * ptrSize);
   const outputNamesOffset = wasm.stackAlloc(outputCount * ptrSize);
 
+  wasm.webgpuStat?.('run_start');
+
   try {
     [runOptionsHandle, runOptionsAllocs] = setRunOptions(options);
 
@@ -722,6 +730,7 @@ export const run = async (
     }
 
     wasm.jsepOnRunStart?.(sessionHandle);
+    //wasm.webgpuStat?.('run_beforeAPI');
 
     let errorCode: number;
     if (!BUILD_DEFS.DISABLE_JSEP && ioBindingState) {
@@ -745,6 +754,8 @@ export const run = async (
       );
     }
 
+    //wasm.webgpuStat?.('run_afterAPI');
+
     if (errorCode !== 0) {
       checkLastError('failed to call OrtRun().');
     }
@@ -926,6 +937,8 @@ export const run = async (
         false,
       ]);
     }
+    wasm.webgpuStat?.('run_end');
+
     return output;
   } finally {
     wasm.stackRestore(beforeRunStack);

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

@@ -294,6 +294,7 @@ export declare namespace WebGpu {
     webgpuUnregisterBuffer(buffer: GPUBuffer): void;
     webgpuGetBuffer(bufferHandle: number): GPUBuffer;
     webgpuCreateDownloader(gpuBuffer: GPUBuffer, size: number, sessionHandle: number): () => Promise<ArrayBuffer>;
+    webgpuStat(label?: string): void;
   }
 }