From b03c9496aa081fa6c07c5b266800694c830afd60 Mon Sep 17 00:00:00 2001
From: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
Date: Mon, 29 Jul 2024 13:39:38 -0700
Subject: [PATCH] [js/web] allow load WebAssembly binary from buffer (#21534)

### Description

This PR adds a new option `ort.env.wasm.wasmBinary`, which allows user
to set to a buffer containing preload .wasm file content.

This PR should resolve the problem from latest discussion in #20876.
---
 cmake/onnxruntime_webassembly.cmake           |  2 +-
 js/common/lib/env.ts                          |  6 +++++
 js/web/lib/wasm/wasm-factory.ts               |  8 ++++++-
 .../e2e/browser-test-wasm-binary-override.js  | 22 +++++++++++++++++++
 js/web/test/e2e/run-data.js                   |  3 +++
 5 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 js/web/test/e2e/browser-test-wasm-binary-override.js

diff --git a/cmake/onnxruntime_webassembly.cmake b/cmake/onnxruntime_webassembly.cmake
index 7a49e90c00bc..0686b66876d9 100644
--- a/cmake/onnxruntime_webassembly.cmake
+++ b/cmake/onnxruntime_webassembly.cmake
@@ -225,7 +225,7 @@ else()
     "SHELL:-s EXPORT_ALL=0"
     "SHELL:-s VERBOSE=0"
     "SHELL:-s FILESYSTEM=0"
-    "SHELL:-s INCOMING_MODULE_JS_API=[preRun,locateFile,arguments,onExit,wasmMemory,buffer,instantiateWasm,mainScriptUrlOrBlob]"
+    "SHELL:-s INCOMING_MODULE_JS_API=[locateFile,instantiateWasm,wasmBinary]"
     "SHELL:-s WASM_BIGINT=1"
     ${WASM_API_EXCEPTION_CATCHING}
     --no-entry
diff --git a/js/common/lib/env.ts b/js/common/lib/env.ts
index dbb5f8118363..1a87569a115a 100644
--- a/js/common/lib/env.ts
+++ b/js/common/lib/env.ts
@@ -74,6 +74,12 @@ export declare namespace Env {
      */
     wasmPaths?: WasmPrefixOrFilePaths;
 
+    /**
+     * Set a custom buffer which contains the WebAssembly binary. If this property is set, the `wasmPaths` property will
+     * be ignored.
+     */
+    wasmBinary?: ArrayBufferLike|Uint8Array;
+
     /**
      * Set or get a boolean value indicating whether to proxy the execution of main thread to a worker thread.
      *
diff --git a/js/web/lib/wasm/wasm-factory.ts b/js/web/lib/wasm/wasm-factory.ts
index fb068ab42d04..0f5f10716a00 100644
--- a/js/web/lib/wasm/wasm-factory.ts
+++ b/js/web/lib/wasm/wasm-factory.ts
@@ -108,6 +108,7 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
   const mjsPathOverride = (mjsPathOverrideFlag as URL)?.href ?? mjsPathOverrideFlag;
   const wasmPathOverrideFlag = (wasmPaths as Env.WasmFilePaths)?.wasm;
   const wasmPathOverride = (wasmPathOverrideFlag as URL)?.href ?? wasmPathOverrideFlag;
+  const wasmBinaryOverride = flags.wasmBinary;
 
   const [objectUrl, ortWasmFactory] = (await importWasmModule(mjsPathOverride, wasmPrefixOverride, numThreads > 1));
 
@@ -135,7 +136,12 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
       numThreads,
     };
 
-    if (wasmPathOverride || wasmPrefixOverride) {
+    if (wasmBinaryOverride) {
+      /**
+       * Set a custom buffer which contains the WebAssembly binary. This will skip the wasm file fetching.
+       */
+      config.wasmBinary = wasmBinaryOverride;
+    } else if (wasmPathOverride || wasmPrefixOverride) {
       /**
        * A callback function to locate the WebAssembly file. The function should return the full path of the file.
        *
diff --git a/js/web/test/e2e/browser-test-wasm-binary-override.js b/js/web/test/e2e/browser-test-wasm-binary-override.js
new file mode 100644
index 000000000000..35d427fa3b72
--- /dev/null
+++ b/js/web/test/e2e/browser-test-wasm-binary-override.js
@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+const documentUrl = document.currentScript.src;
+
+it('Browser E2E testing - WebAssembly backend', async function() {
+  // preload .wasm file binary
+  const wasmUrl = new URL('./node_modules/onnxruntime-web/dist/ort-wasm-simd-threaded.wasm', documentUrl).href;
+  const response = await fetch(wasmUrl);
+
+  // make sure the .wasm file is loaded successfully
+  assert(response.ok);
+  assert(response.headers.get('Content-Type') === 'application/wasm');
+
+  // override wasm binary
+  const binary = await response.arrayBuffer();
+  ort.env.wasm.wasmBinary = binary;
+
+  await testFunction(ort, {executionProviders: ['wasm']});
+});
diff --git a/js/web/test/e2e/run-data.js b/js/web/test/e2e/run-data.js
index 507192f29be9..856f29eac6dd 100644
--- a/js/web/test/e2e/run-data.js
+++ b/js/web/test/e2e/run-data.js
@@ -36,6 +36,9 @@ const BROWSER_TEST_CASES = [
   [true, false, './browser-test-wasm.js', 'ort.bundle.min.mjs', ['num_threads=2', 'proxy=1']],  // 2 threads, proxy
   [true, false, './browser-test-wasm.js', 'ort.bundle.min.mjs', ['num_threads=1', 'proxy=1']],  // 1 thread, proxy
 
+  // wasm binary override:
+  [true, false, './browser-test-wasm-binary-override.js', 'ort.min.js'],
+
   // path override:
   // wasm, path override filenames for both mjs and wasm, same origin
   [true, false, './browser-test-wasm-path-override-filename.js', 'ort.min.js', ['port=9876', 'files=mjs,wasm']],