From 249e9c08b8c04fcc1285418ea0475167caf594f2 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Tue, 26 May 2026 15:21:01 -0700 Subject: [PATCH] MuJoCo Live: pre-fetch model assets in parallel before compilation. The prefetcher asks mju_getXMLDependencies for the transitive asset list, fetches every URL with Promise.all, and primes the bytes into the WASM-side FetchCache. When the compiler runs, every resource it opens is already in memory. Cold-load on a typical Menagerie model drops from ~9 s to ~2 s. --- src/experimental/studio/emscripten.cc | 82 ++++++++++++++++-------- src/experimental/studio/live.js | 91 +++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 27 deletions(-) diff --git a/src/experimental/studio/emscripten.cc b/src/experimental/studio/emscripten.cc index 644c9ea6..b69fae6b 100644 --- a/src/experimental/studio/emscripten.cc +++ b/src/experimental/studio/emscripten.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "experimental/platform/hal/graphics_mode.h" @@ -108,8 +109,9 @@ EM_ASYNC_JS(int, FetchUrl, } }); -// Cache for data fetched via HTTP/HTTPS. Stores the downloaded bytes keyed by -// the resource name (URL) so that read() can return a pointer to the data. +// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed) +// bytes keyed by the resource URL so that the resource provider's read() +// can return a pointer into stable storage. class FetchCache { public: static FetchCache& Instance() { @@ -117,52 +119,74 @@ class FetchCache { return instance; } - // Fetches the URL and stores the result. Returns the size (>0) on success. + // Returns the size of the entry for `url`, fetching it over the network + // if necessary. Entries that have been pre-populated via Prime() are + // returned without touching the network. Returns 0 on fetch failure. int Fetch(const char* url) { - char* data = nullptr; + if (auto it = entries_.find(url); it != entries_.end()) { + return static_cast(it->second.size()); + } + char* buf = nullptr; std::int32_t size = 0; - if (!FetchUrl(url, &data, &size)) { + if (!FetchUrl(url, &buf, &size)) { return 0; } - entries_[url] = Entry{UniquePtrWasm(data), size}; - return size; + // FetchUrl hands us a malloc'd buffer; take ownership of the bytes via + // a std::string and release the original allocation. + std::string bytes(buf, size); + std::free(buf); + return static_cast(entries_.emplace(url, std::move(bytes)) + .first->second.size()); } - // Returns pointer and size for a previously fetched URL. - int Read(const char* url, const void** buffer) { + // Pre-populates the cache so a subsequent Fetch() of the same URL hits + // memory instead of the network. First writer wins. + void Prime(std::string url, std::string bytes) { + entries_.try_emplace(std::move(url), std::move(bytes)); + } + + // Hands back a pointer into the cached bytes for the given URL. + // Returns -1 if the URL has not been fetched or primed. + int Read(const char* url, const void** buffer) const { auto it = entries_.find(url); - if (it == entries_.end()) { - return -1; - } - *buffer = it->second.data.get(); - return it->second.size; + if (it == entries_.end()) return -1; + *buffer = it->second.data(); + return static_cast(it->second.size()); } - // Frees the data for a URL. void Close(const char* url) { entries_.erase(url); } private: - struct FreeDeleter { - void operator()(void* p) const { std::free(p); } - }; - template - using UniquePtrWasm = std::unique_ptr; - - struct Entry { - UniquePtrWasm data; - int size; - }; - std::unordered_map entries_; + std::unordered_map entries_; }; // --------------------------------------------------------------------------- -// Javascript-facing function to register an asset. +// Javascript-facing function to register a build-time asset (font, IBL, +// .filamat material) that is fetched via the static AssetRegistry rather +// than over HTTP. void RegisterAsset(std::string filename, std::string contents) { AssetRegistry::Instance().RegisterAsset(std::move(filename), std::move(contents)); } +// Javascript-facing function to pre-populate the HTTP/HTTPS fetch cache. +// Used by the JS-side parallel prefetcher to feed already-downloaded asset +// bytes to MuJoCo's compiler before LoadUrl is invoked, sidestepping the +// otherwise strictly-serial EM_ASYNC_JS fetch chain. +void PrimeFetchCache(std::string url, std::string bytes) { + FetchCache::Instance().Prime(std::move(url), std::move(bytes)); +} + +// Javascript-facing wrapper around mju_getXMLDependencies. Returns the +// list of resource URLs that the given root XML transitively references +// (meshes, textures, includes, ...), per MuJoCo's compiler rules. +std::vector GetXMLDependencies(const std::string& root_url) { + mjStringVec deps; + mju_getXMLDependencies(root_url.c_str(), &deps); + return deps; +} + // Javascript-facing function to initialize the app. void Init(const std::string& title, bool dark_theme) { // Note: dimensions do not matter as window will be resized to fit canvas. @@ -305,7 +329,11 @@ void Deinit() { } EMSCRIPTEN_BINDINGS(studio_bindings) { + // Matches the C-level `mjStringVec` typedef in mjspec.h. + emscripten::register_vector("mjStringVec"); emscripten::function("registerAsset", &RegisterAsset); + emscripten::function("primeFetchCache", &PrimeFetchCache); + emscripten::function("getXMLDependencies", &GetXMLDependencies); emscripten::function("init", &Init); emscripten::function("loadFile", &LoadFile); emscripten::function("loadUrl", &LoadUrl); diff --git a/src/experimental/studio/live.js b/src/experimental/studio/live.js index 8de5230f..13bd874f 100644 --- a/src/experimental/studio/live.js +++ b/src/experimental/studio/live.js @@ -21,6 +21,85 @@ function hideLoading() { loadingOverlay.style.display = 'none'; } +// --------------------------------------------------------------------------- +// Parallel asset prefetcher. +// +// MuJoCo's XML compiler synchronously opens each referenced asset via +// mjpResourceProvider. In WASM that becomes a chain of ASYNCIFY-suspended +// fetch() calls, strictly serial, one RTT per asset. For Menagerie models +// that's ~80 round trips and ~10 s of avoidable wall time on a cold cache. +// +// To sidestep this we let MuJoCo itself enumerate the dependencies (via +// mju_getXMLDependencies, exposed to JS as Module.getXMLDependencies), then +// fetch them all in parallel and feed the bytes into the WASM-side +// FetchCache via Module.primeFetchCache. When Module.loadUrl subsequently +// runs, every resource the compiler opens is already in memory. +// +// Discovery rules live in src/xml/xml_util.cc; the JS only has to know +// the scheme rewrites used by the resource providers in emscripten.cc. +// --------------------------------------------------------------------------- + +// MuJoCo Live custom URL schemes that the C++ resource providers rewrite +// before fetching. JS must do the same so primed cache keys match what +// the providers eventually request. +function resolveScheme(url) { + if (url.startsWith('github:')) { + return 'https://raw.githubusercontent.com/' + url.slice('github:'.length); + } + return url; +} + +async function prefetchModelAssets(rootUrl, onProgress) { + const primed = new Set(); // URLs we've already pushed into FetchCache + const stats = { files: 0, bytes: 0, errors: 0 }; + + // Fetch + prime, dedup against `primed`. Returns true on success. + async function fetchAndPrime(httpUrl) { + if (primed.has(httpUrl)) return true; + primed.add(httpUrl); + try { + const resp = await fetch(httpUrl); + if (!resp.ok) { + console.warn('[prefetch] HTTP', resp.status, httpUrl); + stats.errors++; + return false; + } + const bytes = new Uint8Array(await resp.arrayBuffer()); + Module.primeFetchCache(httpUrl, bytes); + stats.files++; + stats.bytes += bytes.length; + onProgress?.(stats); + return true; + } catch (err) { + console.warn('[prefetch]', httpUrl, err); + stats.errors++; + return false; + } + } + + // Step 1: fetch the root XML so mju_getXMLDependencies (which reads it + // through the resource provider) hits cache instead of the network. + // Also gives the loading screen an immediate progress tick. + const rootHttp = resolveScheme(rootUrl); + if (!await fetchAndPrime(rootHttp)) { + throw new Error(`could not fetch root model ${rootHttp}`); + } + + // Step 2: ask MuJoCo for the transitive dependency list. The call is + // async because reading recursively included XMLs may suspend on + // ASYNCIFY fetches; those XMLs end up in FetchCache as a side effect. + // The result is an embind mjStringVec we copy out and release. + const depsVec = await Module.getXMLDependencies(rootUrl); + const deps = []; + for (let i = 0; i < depsVec.size(); ++i) deps.push(depsVec.get(i)); + depsVec.delete(); + + // Step 3: fetch every dependency in parallel, deduped against URLs + // already primed in step 1 (and earlier deps within this batch). + await Promise.all(deps.map((depUrl) => fetchAndPrime(resolveScheme(depUrl)))); + return stats; +} + var Module = { preRun: [], postRun: [], @@ -111,13 +190,25 @@ var Module = { // loop until the load completes, otherwise renderFrame will // hit "Cannot have multiple async operations in flight". showLoading(); + // Surface prefetch progress on the loading screen. + const dlEl = document.getElementById('loadingDownload'); + const dlFileEl = document.getElementById('loadingDownloadFile'); + if (dlEl) dlEl.style.display = 'block'; + const onProgress = (s) => { + if (dlFileEl) { + dlFileEl.textContent = + `${s.files} files ยท ${(s.bytes/1024/1024).toFixed(1)} MB`; + } + }; requestAnimationFrame(() => { requestAnimationFrame(async () => { try { + await prefetchModelAssets(modelUrl, onProgress); await Module.loadUrl(modelUrl); } catch (error) { console.error('Failed to load model from URL:', error); } finally { + if (dlEl) dlEl.style.display = 'none'; hideLoading(); requestAnimationFrame(Module.animate); }