Skip to content

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

Open
kylo5aby wants to merge 1 commit into
microsoft:mainfrom
kylo5aby:webgpu/blob
Open

[web] Support Blob-backed external data for on-demand loading in JSPI builds#29477
kylo5aby wants to merge 1 commit into
microsoft:mainfrom
kylo5aby:webgpu/blob

Conversation

@kylo5aby

@kylo5aby kylo5aby commented Jul 2, 2026

Copy link
Copy Markdown

Description

This PR allows externalData to be supplied as a Blob/File in JSPI builds.
Instead of materializing the whole external data file in the JS heap and holding
it through _OrtCreateSession, the WASM external data loader reads each
initializer's byte range from the Blob on demand and releases it right after the
CPU copy / GPU upload. The load-time CPU memory peak drops from the full model
size to roughly the size of the largest single initializer.

This follows the Blob/range-backed source approach suggested by the maintainers
in #29455. ExternalDataFileDescription.data already accepts Blob in its type
union, so there is no public API change

Motivation and Context

Fixes #29455.

Measurements

Phi-4-mini-instruct ONNX (q4f16, WebGPU EP, external data 1896 MB + 492 MB),
Chrome / Windows, PTL

whole-file (current) Blob-backed (this PR)
JS heap peak during load 2476 MB (held until session created) 434 MB
Tab process peak during load 2893 MB ~845 MB
Model load time 10.7 s 3.9 s

load time benefits mainly from skipping the 2.4 GB in-heap materialization.

@kylo5aby kylo5aby marked this pull request as draft July 2, 2026 11:34
@kylo5aby

kylo5aby commented Jul 7, 2026

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree company="Intel"

@kylo5aby kylo5aby marked this pull request as ready for review July 7, 2026 05:14
@kylo5aby

kylo5aby commented Jul 7, 2026

Copy link
Copy Markdown
Author

@qjia7 please take a look.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds support for supplying ONNX external data as a Blob/File in WebAssembly JSPI builds of onnxruntime-web, enabling on-demand range reads during session initialization to reduce JS heap peak (avoiding full external-data materialization).

Changes:

  • Add an async (JSPI-only) WebAssembly external-data loader that can read requested byte ranges from a mounted Blob.
  • Extend the wasm module/type surface to allow mountExternalData() to accept Blob, and adjust session creation to mount Blob directly in JSPI builds (fallback to full materialization otherwise).
  • Add E2E coverage for WebGPU external data supplied as a Blob (both JSPI and non-JSPI bundles).

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
onnxruntime/wasm/pre.js Allow mounting Blob and clean up extra scratch state on unmount.
onnxruntime/core/framework/external_data_loader.cc Add JSPI async loader to stream Blob byte ranges for external data.
js/web/lib/wasm/wasm-core-impl.ts Mount Blob directly in JSPI builds; otherwise load+mount as Uint8Array.
js/web/lib/wasm/wasm-types.ts Update typings/docs: mountExternalData() accepts `Uint8Array
js/web/test/e2e/run-data.js Add new browser E2E entries for Blob-backed external data.
js/web/test/e2e/browser-test-webgpu-external-data-blob.js New E2E test that provides external data via Blob.
cmake/adjust_global_compile_flags.cmake Define ORT_WASM_JSPI when JSPI is enabled to gate the async loader.

Comment thread onnxruntime/core/framework/external_data_loader.cc
Comment thread onnxruntime/wasm/pre.js

@qjia7 qjia7 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR #29477 — [web] Blob-backed external data for on-demand loading (JSPI)

Re-review of head cda69ddb73.

Summary

Adds a JSPI-only EM_ASYNC_JS external-data loader that streams each
initializer's byte range from a mounted Blob/File on demand, rather than
materializing the full external-data file in the JS heap and holding it
through _OrtCreateSession. Bytes land in a reused Module.ortExtDataScratch
buffer, get copied into WASM heap (CPU) or handed to
webgpuUploadExternalBuffer (GPU), and are released after each initializer.
Non-JSPI builds and non-Blob inputs keep the existing whole-file path. A new
ORT_WASM_JSPI compile flag gates the async loader, and pre.js now cleans
up scratch/chunk buffers on unmountExternalData. E2e test runs the same
scenario against both ort.webgpu.min.js (fallback) and ort.jspi.min.js
(on-demand) bundles.

Overall a well-scoped change with a real fallback path and dual-build test
coverage. Comments below; no hard blockers.


Correctness

C1: Scratch busy flag is cleared before the upload consumes the aliased view

data = scratch.subarray(0, length) aliases Module.ortExtDataScratch. The
inner finally clears ortExtDataScratchBusy = false before the
switch (load_type) runs, so HEAPU8.set(data, ...) /
webgpuUploadExternalBuffer(dataIdOrBuffer, data) reads scratch after the
buffer has been marked reusable.

This is safe today because (a) JS is single-threaded and there is no await
between the finally and the consumers, and (b) both consumers are fully
synchronous — verified webgpuUploadExternalBuffer in post-webgpu.js:212-259
runs getMappedRange/Uint8Array.set/unmap/copyBufferToBuffer/
queue.submit without an await, and HEAPU8.set is a synchronous typed-array
method. But the ordering is fragile: any future change that inserts an await
or turns the upload path async silently breaks the aliasing.

Please either (i) move Module.ortExtDataScratchBusy = false below the switch,
or (ii) add a comment above the finally stating that the scratch view must
outlive the consumers and the consumers must be synchronous. Option (i) is
strictly safer and costs nothing.

C2: Chunk take-and-null is atomic; comment is now correct

let chunk = Module.ortExtDataChunk || new ArrayBuffer(1048576); Module.ortExtDataChunk = null;
is atomic under single-threaded JS (no yield between the two statements), and
concurrent BYOB loads degrade to each allocating a private 1 MiB chunk with
last-writer-wins on the restore — no corruption. The "one owner per load"
comment now matches the code. Please also consider skipping the 1 MiB
allocation when length === 0 (rare but wasteful for empty initializers).


Scope / Compatibility

S1: The widened double offset can't actually exceed 4 GB — caller still caps it

The async signature takes data_offset / data_length as double (good — the
old EM_ASM_INT path was int32-only). But the sole caller in
tensorprotoutils.cc:1759 (inside #if defined(__wasm__)) still guards
file_offset + raw_data_safe_len >= 4294967296 and returns before invoking the
loader. So >4 GB is unreachable in this PR. Not a defect — just don't advertise

4 GB support until that upstream guard is relaxed. A one-line comment near the
signature ("callers currently cap at 4 GB; see tensorprotoutils.cc") would
prevent readers from assuming the loader is the bottleneck.

S2: Non-JSPI Blob path still works via loadFile, but silently loses the memory win

In wasm-core-impl.ts, the Blob shortcut is gated on
BUILD_DEFS.ENABLE_JSPI && data instanceof Blob. If a caller passes a Blob to
a non-JSPI build, loadFile(blob) is invoked and materializes the whole file
via await file.arrayBuffer() (wasm-utils-load-file.ts:79-80). The API
docstring in wasm-types.ts says "A Blob is supported in JSPI builds only",
which is accurate for the memory-win path but misleading — the Blob does load
successfully in non-JSPI builds, just without the streaming benefit. Consider
rephrasing: "A Blob is streamed on demand in JSPI builds and fully
materialized in other builds."

Also note that File extends Blob, so data instanceof Blob catches File
inputs too. Presumably intentional and correct (File exposes slice/stream),
but worth calling out in the docstring.


Design / Quality

Q1: Bare catch {} collapses all streaming failures into code 4

} catch { return 4; } at the outer level hides distinct failure modes: stream
read errors, BYOB fallback allocation failures, HEAPU8.set RangeError from
memory growth, WebGPU submit errors from webgpuUploadExternalBuffer. All
surface only as "Unknown error occurred in memory copy." Since this path runs
once per initializer at load time, the cost of catch (e) { console.error('ORT ext data load failed:', e); return 4; } (or a distinct code + message) is
negligible and would make field diagnosis on specific browsers/versions
tractable.

Q2: Async JS body reimplements the sync EM_ASM_INT loader

The non-Blob branch (Uint8Array → bounds check → HEAPU8.set / upload switch)
and the error-code contract (0/1/2/3/4/5) now live in two places:
OrtLoadWebAssemblyExternalDataAsync (JSPI) and the EM_ASM_INT block below
it. EM_ASYNC_JS and EM_ASM_INT can't share a body directly, but the
load_type switch (case 0 → HEAPU8.set, case 1 → webgpuUploadExternalBuffer || jsepUploadExternalBuffer) is pure JS and could live on Module as a
helper called from both sites. Would prevent drift on the error-code table and
the eventual "unified upload interface" TODO.

Q3: CMake has two adjacent if (onnxruntime_ENABLE_WEBASSEMBLY_JSPI) blocks

adjust_global_compile_flags.cmake:48-59 now has two back-to-back guards on
the same condition — one for add_compile_definitions(ORT_WASM_JSPI) and one
for the exception flags. Merging them (keep the exception-catching comment
above) is a trivial cleanup.


Nits

  • Error-code 5 is fragmented across languages. The JS returns 5 only in
    the JSPI build; the C++ switch handles 5 only under
    #if defined(ORT_WASM_JSPI). Consistent, but the two definitions of
    "unexpected byte count" live in different files. A shared comment listing
    the canonical 0-5 mapping (either at the top of the .cc or on Module)
    would help future maintainers.
  • reader scoping in the non-BYOB fallback. In the else branch, if
    getReader() itself threw (before the BYOB try), the outer catch
    handles it — but this is subtle. Explicit if (!reader) return 4; would be
    clearer, though not strictly needed.
  • E2e test coverage. Validates output shape/values but not the memory
    claim. That's fine as a correctness gate — the peak-memory measurement is
    hard to assert in-browser; PR description carries it. Consider a lightweight
    assertion that Module.ortExtDataScratch.length <= largestInitializerSize
    after load (probes the "peak = largest initializer" claim without measuring
    the JS heap).

Verdict

Approve with minor changes. No blockers.

Pre-merge:

  • C1 — reorder the busy=false release below the switch, or document the
    synchronous-consumer contract. This is the only item with correctness
    implications if the upload path ever becomes async.

Follow-ups (nice-to-have, not gating):

  • Q1 — capture and log the exception message.
  • Q2 — extract the shared load-type/upload helper.
  • S1 — comment that the caller currently caps at 4 GB.
  • S2 — clarify the non-JSPI Blob fallback in the mountExternalData
    docstring.

Nice memory result (2476 MB → 434 MB JS heap on Phi-4-mini), and the
Blob-fallback + dual-build e2e coverage is the right call.

@qjia7 qjia7 requested review from edgchen1 and hariharans29 July 8, 2026 01:45
… builds

Signed-off-by: Zhenwei Jin <zhenwei.jin@intel.com>
@kylo5aby

kylo5aby commented Jul 8, 2026

Copy link
Copy Markdown
Author

PR #29477 — [web] Blob-backed external data for on-demand loading (JSPI)

Re-review of head cda69ddb73.

Summary

Adds a JSPI-only EM_ASYNC_JS external-data loader that streams each initializer's byte range from a mounted Blob/File on demand, rather than materializing the full external-data file in the JS heap and holding it through _OrtCreateSession. Bytes land in a reused Module.ortExtDataScratch buffer, get copied into WASM heap (CPU) or handed to webgpuUploadExternalBuffer (GPU), and are released after each initializer. Non-JSPI builds and non-Blob inputs keep the existing whole-file path. A new ORT_WASM_JSPI compile flag gates the async loader, and pre.js now cleans up scratch/chunk buffers on unmountExternalData. E2e test runs the same scenario against both ort.webgpu.min.js (fallback) and ort.jspi.min.js (on-demand) bundles.

Overall a well-scoped change with a real fallback path and dual-build test coverage. Comments below; no hard blockers.

Correctness

C1: Scratch busy flag is cleared before the upload consumes the aliased view

data = scratch.subarray(0, length) aliases Module.ortExtDataScratch. The inner finally clears ortExtDataScratchBusy = false before the switch (load_type) runs, so HEAPU8.set(data, ...) / webgpuUploadExternalBuffer(dataIdOrBuffer, data) reads scratch after the buffer has been marked reusable.

This is safe today because (a) JS is single-threaded and there is no await between the finally and the consumers, and (b) both consumers are fully synchronous — verified webgpuUploadExternalBuffer in post-webgpu.js:212-259 runs getMappedRange/Uint8Array.set/unmap/copyBufferToBuffer/ queue.submit without an await, and HEAPU8.set is a synchronous typed-array method. But the ordering is fragile: any future change that inserts an await or turns the upload path async silently breaks the aliasing.

Please either (i) move Module.ortExtDataScratchBusy = false below the switch, or (ii) add a comment above the finally stating that the scratch view must outlive the consumers and the consumers must be synchronous. Option (i) is strictly safer and costs nothing.

C2: Chunk take-and-null is atomic; comment is now correct

let chunk = Module.ortExtDataChunk || new ArrayBuffer(1048576); Module.ortExtDataChunk = null; is atomic under single-threaded JS (no yield between the two statements), and concurrent BYOB loads degrade to each allocating a private 1 MiB chunk with last-writer-wins on the restore — no corruption. The "one owner per load" comment now matches the code. Please also consider skipping the 1 MiB allocation when length === 0 (rare but wasteful for empty initializers).

Scope / Compatibility

S1: The widened double offset can't actually exceed 4 GB — caller still caps it

The async signature takes data_offset / data_length as double (good — the old EM_ASM_INT path was int32-only). But the sole caller in tensorprotoutils.cc:1759 (inside #if defined(__wasm__)) still guards file_offset + raw_data_safe_len >= 4294967296 and returns before invoking the loader. So >4 GB is unreachable in this PR. Not a defect — just don't advertise

4 GB support until that upstream guard is relaxed. A one-line comment near the
signature ("callers currently cap at 4 GB; see tensorprotoutils.cc") would
prevent readers from assuming the loader is the bottleneck.

S2: Non-JSPI Blob path still works via loadFile, but silently loses the memory win

In wasm-core-impl.ts, the Blob shortcut is gated on BUILD_DEFS.ENABLE_JSPI && data instanceof Blob. If a caller passes a Blob to a non-JSPI build, loadFile(blob) is invoked and materializes the whole file via await file.arrayBuffer() (wasm-utils-load-file.ts:79-80). The API docstring in wasm-types.ts says "A Blob is supported in JSPI builds only", which is accurate for the memory-win path but misleading — the Blob does load successfully in non-JSPI builds, just without the streaming benefit. Consider rephrasing: "A Blob is streamed on demand in JSPI builds and fully materialized in other builds."

Also note that File extends Blob, so data instanceof Blob catches File inputs too. Presumably intentional and correct (File exposes slice/stream), but worth calling out in the docstring.

Design / Quality

Q1: Bare catch {} collapses all streaming failures into code 4

} catch { return 4; } at the outer level hides distinct failure modes: stream read errors, BYOB fallback allocation failures, HEAPU8.set RangeError from memory growth, WebGPU submit errors from webgpuUploadExternalBuffer. All surface only as "Unknown error occurred in memory copy." Since this path runs once per initializer at load time, the cost of catch (e) { console.error('ORT ext data load failed:', e); return 4; } (or a distinct code + message) is negligible and would make field diagnosis on specific browsers/versions tractable.

Q2: Async JS body reimplements the sync EM_ASM_INT loader

The non-Blob branch (Uint8Array → bounds check → HEAPU8.set / upload switch) and the error-code contract (0/1/2/3/4/5) now live in two places: OrtLoadWebAssemblyExternalDataAsync (JSPI) and the EM_ASM_INT block below it. EM_ASYNC_JS and EM_ASM_INT can't share a body directly, but the load_type switch (case 0 → HEAPU8.set, case 1 → webgpuUploadExternalBuffer || jsepUploadExternalBuffer) is pure JS and could live on Module as a helper called from both sites. Would prevent drift on the error-code table and the eventual "unified upload interface" TODO.

Q3: CMake has two adjacent if (onnxruntime_ENABLE_WEBASSEMBLY_JSPI) blocks

adjust_global_compile_flags.cmake:48-59 now has two back-to-back guards on the same condition — one for add_compile_definitions(ORT_WASM_JSPI) and one for the exception flags. Merging them (keep the exception-catching comment above) is a trivial cleanup.

Nits

  • Error-code 5 is fragmented across languages. The JS returns 5 only in
    the JSPI build; the C++ switch handles 5 only under
    #if defined(ORT_WASM_JSPI). Consistent, but the two definitions of
    "unexpected byte count" live in different files. A shared comment listing
    the canonical 0-5 mapping (either at the top of the .cc or on Module)
    would help future maintainers.
  • reader scoping in the non-BYOB fallback. In the else branch, if
    getReader() itself threw (before the BYOB try), the outer catch
    handles it — but this is subtle. Explicit if (!reader) return 4; would be
    clearer, though not strictly needed.
  • E2e test coverage. Validates output shape/values but not the memory
    claim. That's fine as a correctness gate — the peak-memory measurement is
    hard to assert in-browser; PR description carries it. Consider a lightweight
    assertion that Module.ortExtDataScratch.length <= largestInitializerSize
    after load (probes the "peak = largest initializer" claim without measuring
    the JS heap).

Verdict

Approve with minor changes. No blockers.

Pre-merge:

  • C1 — reorder the busy=false release below the switch, or document the
    synchronous-consumer contract. This is the only item with correctness
    implications if the upload path ever becomes async.

Follow-ups (nice-to-have, not gating):

  • Q1 — capture and log the exception message.
  • Q2 — extract the shared load-type/upload helper.
  • S1 — comment that the caller currently caps at 4 GB.
  • S2 — clarify the non-JSPI Blob fallback in the mountExternalData
    docstring.

Nice memory result (2476 MB → 434 MB JS heap on Phi-4-mini), and the Blob-fallback + dual-build e2e coverage is the right call.

thanks for the review. addressed C1/C2/S1/S2/Q1/Q3 and the error-code comment. Left Q2 and the two nits for follow-up as discussed above.

C1

resolved. moved the scratch lock release to an outer finally that runs after copy/upload.

C2

resolved. length === 0 now short-circuits to an empty Uint8Array.

S1

add a comment above the async loader noting the double params can address beyond 4GB, but the wasm caller still rejects currently.

S2

resolved. updated the docstring.

Q1

resolved. the async loader now logs the exception with console.error before returning code 4.

Q2

agree this would reduce drift, but I'd rather keep it as a follow up, it touches pre.js plus both loader paths and would widen the PR. happy to do it in a separate change.

Q3

resolved. merged into a single block.

Nits

  • Error-code 5 is fragmented across languages.

added a shared error-code comment at the top of the wasm loader section,

  • reader scoping in the non-BYOB fallback.

getReader() either returns a reader or throws; a falsy reader without an exception isn't possible. the outer catch now logs the error.

  • E2e test coverage.

considered it, but unmountExternalData() deletes ortExtDataScratch after session creation, so there's nothing to assert post-load. The e2e model is also tiny, so the check wouldn't add much, keep the memory numbers in the PR description for now.

@qjia7

qjia7 commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

The PR looks good to me. I'd like @hariharans29 or @edgchen1 to do a final review. Thanks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[WebGPU] Optimize load-time memory peak by streaming external data instead of fully materializing it

3 participants