Skip to content

Commit 2ba3682

Browse files
fs-eirefs-eire
authored
[web] fix support for subgroup (microsoft#25649)
### Description Currently the subgroup feature is not working correctly for WebGPU EP on web. See microsoft#25595. Before the bug is fixed in upstream (https://issues.chromium.org/issues/435879324), use this workaround to enable subgroup support.
1 parent 7b31a10 commit 2ba3682

File tree

2 files changed

+111
-62
lines changed

2 files changed

+111
-62
lines changed

‎cmake/onnxruntime_webassembly.cmake‎

Lines changed: 16 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,9 @@ if (NOT onnxruntime_USE_VCPKG)
101101
target_compile_options(onnx PRIVATE -Wno-unused-parameter -Wno-unused-variable)
102102
endif()
103103

104+
# Include the Node.js helper for finding and validating Node.js and NPM
105+
include(node_helper.cmake)
106+
104107
if (onnxruntime_BUILD_WEBASSEMBLY_STATIC_LIB)
105108
bundle_static_library(onnxruntime_webassembly
106109
${PROTOBUF_LIB}
@@ -148,11 +151,6 @@ if (onnxruntime_BUILD_WEBASSEMBLY_STATIC_LIB)
148151
GTest::gtest
149152
)
150153

151-
find_program(NODE_EXECUTABLE node required)
152-
if (NOT NODE_EXECUTABLE)
153-
message(FATAL_ERROR "Node is required for a test")
154-
endif()
155-
156154
add_test(NAME onnxruntime_webassembly_test
157155
COMMAND ${NODE_EXECUTABLE} onnxruntime_webassembly_test.js
158156
WORKING_DIRECTORY $<TARGET_FILE_DIR:onnxruntime_webassembly_test>
@@ -343,6 +341,19 @@ else()
343341
)
344342
endif()
345343

344+
#
345+
# Apply post-processing script for the generated JavaScript file
346+
#
347+
list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/wasm_post_build.js")
348+
add_custom_command(
349+
TARGET onnxruntime_webassembly
350+
POST_BUILD
351+
# Backup file at $<TARGET_FILE_NAME:onnxruntime_webassembly>.bak
352+
COMMAND ${CMAKE_COMMAND} -E copy_if_different "$<TARGET_FILE_NAME:onnxruntime_webassembly>" "$<TARGET_FILE_NAME:onnxruntime_webassembly>.bak"
353+
COMMAND ${CMAKE_COMMAND} -E echo "Performing post-process for $<TARGET_FILE_NAME:onnxruntime_webassembly>"
354+
COMMAND ${NODE_EXECUTABLE} "${ONNXRUNTIME_ROOT}/wasm/wasm_post_build.js" "$<TARGET_FILE_NAME:onnxruntime_webassembly>"
355+
)
356+
346357
set_target_properties(onnxruntime_webassembly PROPERTIES LINK_DEPENDS "${onnxruntime_webassembly_script_deps}")
347358

348359
set(target_name_list ort)
@@ -373,61 +384,4 @@ else()
373384
endif()
374385

375386
set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME ${target_name} SUFFIX ".mjs")
376-
377-
if (onnxruntime_ENABLE_WEBASSEMBLY_THREADS)
378-
#
379-
# The following POST_BUILD script is a workaround for enabling:
380-
# - using onnxruntime-web with Multi-threading enabled when import from CDN
381-
# - using onnxruntime-web when consumed in some frameworks like Vite
382-
#
383-
# In the use case mentioned above, the file name of the script may be changed. So we need to replace the line:
384-
# `new Worker(new URL("ort-wasm-*.mjs", import.meta.url),`
385-
# with
386-
# `new Worker(new URL(import.meta.url),`
387-
#
388-
# This behavior is introduced in https://github.com/emscripten-core/emscripten/pull/22165. Since it's unlikely to be
389-
# reverted, and there is no config to disable this behavior, we have to use a post-build script to workaround it.
390-
#
391-
392-
# Generate a script to do the post-build work
393-
file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/wasm_post_build.js "
394-
const fs = require('fs');
395-
const path = require('path');
396-
397-
// node wasm_post_build.js <mjsFilePath>
398-
const mjsFilePath = process.argv[2];
399-
let contents = fs.readFileSync(mjsFilePath).toString();
400-
401-
const regex = 'new Worker\\\\(new URL\\\\(\".+?\", ?import\\\\.meta\\\\.url\\\\),';
402-
const matches = [...contents.matchAll(new RegExp(regex, 'g'))];
403-
if (matches.length !== 1) {
404-
throw new Error(
405-
`Unexpected number of matches for \"\${regex}\" in \"\${mjsFilePath}\": \${matches.length}.`,
406-
);
407-
}
408-
409-
// Replace the only occurrence.
410-
contents = contents.replace(
411-
new RegExp(regex),
412-
`new Worker(new URL(import.meta.url),`,
413-
);
414-
415-
fs.writeFileSync(mjsFilePath, contents);
416-
"
417-
)
418-
419-
find_program(NODE_EXECUTABLE node required)
420-
if (NOT NODE_EXECUTABLE)
421-
message(FATAL_ERROR "Node is required to run the post-build script")
422-
endif()
423-
424-
add_custom_command(
425-
TARGET onnxruntime_webassembly
426-
POST_BUILD
427-
# Backup file at $<TARGET_FILE_NAME:onnxruntime_webassembly>.bak
428-
COMMAND ${CMAKE_COMMAND} -E copy_if_different "$<TARGET_FILE_NAME:onnxruntime_webassembly>" "$<TARGET_FILE_NAME:onnxruntime_webassembly>.bak"
429-
COMMAND ${CMAKE_COMMAND} -E echo "Performing post-process for $<TARGET_FILE_NAME:onnxruntime_webassembly>"
430-
COMMAND ${NODE_EXECUTABLE} "${CMAKE_CURRENT_BINARY_DIR}/wasm_post_build.js" "$<TARGET_FILE_NAME:onnxruntime_webassembly>"
431-
)
432-
endif()
433387
endif()

‎onnxruntime/wasm/wasm_post_build.js‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
5+
const fs = require('node:fs');
6+
const path = require('node:path');
7+
8+
//
9+
// Post-process the mjs file.
10+
//
11+
12+
//
13+
// USAGE: node wasm_post_build.js <mjsFilePath>
14+
//
15+
16+
const mjsFilePath = process.argv[2];
17+
let contents = fs.readFileSync(mjsFilePath).toString();
18+
19+
// STEP.1 - Apply workaround for spawning workers using `import.meta.url`.
20+
21+
// This script is a workaround for enabling:
22+
// - using onnxruntime-web with Multi-threading enabled when import from CDN
23+
// - using onnxruntime-web when consumed in some frameworks like Vite
24+
//
25+
// In the use case mentioned above, the file name of the script may be changed. So we need to replace the line:
26+
// `new Worker(new URL("ort-wasm-*.mjs", import.meta.url),`
27+
// with
28+
// `new Worker(new URL(import.meta.url),`
29+
//
30+
// This behavior is introduced in https://github.com/emscripten-core/emscripten/pull/22165. Since it's unlikely to be
31+
// reverted, and there is no config to disable this behavior, we have to use a post-build script to workaround it.
32+
33+
// This step should only be applied for multithreading builds
34+
if (path.basename(mjsFilePath).includes('-threaded')) {
35+
const regex = 'new Worker\\(new URL\\(".+?", ?import\\.meta\\.url\\),';
36+
const matches = [...contents.matchAll(new RegExp(regex, 'g'))];
37+
if (matches.length !== 1) {
38+
throw new Error(
39+
`Unexpected number of matches for "${regex}" in "${mjsFilePath}": ${matches.length}.`,
40+
);
41+
}
42+
43+
// Replace the only occurrence.
44+
contents = contents.replace(
45+
new RegExp(regex),
46+
`new Worker(new URL(import.meta.url),`,
47+
);
48+
}
49+
50+
// STEP.2 - Workaround the issue referred in https://issues.chromium.org/issues/435879324
51+
52+
// Closure compiler will minimize the key of object `FeatureNameString2Enum`, turning `subgroup` into something else.
53+
54+
// This workaround is to replace the generated code as following:
55+
//
56+
// (for debug build)
57+
//
58+
// > subgroups: "17",
59+
// --- change to -->
60+
// > "subgroups": "17",
61+
//
62+
// (for release build)
63+
//
64+
// > Pe:"17",
65+
// --- change to -->
66+
// > "subgroups":"17",
67+
//
68+
69+
// This step should only be applied for WebGPU EP builds
70+
if (path.basename(mjsFilePath).includes('.async')) {
71+
const regexDebug = 'subgroups: "17"';
72+
const regexRelease = '[a-zA-Z_$][a-zA-Z0-9_$]*:"17"';
73+
74+
const matchesDebug = [...contents.matchAll(new RegExp(regexDebug, 'g'))];
75+
const matchesRelease = [...contents.matchAll(new RegExp(regexRelease, 'g'))];
76+
77+
if (matchesDebug.length === 1 && matchesRelease.length === 0) {
78+
contents = contents.replace(
79+
new RegExp(regexDebug),
80+
'"subgroups": "17"',
81+
);
82+
} else if (matchesDebug.length === 0 && matchesRelease.length === 1) {
83+
contents = contents.replace(
84+
new RegExp(regexRelease),
85+
'"subgroups":"17"',
86+
);
87+
} else {
88+
throw new Error(
89+
`Unexpected number of matches for "${regexDebug}" and "${regexRelease}" in "${mjsFilePath}": Debug=${matchesDebug.length}, Release=${matchesRelease.length}.`,
90+
);
91+
}
92+
}
93+
94+
95+
fs.writeFileSync(mjsFilePath, contents);

0 commit comments

Comments
 (0)