Skip to content

Commit ba2999c

Browse files
jing-baofs-eire
andauthored
[js/web] Add Wasm Relaxed SIMD support to wasm backend (#22794)
### Description <!-- Describe your changes. --> Add Wasm Relaxed SIMD support. Use integer dot product instructions for QGemmU8X8. 1. Build with --enable_wasm_relaxed_simd 2. Use env.wasm.relaxedSimd to run it ### Motivation and Context #22533 --------- Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
1 parent 528f29a commit ba2999c

File tree

11 files changed

+611
-5
lines changed

11 files changed

+611
-5
lines changed

cmake/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ option(onnxruntime_ENABLE_WEBASSEMBLY_DEBUG_INFO "Enable this option to turn on
193193
option(onnxruntime_ENABLE_WEBASSEMBLY_PROFILING "Enable this option to turn on WebAssembly profiling and preserve function names" OFF)
194194
option(onnxruntime_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL "Enable this option to allow WebAssembly to output optimized model" OFF)
195195
option(onnxruntime_ENABLE_WEBASSEMBLY_MEMORY64 "Enable this option to allow WebAssembly to use 64bit memory" OFF)
196+
option(onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD "Enable WebAssembly Relaxed SIMD" OFF)
196197

197198
# Enable bitcode for iOS
198199
option(onnxruntime_ENABLE_BITCODE "Enable bitcode for iOS only" OFF)

cmake/adjust_global_compile_flags.cmake

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,10 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
3535
set(CMAKE_CXX_FLAGS_DEBUG "-g2")
3636
endif()
3737

38-
if (onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
38+
if (onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD)
39+
string(APPEND CMAKE_C_FLAGS " -msimd128 -mrelaxed-simd")
40+
string(APPEND CMAKE_CXX_FLAGS " -msimd128 -mrelaxed-simd")
41+
elseif (onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
3942
string(APPEND CMAKE_C_FLAGS " -msimd128")
4043
string(APPEND CMAKE_CXX_FLAGS " -msimd128")
4144
endif()

cmake/external/xnnpack.cmake

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,12 @@ if(CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
143143
list(APPEND wasm_srcs ${XNNPACK_DIR}/src/amalgam/gen/scalar.c)
144144
list(APPEND wasm_srcs ${XNNPACK_DIR}/src/amalgam/gen/wasm.c)
145145

146-
if(onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
146+
if(onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD)
147+
list(APPEND wasm_srcs ${XNNPACK_DIR}/src/amalgam/gen/wasmsimd.c)
148+
list(APPEND wasm_srcs ${XNNPACK_DIR}/src/amalgam/gen/wasmrelaxedsimd.c)
149+
target_compile_options(XNNPACK PRIVATE "-msimd128")
150+
target_compile_options(XNNPACK PRIVATE "-mrelaxed-simd")
151+
elseif(onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
147152
list(APPEND wasm_srcs ${XNNPACK_DIR}/src/amalgam/gen/wasmsimd.c)
148153
target_compile_options(XNNPACK PRIVATE "-msimd128")
149154
endif()

cmake/onnxruntime_mlas.cmake

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,12 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
287287
${mlas_platform_srcs}
288288
${MLAS_SRC_DIR}/qgemm_kernel_wasmsimd.cpp
289289
)
290+
if (onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD)
291+
set(mlas_platform_srcs
292+
${mlas_platform_srcs}
293+
${MLAS_SRC_DIR}/qgemm_kernel_wasmrelaxedsimd.cpp
294+
)
295+
endif()
290296
else()
291297
file(GLOB_RECURSE mlas_platform_srcs
292298
"${MLAS_SRC_DIR}/scalar/*.cpp"

cmake/onnxruntime_unittests.cmake

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,8 +222,11 @@ function(AddTest)
222222
else()
223223
set(TEST_NODE_FLAGS)
224224

225+
if (onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD)
226+
message(WARNING "Use system `node` to test Wasm relaxed SIMD. Please make sure to install node v21 or newer.")
227+
set(NODE_EXECUTABLE node)
225228
# prefer Node from emsdk so the version is more deterministic
226-
if (DEFINED ENV{EMSDK_NODE})
229+
elseif (DEFINED ENV{EMSDK_NODE})
227230
set(NODE_EXECUTABLE $ENV{EMSDK_NODE})
228231
else()
229232
message(WARNING "EMSDK_NODE environment variable was not set. Falling back to system `node`.")

cmake/onnxruntime_webassembly.cmake

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -485,7 +485,9 @@ jsepDownload:_pp_")
485485

486486
list(APPEND target_name_list "wasm")
487487

488-
if (onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
488+
if (onnxruntime_ENABLE_WEBASSEMBLY_RELAXED_SIMD)
489+
list(APPEND target_name_list "relaxedsimd")
490+
elseif (onnxruntime_ENABLE_WEBASSEMBLY_SIMD)
489491
list(APPEND target_name_list "simd")
490492
endif()
491493

onnxruntime/core/mlas/inc/mlas.h

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,10 @@ Module Name:
6363
#endif
6464
#if defined(__wasm__)
6565
#define MLAS_TARGET_WASM
66-
#if defined(__wasm_simd128__)
66+
#if defined(__wasm_relaxed_simd__)
67+
#define MLAS_TARGET_WASM_RELAXED_SIMD
68+
#define MLAS_TARGET_WASM_SIMD
69+
#elif defined(__wasm_simd128__)
6770
#define MLAS_TARGET_WASM_SIMD
6871
#else
6972
#define MLAS_TARGET_WASM_SCALAR

onnxruntime/core/mlas/lib/mlasi.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -996,9 +996,14 @@ extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchSdot;
996996
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchUmmla;
997997
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchSmmla;
998998
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchWasmSimd;
999+
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchWasmRelaxedSimd;
9991000
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmQuantDispatchDefault;
10001001
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemm8X8DispatchPOWER10;
10011002

1003+
#if defined(MLAS_TARGET_WASM_RELAXED_SIMD)
1004+
extern bool HasUSDot();
1005+
#endif
1006+
10021007
//
10031008
// Symmetric quantized qgemm dispatch structure
10041009
//

onnxruntime/core/mlas/lib/qgemm.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -886,6 +886,14 @@ MlasGemmQuantGetDispatch(
886886
if(BIsSigned || !AIsSigned) {
887887
GemmQuantDispatch = &MlasGemmU8X8DispatchNeon;
888888
}
889+
#elif defined(MLAS_TARGET_WASM_RELAXED_SIMD)
890+
if (!AIsSigned) {
891+
if (HasUSDot()) {
892+
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmRelaxedSimd;
893+
} else {
894+
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmSimd;
895+
}
896+
}
889897
#elif defined(MLAS_TARGET_WASM_SIMD)
890898
if (!AIsSigned) {
891899
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmSimd;

0 commit comments

Comments
 (0)