[Web] How to use JSEP and WebGPU in static library (missing jsepAlloc or jsepInit)

[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? @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 @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