[web] Support Blob-backed external data for on-demand loading in JSPI builds

cmake/adjust_global_compile_flags.cmake

@@ -45,8 +45,9 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
     string(APPEND CMAKE_CXX_FLAGS " -msimd128")
   endif()
 
-  # Enable WebAssembly exception catching.
   if (onnxruntime_ENABLE_WEBASSEMBLY_JSPI)
+    add_compile_definitions(ORT_WASM_JSPI)
+    # Enable WebAssembly exception catching.
     string(APPEND CMAKE_C_FLAGS " -fwasm-exceptions -s WASM_LEGACY_EXCEPTIONS=0")
     string(APPEND CMAKE_CXX_FLAGS " -fwasm-exceptions -s WASM_LEGACY_EXCEPTIONS=0")
   elseif (onnxruntime_ENABLE_WEBASSEMBLY_EXCEPTION_CATCHING)

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

@@ -358,9 +358,14 @@ export const createSession = async (
       const loadingPromises = [];
       for (const file of options.externalData) {
         const path = typeof file === 'string' ? file : file.path;
+        const data = typeof file === 'string' ? file : file.data;
+        if (BUILD_DEFS.ENABLE_JSPI && data instanceof Blob) {
+          wasm.mountExternalData(path, data);
+          continue;
+        }
         loadingPromises.push(
-          loadFile(typeof file === 'string' ? file : file.data).then((data) => {
-            wasm.mountExternalData(path, data);
+          loadFile(data).then((fileData) => {
+            wasm.mountExternalData(path, fileData);
           }),
         );
       }

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

@@ -468,9 +468,11 @@ export interface OrtWasmModule
    * Mount the external data file to an internal map, which will be used during session initialization.
    *
    * @param externalDataFilePath - specify the relative path of the external data file.
-   * @param externalDataFileData - specify the content data.
+   * @param externalDataFileData - specify the content data, either as the whole file content (Uint8Array) or as a
+   * reference to the file (Blob, including File). In JSPI builds a Blob is streamed on demand during session
+   * initialization; in other builds it is fully materialized into memory first.
    */
-  mountExternalData(externalDataFilePath: string, externalDataFileData: Uint8Array): void;
+  mountExternalData(externalDataFilePath: string, externalDataFileData: Uint8Array | Blob): void;
   /**
    * Unmount all external data files from the internal map.
    */

js/web/test/e2e/browser-test-webgpu-external-data-blob.js

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+it('Browser E2E testing - WebGPU backend with external data as Blob', async function () {
+  // supplying external data as a Blob. In JSPI builds, the Blob is read on demand during session initialization
+  // instead of being materialized in memory at once; in other builds it falls back to loading the whole content.
+  const blob = await (await fetch('./model_with_orig_ext_data.bin')).blob();
+  const session = await ort.InferenceSession.create('./model_with_orig_ext_data.onnx', {
+    executionProviders: ['webgpu'],
+    externalData: [{ data: blob, path: 'model_with_orig_ext_data.bin' }],
+  });
+
+  const fetches = await session.run({ X: new ort.Tensor('float32', [1, 1], [1, 2]) });
+
+  const Y = fetches.Y;
+
+  assert(Y instanceof ort.Tensor);
+  assert(Y.dims.length === 2 && Y.dims[0] === 2 && Y.dims[1] === 3);
+  assert(Y.data[0] === 1);
+  assert(Y.data[1] === 1);
+  assert(Y.data[2] === 0);
+  assert(Y.data[3] === 0);
+  assert(Y.data[4] === 0);
+  assert(Y.data[5] === 0);
+});

js/web/test/e2e/run-data.js

@@ -67,6 +67,8 @@ const BROWSER_TEST_CASES = [
 
   [true, true, './browser-test-wasm-image-tensor-image.js', 'ort.min.js'], // pre-post-process
   [true, true, './browser-test-webgpu-external-data.js', 'ort.webgpu.min.js'], // external data
+  [true, true, './browser-test-webgpu-external-data-blob.js', 'ort.webgpu.min.js'], // external data as Blob (fallback)
+  [true, true, './browser-test-webgpu-external-data-blob.js', 'ort.jspi.min.js'], // external data as Blob (on-demand)
 ];
 
 // [bundle_path, format]

onnxruntime/core/framework/external_data_loader.cc

@@ -21,12 +21,170 @@
 
 #if defined(__wasm__)
 
+// Error codes returned by the JS loader bodies below. The async and sync
+// loaders share this contract - keep them in sync:
+//   0 - success
+//   1 - Module.MountedFiles is not available
+//   2 - file not found in preloaded files
+//   3 - out-of-bounds
+//   4 - unknown error during memory copy / GPU upload
+//   5 - Blob stream returned an unexpected number of bytes
+
+#if defined(ORT_WASM_JSPI)
+
+// Asynchronous external data loader, available in JSPI builds only.
+//
+// The offset/length parameters are doubles so the loader itself can address beyond
+// 4GB, but the current wasm caller still rejects offset + length >= 4GB before reaching here.
+
+// clang-format off
+EM_ASYNC_JS(int, OrtLoadWebAssemblyExternalDataAsync,
+            (const char* data_file_path, double data_offset, double data_length,
+             void* tensor_data, int32_t load_type), {
+
+  if (typeof Module == 'undefined' || !Module.MountedFiles) {
+    return 1;  // "Module.MountedFiles" is not available.
+  }
+  let fileName = UTF8ToString(Number(data_file_path >>> 0));
+  if (fileName.startsWith('./')) {
+    fileName = fileName.substring(2);
+  }
+  const fileData = Module.MountedFiles.get(fileName);
+  if (!fileData) {
+    return 2;  // file not found in preloaded files.
+  }
+  const offset = data_offset;
+  const length = data_length;
+  const dataIdOrBuffer = Number(tensor_data >>> 0);
+
+  if (offset < 0 || length < 0) {
+    return 3;
+  }
+
+  let ownScratchLock = false;
+  try {
+    let data;
+    if (typeof Blob !== 'undefined' && fileData instanceof Blob) {
+      if (offset + length > fileData.size) {
+        return 3;  // Out of bounds.
+      }
+      if (length === 0) {
+        data = new Uint8Array(0);
+      } else {
+        // Stream into a reused scratch buffer rather than slice().arrayBuffer(), which
+        // creates a large short-lived buffer per initializer and inflates peak memory.
+        let scratch;
+        if (!Module.ortExtDataScratchBusy) {
+          Module.ortExtDataScratchBusy = true;
+          ownScratchLock = true;
+          scratch = Module.ortExtDataScratch;
+          if (!scratch || scratch.length < length) {
+            scratch = Module.ortExtDataScratch = new Uint8Array(length);
+          }
+        } else {
+          scratch = new Uint8Array(length);
+        }
+        const stream = fileData.slice(offset, offset + length).stream();
+        // BYOB reader avoids per-chunk allocations; fall back if unavailable.
+        let reader;
+        let byob = false;
+        try {
+          reader = stream.getReader({ 'mode': 'byob' });
+          byob = true;
+        } catch (e) {
+          reader = stream.getReader();
+        }
+        let pos = 0;
+        if (byob) {
+          let chunk = Module.ortExtDataChunk || new ArrayBuffer(1048576);
+          Module.ortExtDataChunk = null;  // BYOB read detaches; one owner per load.
+          for (;;) {
+            const r = await reader.read(new Uint8Array(chunk, 0, chunk.byteLength));
+            if (r.value) {
+              if (r.value.byteLength > 0) {
+                if (pos + r.value.byteLength > length) {
+                  Module.ortExtDataChunk = r.value.buffer;
+                  return 5;  // Stream yielded more bytes than requested.
+                }
+                scratch.set(r.value, pos);
+                pos += r.value.byteLength;
+              }
+              chunk = r.value.buffer;  // reclaim the (detached) chunk buffer
+            }
+            if (r.done) break;
+          }
+          Module.ortExtDataChunk = chunk;
+        } else {
+          for (;;) {
+            const r = await reader.read();
+            if (r.done) break;
+            if (pos + r.value.byteLength > length) {
+              return 5;  // Stream yielded more bytes than requested.
+            }
+            scratch.set(r.value, pos);
+            pos += r.value.byteLength;
+          }
+        }
+        if (pos !== length) {
+          return 5;  // Reading from the Blob returned an unexpected number of bytes.
+        }
+        data = scratch.subarray(0, length);
+      }
+    } else {
+      if (offset + length > fileData.byteLength) {
+        return 3;  // Out of bounds.
+      }
+      data = fileData.subarray(offset, offset + length);
+    }
+
+    switch (load_type) {
+      case 0:
+        // Load external data to CPU memory.
+        // Copy the file data (fileData,offset,length) into WebAssembly memory
+        // (HEAPU8,buffer,length).
+        HEAPU8.set(data, dataIdOrBuffer);
+        break;
+      case 1:
+        // Load external data to GPU.
+        // TODO: use a unified interface for upload external buffer.
+        if (Module.webgpuUploadExternalBuffer) {
+          Module.webgpuUploadExternalBuffer(dataIdOrBuffer, data);
+        } else {
+          Module.jsepUploadExternalBuffer(dataIdOrBuffer, data);
+        }
+        break;
+      default:
+        return 4;  // Unknown error occurred in memory copy.
+    }
+    return 0;
+  } catch (e) {
+    console.error('Failed to load external data "' + fileName + '":', e);
+    return 4;
+  } finally {
+    // `data` may alias the shared scratch; release it only after the copy/upload
+    // above has consumed it.
+    if (ownScratchLock) {
+      Module.ortExtDataScratchBusy = false;
+    }
+  }
+});
+// clang-format on
+
+#endif
+
 common::Status LoadWebAssemblyExternalData(const Env& env,
                                            const std::filesystem::path& data_file_path,
                                            FileOffsetType data_offset,
                                            SafeInt<size_t> data_length,
                                            ExternalDataLoadType load_type,
                                            void* tensor_data) {
+#if defined(ORT_WASM_JSPI)
+  auto err_code = OrtLoadWebAssemblyExternalDataAsync(data_file_path.c_str(),
+                                                      static_cast<double>(data_offset),
+                                                      static_cast<double>(static_cast<size_t>(data_length)),
+                                                      tensor_data,
+                                                      static_cast<int32_t>(load_type));
+#else
   auto err_code = EM_ASM_INT(({
                                // If available, "Module.MountedFiles" is a Map for all preloaded files.
                                if (typeof Module == 'undefined' || !Module.MountedFiles) {
@@ -80,6 +238,7 @@
                              static_cast<int32_t>(data_length),
                              tensor_data,
                              static_cast<int32_t>(load_type));
+#endif
   const char* err_msg;
   switch (err_code) {
     case 0:
@@ -93,6 +252,11 @@
     case 3:
       err_msg = "Out of bounds.";
       break;
+#if defined(ORT_WASM_JSPI)
+    case 5:
+      err_msg = "Reading from the Blob returned an unexpected number of bytes.";
+      break;
+#endif
     default:
       err_msg = "Unknown error occurred in memory copy.";
   }

onnxruntime/wasm/pre.js

@@ -9,7 +9,7 @@
  * Mount external data files of a model to an internal map, which will be used during session initialization.
  *
  * @param {string} externalDataFilesPath
- * @param {Uint8Array} externalDataFilesData
+ * @param {Uint8Array|Blob} externalDataFilesData
  */
 Module["mountExternalData"] = (externalDataFilePath, externalDataFileData) => {
   if (externalDataFilePath.startsWith("./")) {
@@ -24,6 +24,10 @@ Module["mountExternalData"] = (externalDataFilePath, externalDataFileData) => {
  */
 Module["unmountExternalData"] = () => {
   delete Module.MountedFiles;
+  // Release the buffers used by the Blob-backed external data loader
+  delete Module.ortExtDataScratch;
+  delete Module.ortExtDataChunk;
+  delete Module.ortExtDataScratchBusy;
 };
 
 /**