@@ -21,12 +21,170 @@ common::Status IExternalDataLoader::LoadTensor([[maybe_unused]] const Env& env,
2121
2222#if defined(__wasm__)
2323
24+ // Error codes returned by the JS loader bodies below. The async and sync
25+ // loaders share this contract - keep them in sync:
26+ // 0 - success
27+ // 1 - Module.MountedFiles is not available
28+ // 2 - file not found in preloaded files
29+ // 3 - out-of-bounds
30+ // 4 - unknown error during memory copy / GPU upload
31+ // 5 - Blob stream returned an unexpected number of bytes
32+
33+ #if defined(ORT_WASM_JSPI)
34+
35+ // Asynchronous external data loader, available in JSPI builds only.
36+ //
37+ // The offset/length parameters are doubles so the loader itself can address beyond
38+ // 4GB, but the current wasm caller still rejects offset + length >= 4GB before reaching here.
39+
40+ // clang-format off
41+ EM_ASYNC_JS (int , OrtLoadWebAssemblyExternalDataAsync,
42+ (const char * data_file_path, double data_offset, double data_length,
43+ void * tensor_data, int32_t load_type), {
44+
45+ if (typeof Module == ' undefined' MountedFiles ) {
46+ return 1 ; // "Module.MountedFiles" is not available.
47+ }
48+ let fileName = UTF8ToString (Number (data_file_path >>> 0 ));
49+ if (fileName.startsWith (' ./'
50+ fileName = fileName.substring (2 );
51+ }
52+ const fileData = Module.MountedFiles .get (fileName);
53+ if (!fileData) {
54+ return 2 ; // file not found in preloaded files.
55+ }
56+ const offset = data_offset;
57+ const length = data_length;
58+ const dataIdOrBuffer = Number (tensor_data >>> 0 );
59+
60+ if (offset < 0 || length < 0 ) {
61+ return 3 ;
62+ }
63+
64+ let ownScratchLock = false ;
65+ try {
66+ let data;
67+ if (typeof Blob !== ' undefined'
68+ if (offset + length > fileData.size ) {
69+ return 3 ; // Out of bounds.
70+ }
71+ if (length === 0 ) {
72+ data = new Uint8Array (0 );
73+ } else {
74+ // Stream into a reused scratch buffer rather than slice().arrayBuffer(), which
75+ // creates a large short-lived buffer per initializer and inflates peak memory.
76+ let scratch;
77+ if (!Module.ortExtDataScratchBusy ) {
78+ Module.ortExtDataScratchBusy = true ;
79+ ownScratchLock = true ;
80+ scratch = Module.ortExtDataScratch ;
81+ if (!scratch || scratch.length < length) {
82+ scratch = Module.ortExtDataScratch = new Uint8Array (length);
83+ }
84+ } else {
85+ scratch = new Uint8Array (length);
86+ }
87+ const stream = fileData.slice (offset, offset + length).stream ();
88+ // BYOB reader avoids per-chunk allocations; fall back if unavailable.
89+ let reader;
90+ let byob = false ;
91+ try {
92+ reader = stream.getReader ({ ' mode' ' byob'
93+ byob = true ;
94+ } catch (e) {
95+ reader = stream.getReader ();
96+ }
97+ let pos = 0 ;
98+ if (byob) {
99+ let chunk = Module.ortExtDataChunk || new ArrayBuffer (1048576 );
100+ Module.ortExtDataChunk = null; // BYOB read detaches; one owner per load.
101+ for (;;) {
102+ const r = await reader.read (new Uint8Array (chunk, 0 , chunk.byteLength ));
103+ if (r.value ) {
104+ if (r.value .byteLength > 0 ) {
105+ if (pos + r.value .byteLength > length) {
106+ Module.ortExtDataChunk = r.value .buffer ;
107+ return 5 ; // Stream yielded more bytes than requested.
108+ }
109+ scratch.set (r.value , pos);
110+ pos += r.value .byteLength ;
111+ }
112+ chunk = r.value .buffer ; // reclaim the (detached) chunk buffer
113+ }
114+ if (r.done ) break ;
115+ }
116+ Module.ortExtDataChunk = chunk;
117+ } else {
118+ for (;;) {
119+ const r = await reader.read ();
120+ if (r.done ) break ;
121+ if (pos + r.value .byteLength > length) {
122+ return 5 ; // Stream yielded more bytes than requested.
123+ }
124+ scratch.set (r.value , pos);
125+ pos += r.value .byteLength ;
126+ }
127+ }
128+ if (pos !== length) {
129+ return 5 ; // Reading from the Blob returned an unexpected number of bytes.
130+ }
131+ data = scratch.subarray (0 , length);
132+ }
133+ } else {
134+ if (offset + length > fileData.byteLength ) {
135+ return 3 ; // Out of bounds.
136+ }
137+ data = fileData.subarray (offset, offset + length);
138+ }
139+
140+ switch (load_type) {
141+ case 0 :
142+ // Load external data to CPU memory.
143+ // Copy the file data (fileData,offset,length) into WebAssembly memory
144+ // (HEAPU8,buffer,length).
145+ HEAPU8 .set (data, dataIdOrBuffer);
146+ break ;
147+ case 1 :
148+ // Load external data to GPU.
149+ // TODO: use a unified interface for upload external buffer.
150+ if (Module.webgpuUploadExternalBuffer ) {
151+ Module.webgpuUploadExternalBuffer (dataIdOrBuffer, data);
152+ } else {
153+ Module.jsepUploadExternalBuffer (dataIdOrBuffer, data);
154+ }
155+ break ;
156+ default :
157+ return 4 ; // Unknown error occurred in memory copy.
158+ }
159+ return 0 ;
160+ } catch (e) {
161+ console.error (' Failed to load external data "' ' ":'
162+ return 4 ;
163+ } finally {
164+ // `data` may alias the shared scratch; release it only after the copy/upload
165+ // above has consumed it.
166+ if (ownScratchLock) {
167+ Module.ortExtDataScratchBusy = false ;
168+ }
169+ }
170+ });
171+ // clang-format on
172+
173+ #endif
174+
24175common::Status LoadWebAssemblyExternalData (const Env& env,
25176 const std::filesystem::path& data_file_path,
26177 FileOffsetType data_offset,
27178 SafeInt<size_t > data_length,
28179 ExternalDataLoadType load_type,
29180 void * tensor_data) {
181+ #if defined(ORT_WASM_JSPI)
182+ auto err_code = OrtLoadWebAssemblyExternalDataAsync (data_file_path.c_str (),
183+ static_cast <double >(data_offset),
184+ static_cast <double >(static_cast <size_t >(data_length)),
185+ tensor_data,
186+ static_cast <int32_t >(load_type));
187+ #else
30188 auto err_code = EM_ASM_INT (({
31189 // If available, "Module.MountedFiles" is a Map for all preloaded files.
32190 if (typeof Module == ' undefined' MountedFiles ) {
@@ -80,6 +238,7 @@ common::Status LoadWebAssemblyExternalData(const Env& env,
80238 static_cast <int32_t >(data_length),
81239 tensor_data,
82240 static_cast <int32_t >(load_type));
241+ #endif
83242 const char * err_msg;
84243 switch (err_code) {
85244 case 0 :
@@ -93,6 +252,11 @@ common::Status LoadWebAssemblyExternalData(const Env& env,
93252 case 3 :
94253 err_msg = " Out of bounds."
95254 break ;
255+ #if defined(ORT_WASM_JSPI)
256+ case 5 :
257+ err_msg = " Reading from the Blob returned an unexpected number of bytes."
258+ break ;
259+ #endif
96260 default :
97261 err_msg = " Unknown error occurred in memory copy."
98262 }
0 commit comments