diff --git a/web/emcc/preload.js b/web/emcc/preload.js index ad4d056aa9..3f4547173b 100644 --- a/web/emcc/preload.js +++ b/web/emcc/preload.js @@ -22,6 +22,8 @@ */ var __wasmLib = {}; +var __wasmReadyResolve; +var __wasmReady = new Promise(function(resolve) { __wasmReadyResolve = resolve; }); function __wasmLibInstantiateWasm(imports, successCallback) { __wasmLib.imports = imports; @@ -30,11 +32,13 @@ function __wasmLibInstantiateWasm(imports, successCallback) { function __wasmLibStart(wasmInstance) { __wasmLib.successCallback(wasmInstance); + return __wasmReady; } __wasmLib.start = __wasmLibStart; var Module = { "instantiateWasm": __wasmLibInstantiateWasm, - "wasmLibraryProvider": __wasmLib + "wasmLibraryProvider": __wasmLib, + "onRuntimeInitialized": function() { __wasmReadyResolve(); } }; diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 99d4194db0..15282e04a6 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -49,6 +49,7 @@ // relax setup #include "3rdparty/tvm-ffi/src/ffi/backtrace.cc" #include "3rdparty/tvm-ffi/src/ffi/container.cc" +#include "3rdparty/tvm-ffi/src/ffi/custom_allocator.cc" #include "3rdparty/tvm-ffi/src/ffi/dtype.cc" #include "3rdparty/tvm-ffi/src/ffi/error.cc" #include "3rdparty/tvm-ffi/src/ffi/extra/env_c_api.cc" diff --git a/web/src/environment.ts b/web/src/environment.ts index 19f24bf3e9..cddd09a3a4 100644 --- a/web/src/environment.ts +++ b/web/src/environment.ts @@ -41,8 +41,8 @@ function detectLibraryProvider( // create provider so that we capture imports in the provider. return { imports: item.wasmLibraryProvider.imports, - start: (inst: WebAssembly.Instance): void => { - item.wasmLibraryProvider.start(inst); + start: (inst: WebAssembly.Instance): void | Promise => { + return item.wasmLibraryProvider.start(inst); }, }; } else if (importObject["imports"] && importObject["start"] !== undefined) { @@ -53,8 +53,8 @@ function detectLibraryProvider( imports: { "wasi_snapshot_preview1": importObject["wasiImport"], }, - start: (inst: WebAssembly.Instance): void => { - importObject["start"](inst); + start: (inst: WebAssembly.Instance): void | Promise => { + return importObject["start"](inst); } }; } else { @@ -102,9 +102,9 @@ export class Environment implements LibraryProvider { } /** Mark the start of the instance. */ - start(inst: WebAssembly.Instance): void { + start(inst: WebAssembly.Instance): void | Promise { if (this.libProvider !== undefined) { - this.libProvider.start(inst); + return this.libProvider.start(inst); } } diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 9c83a4d9ac..4efe87dc75 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -895,7 +895,8 @@ export class Instance implements Disposable { wasmModule: WebAssembly.Module, importObject: Record = {}, wasmInstance?: WebAssembly.Instance, - env?: Environment + env?: Environment, + envStarted = false ) { if (wasmInstance instanceof WebAssembly.Instance) { assert( @@ -907,7 +908,7 @@ export class Instance implements Disposable { env = new Environment(importObject); wasmInstance = new WebAssembly.Instance(wasmModule, env.imports); } - env.start(wasmInstance); + if (!envStarted) env.start(wasmInstance); this.env = env; this.lib = new FFILibrary(wasmInstance, env.imports); this.memory = this.lib.memory; @@ -1374,33 +1375,7 @@ export class Instance implements Disposable { }); } - // First download all shards to cache parallely if not yet in cache - const downloadCache = async (start: number, end: number) => { - // Download params [start, end) from `list` - for (let i = start; i < end; i++) { - const shard = list[i]; - const dataUrl = new URL(shard.dataPath, tensorCacheUrl).href; - try { - await artifactCache.addToCache(dataUrl, "arraybuffer", signal); - } catch (err) { - this.env.logger("Error: Cannot fetch " + dataUrl + " err= " + err); - throw err; - } - timeElapsed = Math.ceil((perf.now() - tstart) / 1000); - fetchedBytes += shard.nbytes; - reportCallback(++fetchedShards, /*loading=*/false); - } - } - // We launch 4 parallel for loops to limit the max concurrency to 4 download - if (!cacheOnly) { - const loopSize = Math.floor(list.length / 4); - await Promise.all([ - downloadCache(0, loopSize), - downloadCache(loopSize, 2 * loopSize), - downloadCache(2 * loopSize, 3 * loopSize), - downloadCache(3 * loopSize, list.length) - ]); - } + // TripSurf: fetch/upload one shard at a time; no whole-model RAM prefetch. // Reset for the loading phase to avoid double counting with download phase fetchedBytes = 0; @@ -1426,9 +1401,9 @@ export class Instance implements Disposable { this.empty(rec.shape, rec.dtype, this.cpu()) ) }); - const recSource = buffer.slice(rec.byteOffset, rec.byteOffset + rec.nbytes); + const recSource = new Uint8Array(buffer, rec.byteOffset, rec.nbytes); // first sync copy to cpu. - this.ctx.arrayDecodeStorage(cpu_arr, new Uint8Array(recSource), rec.format, rec.dtype); + this.ctx.arrayDecodeStorage(cpu_arr, recSource, rec.format, rec.dtype); // then async stream into GPU if needed if (device.deviceType === DeviceStrToEnum.cpu) { this.tensorCacheUpdate(rec.name, cpu_arr, false); @@ -2374,8 +2349,10 @@ export function instantiate( const env = new Environment(importObject, logger); return WebAssembly.instantiate(bufferSource, env.imports).then( - (result: WebAssembly.WebAssemblyInstantiatedSource): Instance => { - return new Instance(result.module, {}, result.instance, env); + async (result: WebAssembly.WebAssemblyInstantiatedSource): Promise => { + // Emscripten 6 initializes the runtime asynchronously after instantiation. + await env.start(result.instance); + return new Instance(result.module, {}, result.instance, env, true); } ); } diff --git a/web/src/types.ts b/web/src/types.ts index 621375a23f..2412b36598 100644 --- a/web/src/types.ts +++ b/web/src/types.ts @@ -36,7 +36,7 @@ export interface LibraryProvider { * Callback function to notify the provider the created instance. * @param inst The created instance. */ - start: (inst: WebAssembly.Instance) => void; + start: (inst: WebAssembly.Instance) => void | Promise; } /**