Skip to content

Commit 2c041e0

Browse files
fs-eirefs-eire
authored
use WebGPU EP instead of JSEP in WebAssembly (microsoft#24856)
### Description This PR allows to use WebGPU EP in `onnxruntime-web` NPM package. ### Migration Plan Currently, there are 2 different EPs implemented the WebGPU backend of onnxruntime-web. They are JSEP and WebGPU EP. The migration plan is to replace the JSEP with WebGPU EP and eventually remove the JSEP. The plan contains the following stages: - STAGE 1: enable WebGPU EP on onnxruntime-web in local build. (Done) - **STAGE 2: enable WebGPU EP on onnxruntime-web in the public package. (This PR)** - STAGE 3: remove JSEP from onnxruntime-web. ### Package consumption changes - Default import (`import 'onnxruntime-web'`) and CPU only import (`import 'onnxruntime-web/wasm'`) keeps their previous behaviors. - WebGPU import (`import 'onnxruntime-web/webgpu'`) will now use WebGPU EP instead of JSEP. Previously it was the same as default import. - WebGPU import will use a different suffix for the .mjs and .wasm name (which was `.jsep`): - ort-wasm-simd-threaded<b>.asyncify</b>.mjs - ort-wasm-simd-threaded<b>.asyncify</b>.wasm - The suffix `.asyncify` is used as `.jspi` is planned for future. They are 2 different ways of emscripten's async implementation (sync C++ function calls async JS function)
1 parent dbfbebe commit 2c041e0

File tree

19 files changed

+512
-226
lines changed

19 files changed

+512
-226
lines changed

‎.github/workflows/linux-wasm-ci-build-and-test-workflow.yml‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,6 @@ jobs:
9292
${{ env.common_build_args }} \
9393
--build_dir ${{ github.workspace }}/build/wasm_inferencing_webgpu \
9494
--use_webgpu \
95-
--use_jsep \
9695
--use_webnn \
9796
--target onnxruntime_webassembly \
9897
--skip_tests
@@ -113,8 +112,8 @@ jobs:
113112
if: ${{ inputs.skip_publish != true && inputs.build_webgpu == true }}
114113
run: |
115114
mkdir -p ${{ github.workspace }}/artifacts/wasm_webgpu/
116-
cp ${{ github.workspace }}/build/wasm_inferencing_webgpu/${{ inputs.build_config }}/ort-wasm-simd-threaded.jsep.wasm ${{ github.workspace }}/artifacts/wasm_webgpu/
117-
cp ${{ github.workspace }}/build/wasm_inferencing_webgpu/${{ inputs.build_config }}/ort-wasm-simd-threaded.jsep.mjs ${{ github.workspace }}/artifacts/wasm_webgpu/
115+
cp ${{ github.workspace }}/build/wasm_inferencing_webgpu/${{ inputs.build_config }}/ort-wasm-simd-threaded.asyncify.wasm ${{ github.workspace }}/artifacts/wasm_webgpu/
116+
cp ${{ github.workspace }}/build/wasm_inferencing_webgpu/${{ inputs.build_config }}/ort-wasm-simd-threaded.asyncify.mjs ${{ github.workspace }}/artifacts/wasm_webgpu/
118117
119118
- name: Upload WASM artifacts
120119
if: ${{ inputs.skip_publish != true }}

‎.github/workflows/windows-web-ci-workflow.yml‎

Lines changed: 18 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,6 @@ on:
1616
package_name:
1717
type: string
1818
default: "NPM_packages"
19-
run_webgpu_tests:
20-
type: boolean
21-
default: true
2219

2320
jobs:
2421
build_onnxruntime_web:
@@ -86,6 +83,22 @@ jobs:
8683
run: |
8784
copy ${{ github.workspace }}\artifacts_wasm\ort-*.mjs ${{ github.workspace }}\js\web\dist\
8885
86+
- name: Download WebAssembly WebGPU artifacts
87+
uses: actions/download-artifact@v4
88+
with:
89+
name: ${{ inputs.build_config }}_wasm_webgpu
90+
path: ${{ github.workspace }}/artifacts_wasm_webgpu
91+
92+
- name: Binplace dist files (.wasm) for WebGPU
93+
shell: cmd
94+
run: |
95+
copy ${{ github.workspace }}\artifacts_wasm_webgpu\ort-*.wasm ${{ github.workspace }}\js\web\dist\
96+
97+
- name: Binplace dist files (.mjs) for WebGPU
98+
shell: cmd
99+
run: |
100+
copy ${{ github.workspace }}\artifacts_wasm_webgpu\ort-*.mjs ${{ github.workspace }}\js\web\dist\
101+
89102
- name: npm ci for /js/
90103
run: npm ci
91104
working-directory: ${{ github.workspace }}/js
@@ -115,17 +128,7 @@ jobs:
115128
run: |
116129
Get-WmiObject Win32_Process -Filter "name = 'chrome.exe'" | Format-List CommandLine
117130
118-
- name: Run ort-web tests (wasm,webgl backend)
119-
if: ${{ inputs.run_webgpu_tests != true }}
120-
shell: cmd
121-
run: |
122-
mkdir ${{ runner.temp }}\web\test\01
123-
dir ${{ runner.temp }}\web\test\01
124-
npm test -- -e=chrome -b=webgl,wasm --user-data-dir=${{ runner.temp }}\web\test\01 --chromium-flags=--enable-logging --chromium-flags=--v=1
125-
working-directory: ${{ github.workspace }}\js\web
126-
127131
- name: Run ort-web tests (ALL backends)
128-
if: ${{ inputs.run_webgpu_tests == true }}
129132
shell: cmd
130133
run: |
131134
mkdir ${{ runner.temp }}\web\test\02
@@ -134,7 +137,6 @@ jobs:
134137
working-directory: ${{ github.workspace }}\js\web
135138

136139
- name: Run ort-web tests (Suite1, webgpu, IO-binding=gpu-tensor)
137-
if: ${{ inputs.run_webgpu_tests == true }}
138140
shell: cmd
139141
run: |
140142
mkdir ${{ runner.temp }}\web\test\03
@@ -143,7 +145,6 @@ jobs:
143145
working-directory: ${{ github.workspace }}\js\web
144146

145147
- name: Run ort-web tests (Suite1, webgpu, IO-binding=gpu-location)
146-
if: ${{ inputs.run_webgpu_tests == true }}
147148
shell: cmd
148149
run: |
149150
mkdir ${{ runner.temp }}\web\test\04
@@ -169,27 +170,7 @@ jobs:
169170
working-directory: ${{ github.workspace }}\js\web
170171

171172
# WebGPU EP tests
172-
- name: Download WebAssembly WebGPU artifacts
173-
if: ${{ inputs.run_webgpu_tests == true }}
174-
uses: actions/download-artifact@v4
175-
with:
176-
name: ${{ inputs.build_config }}_wasm_webgpu
177-
path: ${{ github.workspace }}/artifacts_wasm_webgpu
178-
179-
- name: Binplace dist files (.wasm) for WebGPU
180-
if: ${{ inputs.run_webgpu_tests == true }}
181-
shell: cmd
182-
run: |
183-
copy /Y ${{ github.workspace }}\artifacts_wasm_webgpu\ort-*.wasm ${{ github.workspace }}\js\web\dist\
184-
185-
- name: Binplace dist files (.mjs) for WebGPU
186-
if: ${{ inputs.run_webgpu_tests == true }}
187-
shell: cmd
188-
run: |
189-
copy /Y ${{ github.workspace }}\artifacts_wasm_webgpu\ort-*.mjs ${{ github.workspace }}\js\web\dist\
190-
191173
- name: Run ort-web tests - WebGPU EP
192-
if: ${{ inputs.run_webgpu_tests == true }}
193174
continue-on-error: true
194175
shell: cmd
195176
run: |
@@ -199,15 +180,15 @@ jobs:
199180
working-directory: ${{ github.workspace }}\js\web
200181

201182
- name: Validate shader keys - WebGPU EP
202-
if: ${{ inputs.run_webgpu_tests == true && inputs.build_config == 'Debug' }}
183+
if: ${{ inputs.build_config == 'Debug' }}
203184
uses: ./.github/actions/webgpu-validate-shader-key
204185
with:
205186
log_file_path: ${{ runner.temp }}\web\test\07\chrome_debug.log
206187
is_chromium_log: true
207188

208189
# this step is added to help investigate the shader validation failure which is hard to reproduce
209190
- name: Upload WebGPU shader validation log on failure
210-
if: ${{ failure() && inputs.run_webgpu_tests == true && inputs.build_config == 'Debug' }}
191+
if: ${{ failure() && inputs.build_config == 'Debug' }}
211192
uses: actions/upload-artifact@v4
212193
with:
213194
name: webgpu-shader-validation-logs

‎cmake/onnxruntime_webassembly.cmake‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,10 @@ function(bundle_static_library bundled_target_name)
8484
add_dependencies(${bundled_target_name} bundling_target)
8585
endfunction()
8686

87+
if (onnxruntime_USE_JSEP AND onnxruntime_USE_WEBGPU)
88+
message(FATAL_ERROR "onnxruntime_USE_JSEP and onnxruntime_USE_WEBGPU cannot be enabled at the same time.")
89+
endif()
90+
8791
if (NOT onnxruntime_ENABLE_WEBASSEMBLY_THREADS)
8892
add_compile_definitions(
8993
BUILD_MLAS_NO_ONNXRUNTIME
@@ -406,6 +410,16 @@ jsepDownload:_pp_")
406410
list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/post-webgpu.js")
407411
endif()
408412

413+
if (onnxruntime_USE_WEBNN)
414+
target_compile_definitions(onnxruntime_webassembly PRIVATE USE_WEBNN=1)
415+
if (NOT onnxruntime_USE_JSEP)
416+
target_link_options(onnxruntime_webassembly PRIVATE
417+
"SHELL:--post-js \"${ONNXRUNTIME_ROOT}/wasm/post-webnn.js\""
418+
)
419+
list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/post-webnn.js")
420+
endif()
421+
endif()
422+
409423
if (onnxruntime_USE_JSEP OR onnxruntime_USE_WEBGPU OR onnxruntime_USE_WEBNN)
410424
# if any of the above is enabled, we need to use the asyncify library
411425
target_link_options(onnxruntime_webassembly PRIVATE
@@ -499,6 +513,9 @@ jsepDownload:_pp_")
499513

500514
if (onnxruntime_USE_JSEP)
501515
string(APPEND target_name ".jsep")
516+
elseif (onnxruntime_USE_WEBGPU OR onnxruntime_USE_WEBNN)
517+
string(APPEND target_name ".asyncify")
518+
# TODO: support JSPI and add ".jspi" once JSPI build is supported
502519
endif()
503520

504521
set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME ${target_name} SUFFIX ".mjs")

‎js/build_webgpu.bat‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,11 @@ popd
6969
set PATH=C:\Program Files\Git\usr\bin;%PATH%
7070

7171
call %ROOT%build.bat --config %CONFIG% %CONFIG_EXTRA_FLAG% --skip_submodule_sync --build_wasm --target onnxruntime_webassembly --skip_tests^
72-
--enable_wasm_simd --enable_wasm_threads --use_jsep --use_webnn --use_webgpu --build_dir %BUILD_DIR%
72+
--enable_wasm_simd --enable_wasm_threads --use_webnn --use_webgpu --build_dir %BUILD_DIR%
7373

7474
IF NOT "%ERRORLEVEL%" == "0" (
7575
exit /b %ERRORLEVEL%
7676
)
7777

78-
copy /Y %BUILD_DIR%\%CONFIG%\ort-wasm-simd-threaded.jsep.wasm %ROOT%js\web\dist\
79-
copy /Y %BUILD_DIR%\%CONFIG%\ort-wasm-simd-threaded.jsep.mjs %ROOT%js\web\dist\
78+
copy /Y %BUILD_DIR%\%CONFIG%\ort-wasm-simd-threaded.asyncify.wasm %ROOT%js\web\dist\
79+
copy /Y %BUILD_DIR%\%CONFIG%\ort-wasm-simd-threaded.asyncify.mjs %ROOT%js\web\dist\

‎js/build_webgpu.sh‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
#!/bin/bash
2+
# Exit immediately if a command exits with a non-zero status.
3+
set -e
4+
5+
# build_webgpu.sh --- build onnxruntime-web with WebGPU EP
6+
#
7+
# Usage:
8+
# build_webgpu.sh config [clean]
9+
#
10+
# Options:
11+
# config Build configuration, "d" (Debug) or "r" (Release)
12+
# clean Perform a clean build (optional)
13+
14+
# Determine the root directory of the project (one level up from the script's directory)
15+
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
16+
BUILD_DIR="$ROOT_DIR/build_webgpu"
17+
18+
CONFIG=""
19+
CONFIG_EXTRA_FLAG="" # This will be empty by default
20+
21+
# Parse config argument
22+
if [ "$1" = "d" ]; then
23+
CONFIG="Debug"
24+
CONFIG_EXTRA_FLAG="--enable_wasm_profiling --wasm_run_tests_in_browser --cmake_extra_defines onnxruntime_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL=1 --enable_wasm_debug_info"
25+
elif [ "$1" = "r" ]; then
26+
CONFIG="Release"
27+
CONFIG_EXTRA_FLAG="--enable_wasm_api_exception_catching --disable_rtti"
28+
else
29+
echo "Error: Invalid configuration \"$1\"."
30+
echo "Configuration must be 'd' (Debug) or 'r' (Release)."
31+
echo "Usage: $0 [d|r] [clean]"
32+
exit 1
33+
fi
34+
35+
CLEAN_BUILD_REQUESTED=false
36+
if [ "$2" = "clean" ]; then
37+
CLEAN_BUILD_REQUESTED=true
38+
fi
39+
40+
# Perform clean if requested
41+
if [ "$CLEAN_BUILD_REQUESTED" = true ]; then
42+
echo "--- Performing clean build ---"
43+
if [ -d "$BUILD_DIR" ]; then
44+
echo "Removing build directory: $BUILD_DIR"
45+
rm -rf "$BUILD_DIR"
46+
fi
47+
48+
echo "Synchronizing and updating submodules..."
49+
pushd "$ROOT_DIR" > /dev/null
50+
git submodule sync --recursive
51+
git submodule update --init --recursive
52+
popd > /dev/null
53+
fi
54+
55+
# Determine if npm ci needs to be run
56+
# It needs to run if:
57+
# 1. A clean build was requested (which implies js/web/dist will be missing or stale)
58+
# 2. The js/web/dist directory does not exist (e.g., first build or manually removed)
59+
PERFORM_NPM_CI=false
60+
if [ "$CLEAN_BUILD_REQUESTED" = true ]; then
61+
PERFORM_NPM_CI=true
62+
elif [ ! -d "$ROOT_DIR/js/web/dist" ]; then
63+
echo "Directory $ROOT_DIR/js/web/dist not found."
64+
PERFORM_NPM_CI=true
65+
fi
66+
67+
if [ "$PERFORM_NPM_CI" = true ]; then
68+
echo "--- Running npm ci and pulling WASM artifacts ---"
69+
echo "Running npm ci in $ROOT_DIR/js"
70+
pushd "$ROOT_DIR/js" > /dev/null
71+
npm ci
72+
popd > /dev/null
73+
74+
echo "Running npm ci in $ROOT_DIR/js/common"
75+
pushd "$ROOT_DIR/js/common" > /dev/null
76+
npm ci
77+
popd > /dev/null
78+
79+
echo "Running npm ci and pull:wasm in $ROOT_DIR/js/web"
80+
pushd "$ROOT_DIR/js/web" > /dev/null
81+
npm ci
82+
npm run pull:wasm
83+
popd > /dev/null
84+
fi
85+
86+
echo "--- Building WebAssembly modules ---"
87+
88+
echo "Calling $ROOT_DIR/build.sh to build WebAssembly..."
89+
# Note: If $CONFIG_EXTRA_FLAG is empty, it will be omitted from the command due to shell expansion.
90+
"$ROOT_DIR/build.sh" \
91+
--config "$CONFIG" \
92+
--parallel \
93+
${CONFIG_EXTRA_FLAG} \
94+
--skip_submodule_sync \
95+
--build_wasm \
96+
--target onnxruntime_webassembly \
97+
--skip_tests \
98+
--enable_wasm_simd \
99+
--enable_wasm_threads \
100+
--use_webnn \
101+
--use_webgpu \
102+
--build_dir "$BUILD_DIR"
103+
104+
# The 'set -e' command at the beginning of the script ensures that the script will exit
105+
# immediately if the build.sh command (or any other command) fails.
106+
107+
echo "--- Copying build artifacts ---"
108+
# Ensure the dist directory exists before copying files
109+
mkdir -p "$ROOT_DIR/js/web/dist"
110+
111+
echo "Copying ort-wasm-simd-threaded.asyncify.wasm to $ROOT_DIR/js/web/dist/"
112+
cp -f "$BUILD_DIR/$CONFIG/ort-wasm-simd-threaded.asyncify.wasm" "$ROOT_DIR/js/web/dist/"
113+
114+
echo "Copying ort-wasm-simd-threaded.asyncify.mjs to $ROOT_DIR/js/web/dist/"
115+
cp -f "$BUILD_DIR/$CONFIG/ort-wasm-simd-threaded.asyncify.mjs" "$ROOT_DIR/js/web/dist/"
116+
117+
echo "--- WebGPU build process completed successfully ---"

‎js/common/lib/env.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ export declare namespace Env {
1515
* If not modified, the filename of the .wasm file is:
1616
* - `ort-wasm-simd-threaded.wasm` for default build
1717
* - `ort-wasm-simd-threaded.jsep.wasm` for JSEP build (with WebGPU and WebNN)
18+
* - `ort-wasm-simd-threaded.asyncify.wasm` for WebGPU build with Asyncify (with WebNN)
1819
*/
1920
wasm?: URL | string;
2021
/**
@@ -25,6 +26,7 @@ export declare namespace Env {
2526
* If not modified, the filename of the .mjs file is:
2627
* - `ort-wasm-simd-threaded.mjs` for default build
2728
* - `ort-wasm-simd-threaded.jsep.mjs` for JSEP build (with WebGPU and WebNN)
29+
* - `ort-wasm-simd-threaded.asyncify.mjs` for WebGPU build with Asyncify (with WebNN)
2830
*/
2931
mjs?: URL | string;
3032
}

‎js/web/lib/build-def.d.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,21 @@ interface BuildDefinitions {
1717
*/
1818
readonly DISABLE_WEBGL: boolean;
1919
/**
20-
* defines whether to disable the whole WebGpu/WebNN backend in the build.
20+
* defines whether to disable the JSEP support in the build.
2121
*/
2222
readonly DISABLE_JSEP: boolean;
23+
/**
24+
* defines whether to disable the WebGPU EP support in the build.
25+
*/
26+
readonly DISABLE_WEBGPU: boolean;
27+
/**
28+
* defines whether to disable the WebNN EP support in the build.
29+
*/
30+
readonly DISABLE_WEBNN: boolean;
2331
/**
2432
* defines whether to disable the whole WebAssembly backend in the build.
33+
*
34+
* When this build flag is set to `true`, only WebGL backend will be available.
2535
*/
2636
readonly DISABLE_WASM: boolean;
2737
/**
@@ -35,18 +45,12 @@ interface BuildDefinitions {
3545
* It is usually one of the following files:
3646
* - `ort-wasm-simd-threaded.mjs`
3747
* - `ort-wasm-simd-threaded.jsep.mjs`
48+
* - `ort-wasm-simd-threaded.asyncify.mjs`
3849
*
3950
* The value is valid only when it's an ESM build.
4051
*/
4152
readonly ENABLE_BUNDLE_WASM_JS: boolean;
4253

43-
/**
44-
* defines whether to use WebGPU EP instead of JSEP for WebGPU backend.
45-
*
46-
* This flag requires the corresponding WebAssembly artifact to be built with `--use_webgpu` flag.
47-
*/
48-
readonly USE_WEBGPU_EP: boolean;
49-
5054
// #endregion
5155

5256
// #region Build definitions for ESM

‎js/web/lib/index.ts‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,26 @@ if (!BUILD_DEFS.DISABLE_WEBGL) {
1919
registerBackend('webgl', onnxjsBackend, -10);
2020
}
2121

22+
if (!BUILD_DEFS.DISABLE_JSEP && !BUILD_DEFS.DISABLE_WEBGPU) {
23+
throw new Error(
24+
'The current build is specified to enable both JSEP and WebGPU EP. This is not a valid configuration. ' +
25+
'JSEP and WebGPU EPs cannot be enabled at the same time.',
26+
);
27+
}
28+
29+
if (!BUILD_DEFS.DISABLE_WEBNN && BUILD_DEFS.DISABLE_JSEP && BUILD_DEFS.DISABLE_WEBGPU) {
30+
throw new Error(
31+
'The current build is specified to enable WebNN EP without JSEP or WebGPU EP. This is not a valid configuration. ' +
32+
'WebNN EP requires either JSEP or WebGPU EP to be enabled.',
33+
);
34+
}
35+
2236
if (!BUILD_DEFS.DISABLE_WASM) {
2337
const wasmBackend = require('./backend-wasm').wasmBackend;
24-
if (!BUILD_DEFS.DISABLE_JSEP) {
38+
if (!BUILD_DEFS.DISABLE_JSEP || !BUILD_DEFS.DISABLE_WEBGPU) {
2539
registerBackend('webgpu', wasmBackend, 5);
40+
}
41+
if (!BUILD_DEFS.DISABLE_WEBNN) {
2642
registerBackend('webnn', wasmBackend, 5);
2743
}
2844
registerBackend('cpu', wasmBackend, 10);

0 commit comments

Comments
 (0)