enable WebGPU EP in WebAssembly build

microsoft · Feb 16, 2025 · 78b8823 · 78b8823
1 parent c7aa9a7
commit 78b8823
Show file tree

Hide file tree

Showing 30 changed files with 1,783 additions and 389 deletions.
diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake
@@ -634,6 +634,13 @@ if (onnxruntime_USE_WEBGPU)
   set(DAWN_BUILD_TESTS OFF CACHE BOOL "" FORCE)
   if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
     set(DAWN_EMSCRIPTEN_TOOLCHAIN "${REPO_ROOT}/cmake/external/emsdk/upstream/emscripten" CACHE STRING "" FORCE)
+
+    # Update a few files in Emscripten
+    #
+    # The following files should be updated in Emscripten. We are waiting for the next Emscripten release to include
+    # these changes. For now, we apply the changes manually.
+    # - ${DAWN_EMSCRIPTEN_TOOLCHAIN}/src/closure-externs/webgpu-externs.js
+    execute_process(COMMAND ${CMAKE_COMMAND} -E copy "${PROJECT_SOURCE_DIR}/patches/emscripten/webgpu-externs.js" "${DAWN_EMSCRIPTEN_TOOLCHAIN}/src/closure-externs/webgpu-externs.js")
   else()
     if (onnxruntime_BUILD_DAWN_MONOLITHIC_LIBRARY)
       set(DAWN_BUILD_MONOLITHIC_LIBRARY ON CACHE BOOL "" FORCE)

diff --git a/cmake/onnxruntime_webassembly.cmake b/cmake/onnxruntime_webassembly.cmake
@@ -211,10 +211,14 @@ else()
     target_link_libraries(onnxruntime_webassembly PRIVATE tensorboard)
   endif()
 
+  set(onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/pre.js")
+
+  set(EXPORTED_FUNCTIONS "_malloc,_free")
   if (onnxruntime_USE_JSEP)
-    set(EXPORTED_FUNCTIONS "_malloc,_free,_JsepOutput,_JsepGetNodeName")
-  else()
-    set(EXPORTED_FUNCTIONS "_malloc,_free")
+    string(APPEND EXPORTED_FUNCTIONS ",_JsepOutput,_JsepGetNodeName")
+  endif()
+  if (onnxruntime_USE_WEBGPU)
+    string(APPEND EXPORTED_FUNCTIONS ",_wgpuBufferRelease,_wgpuCreateInstance")
   endif()
 
   if (onnxruntime_ENABLE_WEBASSEMBLY_MEMORY64)
@@ -312,13 +316,15 @@ else()
         target_compile_options(noexcep_operators PRIVATE ${SMEMORY_FLAG} -Wno-experimental)
     endif()
     target_link_options(onnxruntime_webassembly PRIVATE
-      --post-js "${ONNXRUNTIME_ROOT}/wasm/js_post_js_64.js"
+      "SHELL:--post-js \"${ONNXRUNTIME_ROOT}/wasm/js_post_js_64.js\""
     )
+    list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/js_post_js_64.js")
   else ()
     set(MAXIMUM_MEMORY "4294967296")
     target_link_options(onnxruntime_webassembly PRIVATE
-      --post-js "${ONNXRUNTIME_ROOT}/wasm/js_post_js.js"
+      "SHELL:--post-js \"${ONNXRUNTIME_ROOT}/wasm/js_post_js.js\""
     )
+    list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/js_post_js.js")
   endif ()
 
   target_link_options(onnxruntime_webassembly PRIVATE
@@ -372,7 +378,6 @@ jsepDownload:_pp_")
       "SHELL:-s SIGNATURE_CONVERSIONS='${SIGNATURE_CONVERSIONS}'"
     )
   endif ()
-  set_target_properties(onnxruntime_webassembly PROPERTIES LINK_DEPENDS ${ONNXRUNTIME_ROOT}/wasm/pre.js)
 
   if (onnxruntime_USE_JSEP)
     # NOTE: "-s ASYNCIFY=1" is required for JSEP to work with WebGPU
@@ -382,10 +387,8 @@ jsepDownload:_pp_")
     target_compile_definitions(onnxruntime_webassembly PRIVATE USE_JSEP=1)
     target_link_options(onnxruntime_webassembly PRIVATE
       "SHELL:--pre-js \"${ONNXRUNTIME_ROOT}/wasm/pre-jsep.js\""
-      "SHELL:-s ASYNCIFY=1"
-      "SHELL:-s ASYNCIFY_STACK_SIZE=65536"
     )
-    set_target_properties(onnxruntime_webassembly PROPERTIES LINK_DEPENDS ${ONNXRUNTIME_ROOT}/wasm/pre-jsep.js)
+    list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/pre-jsep.js")
 
     if (onnxruntime_ENABLE_WEBASSEMBLY_MEMORY64)
       target_link_options(onnxruntime_webassembly PRIVATE
@@ -397,6 +400,20 @@ jsepDownload:_pp_")
 
   if (onnxruntime_USE_WEBGPU)
     target_compile_definitions(onnxruntime_webassembly PRIVATE USE_WEBGPU=1)
+    target_link_options(onnxruntime_webassembly PRIVATE
+      "SHELL:--post-js \"${ONNXRUNTIME_ROOT}/wasm/post-webgpu.js\""
+    )
+    list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/post-webgpu.js")
+  endif()
+
+  if (onnxruntime_USE_JSEP OR onnxruntime_USE_WEBGPU OR onnxruntime_USE_WEBNN)
+    # if any of the above is enabled, we need to use the asyncify library
+    target_link_options(onnxruntime_webassembly PRIVATE
+      "SHELL:--pre-js \"${ONNXRUNTIME_ROOT}/wasm/pre-async.js\""
+      "SHELL:-s ASYNCIFY=1"
+      "SHELL:-s ASYNCIFY_STACK_SIZE=65536"
+    )
+    list(APPEND onnxruntime_webassembly_script_deps "${ONNXRUNTIME_ROOT}/wasm/pre-async.js")
   endif()
 
   if (onnxruntime_EMSCRIPTEN_SETTINGS)
@@ -458,6 +475,8 @@ jsepDownload:_pp_")
     )
   endif()
 
+  set_target_properties(onnxruntime_webassembly PROPERTIES LINK_DEPENDS "${onnxruntime_webassembly_script_deps}")
+
   set(target_name_list ort)
 
   if (onnxruntime_ENABLE_TRAINING_APIS)

diff --git a/cmake/patches/dawn/dawn.patch b/cmake/patches/dawn/dawn.patch
@@ -21,3 +21,76 @@ index 6e8ae37593..633af91eef 100644
              -q
              "${EM_BUILD_GEN_DIR}/struct_info_webgpu.json"
              "-I=${EM_BUILD_GEN_DIR}/include"
+diff --git a/src/emdawnwebgpu/README.md b/src/emdawnwebgpu/README.md
+index efd6491cd6..8ebc5d28b6 100644
+--- a/src/emdawnwebgpu/README.md
++++ b/src/emdawnwebgpu/README.md
+@@ -56,7 +56,7 @@ Set up the build directory using emcmake
+ mkdir out/cmake-wasm
+ cd out/cmake-wasm
+
+-# Make sure the path is to the source checkout of Emscripten, not emsdk's release.
++# If using Emscripten v4.0.2 or lower, make sure the path is to the source checkout of Emscripten, not emsdk's release.
+ emcmake cmake -GNinja -DDAWN_EMSCRIPTEN_TOOLCHAIN="path/to/emscripten" ../..
+
+ ninja
+diff --git a/third_party/emdawnwebgpu/webgpu.cpp b/third_party/emdawnwebgpu/webgpu.cpp
+index ca52b1237b..b11462fb87 100644
+--- a/third_party/emdawnwebgpu/webgpu.cpp
++++ b/third_party/emdawnwebgpu/webgpu.cpp
+@@ -131,7 +131,6 @@ class RefCounted : NonMovable {
+   bool Release() {
+     if (mRefCount.fetch_sub(1u, std::memory_order_release) == 1u) {
+       std::atomic_thread_fence(std::memory_order_acquire);
+-      emwgpuDelete(this);
+       return true;
+     }
+     return false;
+@@ -234,6 +233,7 @@ class Ref {
+   static void Release(T value) {
+     if (value != nullptr && value->RefCounted::Release()) {
+       delete value;
++      emwgpuDelete(value);
+     }
+   }
+
+@@ -642,6 +642,7 @@ struct WGPUBufferImpl final : public EventSource,
+                               public RefCountedWithExternalCount {
+  public:
+   WGPUBufferImpl(const EventSource* source, bool mappedAtCreation);
++  ~WGPUBufferImpl();
+
+   void Destroy();
+   const void* GetConstMappedRange(size_t offset, size_t size);
+@@ -1168,7 +1169,11 @@ WGPUBuffer emwgpuCreateBuffer(const EventSource* source,
+ }
+
+ WGPUDevice emwgpuCreateDevice(const EventSource* source, WGPUQueue queue) {
+-  return new WGPUDeviceImpl(source, queue);
++  // This function is only called from JS via `importJsDevice()`, which
++  // needs to increment the external ref count to fix the behavior.
++  WGPUDeviceImpl* device = new WGPUDeviceImpl(source, queue);
++  device->AddExternalRef();
++  return device;
+ }
+
+ WGPUQueue emwgpuCreateQueue(const EventSource* source) {
+@@ -1284,6 +1289,10 @@ WGPUBufferImpl::WGPUBufferImpl(const EventSource* source, bool mappedAtCreation)
+   }
+ }
+
++WGPUBufferImpl::~WGPUBufferImpl() {
++  Destroy();
++}
++
+ void WGPUBufferImpl::Destroy() {
+   emwgpuBufferDestroy(this);
+   AbortPendingMap("Buffer was destroyed before mapping was resolved.");
+@@ -1504,6 +1513,7 @@ WGPUFuture WGPUShaderModuleImpl::GetCompilationInfo(
+   void wgpu##Name##Release(WGPU##Name o) {       \
+     if (o->Release()) {                          \
+       delete o;                                  \
++      emwgpuDelete(o);                           \
+     }                                            \
+   }
+ WGPU_OBJECTS(DEFINE_WGPU_DEFAULT_ADDREF_RELEASE)