From aec3b45f1ce5b3d0593c84212b97664589ce606a Mon Sep 17 00:00:00 2001 From: Zhenwei Jin Date: Thu, 2 Jul 2026 19:21:36 +0800 Subject: [PATCH] [web] Support Blob-backed external data for on-demand loading in JSPI builds Signed-off-by: Zhenwei Jin --- cmake/adjust_global_compile_flags.cmake | 3 +- js/web/lib/wasm/wasm-core-impl.ts | 9 +- js/web/lib/wasm/wasm-types.ts | 6 +- .../browser-test-webgpu-external-data-blob.js | 27 +++ js/web/test/e2e/run-data.js | 2 + .../core/framework/external_data_loader.cc | 164 ++++++++++++++++++ onnxruntime/wasm/pre.js | 6 +- 7 files changed, 211 insertions(+), 6 deletions(-) create mode 100644 js/web/test/e2e/browser-test-webgpu-external-data-blob.js diff --git a/cmake/adjust_global_compile_flags.cmake b/cmake/adjust_global_compile_flags.cmake index 8e45924db4d6d..6e548acbce3a0 100644 --- a/cmake/adjust_global_compile_flags.cmake +++ b/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) diff --git a/js/web/lib/wasm/wasm-core-impl.ts b/js/web/lib/wasm/wasm-core-impl.ts index 42250e5ef9bfd..0618a0db2b711 100644 --- a/js/web/lib/wasm/wasm-core-impl.ts +++ b/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); }), ); } diff --git a/js/web/lib/wasm/wasm-types.ts b/js/web/lib/wasm/wasm-types.ts index 9e468618c47e1..c2cd1b3c4550e 100644 --- a/js/web/lib/wasm/wasm-types.ts +++ b/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. */ diff --git a/js/web/test/e2e/browser-test-webgpu-external-data-blob.js b/js/web/test/e2e/browser-test-webgpu-external-data-blob.js new file mode 100644 index 0000000000000..47c93c2262d62 --- /dev/null +++ b/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); +}); diff --git a/js/web/test/e2e/run-data.js b/js/web/test/e2e/run-data.js index 92aaecc1398a2..ad2fa693f7763 100644 --- a/js/web/test/e2e/run-data.js +++ b/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] diff --git a/onnxruntime/core/framework/external_data_loader.cc b/onnxruntime/core/framework/external_data_loader.cc index c577805e69cc4..7b4dda4825bed 100644 --- a/onnxruntime/core/framework/external_data_loader.cc +++ b/onnxruntime/core/framework/external_data_loader.cc @@ -21,12 +21,170 @@ common::Status IExternalDataLoader::LoadTensor([[maybe_unused]] const Env& env, #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 data_length, ExternalDataLoadType load_type, void* tensor_data) { +#if defined(ORT_WASM_JSPI) + auto err_code = OrtLoadWebAssemblyExternalDataAsync(data_file_path.c_str(), + static_cast(data_offset), + static_cast(static_cast(data_length)), + tensor_data, + static_cast(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 @@ common::Status LoadWebAssemblyExternalData(const Env& env, static_cast(data_length), tensor_data, static_cast(load_type)); +#endif const char* err_msg; switch (err_code) { case 0: @@ -93,6 +252,11 @@ common::Status LoadWebAssemblyExternalData(const Env& env, 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."; } diff --git a/onnxruntime/wasm/pre.js b/onnxruntime/wasm/pre.js index f8a7fa0286403..70e2f62bebc17 100644 --- a/onnxruntime/wasm/pre.js +++ b/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; }; /**