[Web] ORT format model not working on WebGPU EP + Wasm Static lib

[sevagh]

Describe the issue

Hello,

I have C++ code that uses the ONNXRuntime CXX API which I compile against static onnxruntime wasm (--build_wasm_static_lib - full build commands are shown at the bottom). I link against the static lib in my CMakeLists.txt file like so:

set(ONNX_RUNTIME_WASM_LIB_WEBGPU ${CMAKE_SOURCE_DIR}/../build/build-ort-wasm/Release/libonnxruntime_webassembly.a)

When building with onnxruntime only using simd, my code works fine. When building with webgpu support through jsep
and using the js execution provider in the C++ code:

// set provider to WebGPU via "JsExecutionProvider"
session_options.AppendExecutionProvider("JS", {});

This is the javascript webworker code which imports the wasm module:

function loadWASMModule() {
    // Check if WebGPU is supported
    const webGPUSupported = isWebGPUSupported();
    let scriptToImport = "";
    //const scriptToImport = webGPUSupported ? 'demucs_onnx_webgpu.js' : 'demucs_onnx_simd.js';
    if (webGPUSupported) {
        console.log("WebGPU is supported, using WebGPU module");
        scriptToImport = "demucs_onnx_webgpu.js";
    } else {
        console.log("WebGPU is not supported, using SIMD module");
        scriptToImport = "demucs_onnx_simd.js";
    }

I get an error when using my code, saying Module.jsepAlloc is not a function. This is printed by my wasm module in the Chrome developer console:

WebGPU is supported, using WebGPU module
demucs_onnx_webgpu.js:9 Using JsExecutionProvider
demucs_onnx_webgpu.js:9 JsExecutionProvider
demucs_onnx_webgpu.js:9 Uncaught (in promise) TypeError: Module.jsepAlloc is not a function
    at 738632 (demucs_onnx_webgpu.js:9:11657)
    at runEmAsmFunction (demucs_onnx_webgpu.js:9:89035)
    at _emscripten_asm_const_ptr (demucs_onnx_webgpu.js:9:89275)
    at demucs_onnx_webgpu.wasm:0x2877e6

I believe it's because I'm using the static lib which is missing the jsepInit step we find in the file ort-wasm-simd.jsep.mjs is built with the non-static --build_wasm. I have also read this guide https://gist.github.com/fs-eire/a55b2c7e10a6864b9602c279b8b75dce

I tried to include pre-jsep.js in my link flags:

LINK_FLAGS "${COMMON_LINK_FLAGS} -s USE_WEBGPU=1 \
-s ALLOW_MEMORY_GROWTH=1 -s INITIAL_MEMORY=2048mb -s MAXIMUM_MEMORY=4gb \
--pre-js \"${ONNXRUNTIME_ROOT}/wasm/pre-jsep.js\""
)

It's probably missing other ts or js files. Any advice? Yulong Wang (@fs-eire) sorry to ping you directly but you seem to be the WebGPU/JSEP implementor

Build commands:

Simd no webgpu:

python ./vendor/onnxruntime/tools/ci_build/build.py \
    --build_dir="./build/build-ort-wasm-simd" \
    --config=MinSizeRel \
    --build_wasm_static_lib \
    --parallel \
    --minimal_build \
    --disable_ml_ops \
    --disable_rtti \
    --use_preinstalled_eigen \
    --eigen_path=$(realpath "./vendor/eigen") \
    --skip_tests \
    --skip_onnx_tests \
    --enable_wasm_simd \
    --enable_wasm_exception_throwing_override \
    --disable_exceptions \
    --include_ops_by_config="./onnx-models/required_operators_and_types.config" \
    --enable_reduced_operator_type_support
;;

Webgpu jsep:

python ./vendor/onnxruntime/tools/ci_build/build.py \
    --build_dir="./build/build-ort-wasm-webgpu" \
    --config=Release \
    --build_wasm_static_lib \
    --parallel \
    --use_preinstalled_eigen \
    --eigen_path=$(realpath "./vendor/eigen") \
    --skip_tests \
    --skip_onnx_tests \
    --enable_wasm_simd \
    --use_jsep \
    --include_ops_by_config="./onnx-models/required_operators_and_types.config" \
    --enable_reduced_operator_type_support

To reproduce

Write custom C++ code using ORT CXX API, build onnxruntime WASM static lib with jsep and webgpu support, you might get a Module.jsepAlloc is not a function

Urgency

No response

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

deee480

Execution Provider

'webgpu' (WebGPU)

[fs-eire]

JSEP does not work in ONNX Runtime WebAssembly static lib, because it depends on the inter-op between JS and C++ specifically defined for the JavaScript library.

We are currently working on a C++ implemented WebGPU EP, which is technically able to work in wasm static lib. It is still being worked on.

[sevagh]

Thanks!

[sevagh]

Is there an alternative, for example if I use webgpu without JSEP? But this involves another component called "dawn"

[fs-eire]

if you are building C++ targeting web (wasm), then your application should not depend the native implementation of dawn. instead, you should use the dawn's fork of Emscripten in order to compile the code that uses WebGPU API.

However, as I explained above, we've not get to there yet. We will work on it but since that involves quite code and build script changes it may take a while.

[elliotsayes]

Thanks for the clear explanation Yulong Wang (@fs-eire), do you happen to have a rough estimate of when this is expected to be completed?

[fs-eire]

it's about 3-4 months or so

[raphaelmenges]

Looking forward for WebGPU being available as execution runner in the Emscripten environment! May I ask how is it possible that WebGPU is already available in onnxruntime-web?

PS: My running assumption is that --use_jsep creates the execution providers for the Web, which, however, are not available in C++/Emscripten. Is that correct?

[sevagh]

Is this available now? #23364 (comment)

[fs-eire]

Is this available now? #23364 (comment)

Please track #23697

[sevagh]

Not to be a bother, but it looks like this got more complicated than originally anticipated (judging by the PRs made that are relevant to WebGPU/CXX API).

I'm wondering if I should instead focus on re-thinking my design to be able to use ORT's own javascript wrapper (with webgpu support) vs waiting for the ability to do a custom CXX WASM+WebGPU compilation.

[fs-eire]

There are a lot of changes ongoing, including upstreams (dawn/emsdk). If you can use JavaScript, it would not be a bad idea to use ORT JavaScript API because even if we replace JSEP by WebGPU EP in future the JavaScript API is expected to still work.

If you want to use WebGPU EP in WASM, it is functional working today in ORT web. However there are a few performance issues and I am working on it for a while. If you want to try I can help.

[sevagh]

Latest main branch should work for WebGPU EP or do I need a PR branch? Let me test it.