[WebGPU] allow WGSL template generation (#25130)

### Description

This PR introduces the
[WGSL-Template](https://github.com/fs-eire/wgsl-template) into WebGPU EP
to enable the capability of template based WGSL source generation.

The purpose of this change is to improve the readability of the WGSL
source code and reduce the future maintenance cost due to the difficulty
of understanding the current WGSL code generation logic.

### Explanation

- build
  - for static code generator: `--use_webgpu`
- for dynamic code generator: `--use_webgpu --use_wgsl_template dynamic`

- the template files use extension `.wgsl.template` and are located in
`core/providers/webgpu/` folder.

- the corresponding C++ file should use
`shader_helper.ApplyTemplate<...>()` to apply the template. (check
pad.cc for example)

- for documentation - https://github.com/fs-eire/wgsl-template

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: fs-eire

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

@@ -45,11 +45,28 @@ jobs:
         with:
           submodules: recursive
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
+      # The image used in this workflow has an old version of Node.js installed at /bin/node and npm at /bin/npm.
+      # We need to remove the system Node.js and npm, otherwise CMake will not be able to find the correct Node.js and npm executables.
+      - name: Remove the system Node.js
+        run: |
+          if [ -f /bin/node ]; then
+            rm /bin/node
+          fi
+          if [ -f /bin/npm ]; then
+            rm /bin/npm
+          fi
+
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: "3.12"
           architecture: ${{ env.buildArch }}
+
       - uses: microsoft/onnxruntime-github-actions/[email protected]
         with:
           vcpkg-version: '2025.06.13'

cmake/CMakeLists.txt

@@ -134,6 +134,7 @@ option(onnxruntime_USE_CANN "Build with CANN support" OFF)
 option(onnxruntime_USE_XNNPACK "Build with XNNPACK support. Provides an alternative math library on ARM, WebAssembly and x86." OFF)
 option(onnxruntime_USE_WEBNN "Build with WebNN support. Enable hardware acceleration in web browsers." OFF)
 option(onnxruntime_USE_WEBGPU "Build with WebGPU support. Enable WebGPU via C/C++ interface." OFF)
+option(onnxruntime_WGSL_TEMPLATE "Specify the code generator for WGSL template. Default is static." "static")
 option(onnxruntime_USE_EXTERNAL_DAWN "Build with treating Dawn as external dependency. Will not link Dawn at build time." OFF)
 option(onnxruntime_CUSTOM_DAWN_SRC_PATH "Path to custom Dawn src dir.")
 option(onnxruntime_BUILD_DAWN_MONOLITHIC_LIBRARY "Build Dawn as a monolithic library" OFF)

cmake/deps.txt

@@ -57,3 +57,4 @@ directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v
 cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.12.0.zip;7e733cfdc410d777b76122d64232499205589a96
 dawn;https://github.com/google/dawn/archive/9733be39e18186961d503e064874afe3e9ceb8d1.zip;2a4017c32892b90d072a9102eba90ae691fae36d
 kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.4.0.tar.gz;22d3b57b54a61c194ab256ff11b0353a3b220244
+duktape;https://github.com/svaarala/duktape/releases/download/v2.7.0/duktape-2.7.0.tar.xz;8200c8e417dbab7adcc12c4dbdef7651cfc55794

cmake/external/onnxruntime_external_deps.cmake

@@ -772,6 +772,16 @@ if (onnxruntime_USE_WEBGPU)
   if (onnxruntime_ENABLE_PIX_FOR_WEBGPU_EP)
     list(APPEND onnxruntime_EXTERNAL_LIBRARIES webgpu_glfw glfw)
   endif()
+
+  if (NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic")
+    onnxruntime_fetchcontent_declare(
+      duktape
+      URL ${DEP_URL_duktape}
+      URL_HASH SHA1=${DEP_SHA1_duktape}
+      EXCLUDE_FROM_ALL
+    )
+    onnxruntime_fetchcontent_makeavailable(duktape)
+  endif()
 endif()
 
 if(onnxruntime_USE_COREML)

cmake/onnxruntime_providers_webgpu.cmake

@@ -9,6 +9,17 @@
   if (onnxruntime_ENABLE_WEBASSEMBLY_THREADS)
     add_definitions(-DENABLE_WEBASSEMBLY_THREADS=1)
   endif()
+  if (onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic")
+    if (onnxruntime_DISABLE_EXCEPTIONS)
+      message(FATAL_ERROR "Dynamic WGSL template generation requires exception handling to be enabled.")
+    endif()
+    if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
+      message(FATAL_ERROR "Dynamic WGSL template generation is not supported when targeting WebAssembly.")
+    endif()
+    add_definitions(-DORT_WGSL_TEMPLATE_DYNAMIC=1)
+  elseif (NOT onnxruntime_WGSL_TEMPLATE STREQUAL "static")
+    message(FATAL_ERROR "Unsupported value for onnxruntime_WGSL_TEMPLATE: ${onnxruntime_WGSL_TEMPLATE}. Supported values are 'static' or 'dynamic'.")
+  endif()
   file(GLOB_RECURSE onnxruntime_providers_webgpu_cc_srcs CONFIGURE_DEPENDS
     "${ONNXRUNTIME_ROOT}/core/providers/webgpu/*.h"
     "${ONNXRUNTIME_ROOT}/core/providers/webgpu/*.cc"
@@ -20,6 +31,7 @@
 
   source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_providers_webgpu_cc_srcs})
   onnxruntime_add_static_library(onnxruntime_providers_webgpu ${onnxruntime_providers_webgpu_cc_srcs})
+  target_compile_features(onnxruntime_providers_webgpu PRIVATE cxx_std_20)
   onnxruntime_add_include_to_target(onnxruntime_providers_webgpu
     onnxruntime_common onnx onnx_proto flatbuffers::flatbuffers Boost::mp11 safeint_interface)
 
@@ -117,4 +129,95 @@
   endif()
 
   add_dependencies(onnxruntime_providers_webgpu ${onnxruntime_EXTERNAL_DEPENDENCIES})
+
+  if (onnxruntime_WGSL_TEMPLATE)
+    # Define the WGSL templates directory and output directory
+    set(WGSL_TEMPLATES_DIR "${ONNXRUNTIME_ROOT}/core/providers/webgpu/wgsl_templates")
+    set(WGSL_GENERATED_ROOT "${CMAKE_CURRENT_BINARY_DIR}/wgsl_generated")
+
+    # Find npm and node executables
+    find_program(NPM_EXECUTABLE "npm.cmd" "npm" REQUIRED)
+    if(NOT NPM_EXECUTABLE)
+      message(FATAL_ERROR "npm is required for WGSL template generation but was not found")
+    endif()
+    find_program(NODE_EXECUTABLE "node" REQUIRED)
+    if (NOT NODE_EXECUTABLE)
+      message(FATAL_ERROR "Node is required for WGSL template generation but was not found")
+    endif()
+
+    # Install npm dependencies
+    add_custom_command(
+      OUTPUT "${WGSL_TEMPLATES_DIR}/node_modules/.install_complete"
+      COMMAND ${NPM_EXECUTABLE} ci
+      COMMAND ${CMAKE_COMMAND} -E touch "${WGSL_TEMPLATES_DIR}/node_modules/.install_complete"
+      DEPENDS "${WGSL_TEMPLATES_DIR}/package.json" "${WGSL_TEMPLATES_DIR}/package-lock.json"
+      WORKING_DIRECTORY ${WGSL_TEMPLATES_DIR}
+      COMMENT "Installing npm dependencies for WGSL template generation"
+      VERBATIM
+    )
+
+    if (onnxruntime_WGSL_TEMPLATE STREQUAL "static")
+      set(WGSL_GENERATED_DIR "${WGSL_GENERATED_ROOT}/wgsl_template_gen")
+      # set(WGSL_GEN_OUTPUTS "${WGSL_GENERATED_DIR}/index.h" "${WGSL_GENERATED_DIR}/index_impl.h")
+      # Define the output files that will be generated
+      set(WGSL_GENERATED_INDEX_H "${WGSL_GENERATED_DIR}/index.h")
+      set(WGSL_GENERATED_INDEX_IMPL_H "${WGSL_GENERATED_DIR}/index_impl.h")
+    elseif(onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic")
+      set(WGSL_GENERATED_DIR "${WGSL_GENERATED_ROOT}/dynamic")
+      # set(WGSL_GEN_OUTPUTS "${WGSL_GENERATED_DIR}/templates.js")
+      set(WGSL_GENERATED_TEMPLATES_JS "${WGSL_GENERATED_DIR}/templates.js")
+    endif()
+
+    # Ensure the output directory exists
+    file(MAKE_DIRECTORY ${WGSL_GENERATED_DIR})
+
+    # Find all WGSL template input files
+    file(GLOB_RECURSE WGSL_TEMPLATE_FILES "${ONNXRUNTIME_ROOT}/core/providers/webgpu/*.wgsl.template")
+
+    # Set wgsl-gen command line options as a list
+    set(WGSL_GEN_OPTIONS "-i" "../" "--output" "${WGSL_GENERATED_DIR}" "-I" "wgsl_template_gen/" "--preserve-code-ref" "--verbose")
+    if (onnxruntime_WGSL_TEMPLATE STREQUAL "static")
+      if (CMAKE_BUILD_TYPE STREQUAL "Debug")
+        list(APPEND WGSL_GEN_OPTIONS "--generator" "static-cpp-literal")
+      else()
+        list(APPEND WGSL_GEN_OPTIONS "--generator" "static-cpp")
+      endif()
+    elseif(onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic")
+      list(APPEND WGSL_GEN_OPTIONS "--generator" "dynamic")
+    endif()
+
+    # Generate WGSL templates
+    add_custom_command(
+      OUTPUT ${WGSL_GENERATED_INDEX_H} ${WGSL_GENERATED_INDEX_IMPL_H} ${WGSL_GENERATED_TEMPLATES_JS}
+      COMMAND ${NPM_EXECUTABLE} run gen -- ${WGSL_GEN_OPTIONS}
+      DEPENDS "${WGSL_TEMPLATES_DIR}/node_modules/.install_complete" ${WGSL_TEMPLATE_FILES}
+      WORKING_DIRECTORY ${WGSL_TEMPLATES_DIR}
+      COMMENT "Generating WGSL templates from *.wgsl.template files"
+      COMMAND_EXPAND_LISTS
+      VERBATIM
+    )
+
+    # Create a target to represent the generation step
+    add_custom_target(onnxruntime_webgpu_wgsl_generation
+      DEPENDS ${WGSL_GENERATED_INDEX_H} ${WGSL_GENERATED_INDEX_IMPL_H} ${WGSL_GENERATED_TEMPLATES_JS}
+      SOURCES ${WGSL_TEMPLATE_FILES}
+    )
+
+    if (onnxruntime_WGSL_TEMPLATE STREQUAL "static")
+      # Add the generated directory to include paths
+      target_include_directories(onnxruntime_providers_webgpu PRIVATE ${WGSL_GENERATED_ROOT})
+    elseif(onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic")
+      add_library(duktape_static STATIC "${duktape_SOURCE_DIR}/src/duktape.c")
+      target_compile_features(duktape_static PRIVATE c_std_99)
+      target_link_libraries(onnxruntime_providers_webgpu duktape_static)
+      target_include_directories(onnxruntime_providers_webgpu PRIVATE ${duktape_SOURCE_DIR}/src)
+      # Define the path to the generated templates.js file
+      target_compile_definitions(onnxruntime_providers_webgpu PRIVATE
+        "ORT_WGSL_TEMPLATES_JS_PATH=\"${WGSL_GENERATED_TEMPLATES_JS}\"")
+    endif()
+
+    # Make sure generation happens before building the provider
+    add_dependencies(onnxruntime_providers_webgpu onnxruntime_webgpu_wgsl_generation)
+  endif()
+
   set_target_properties(onnxruntime_providers_webgpu PROPERTIES FOLDER "ONNXRuntime")

onnxruntime/core/providers/webgpu/shader_helper.h

@@ -13,6 +13,8 @@
 #include "core/providers/webgpu/shader_variable.h"
 #include "core/providers/webgpu/string_utils.h"
 
+#include "core/providers/webgpu/wgsl_templates/wgsl_gen.h"
+
 namespace onnxruntime {
 namespace webgpu {
 

onnxruntime/core/providers/webgpu/tensor/pad.cc

@@ -19,73 +19,11 @@ Status PadProgram::GenerateShaderCode(ShaderHelper& shader) const {
   }
   const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseShapeAndStride | ShaderUsage::UseValueTypeAlias);
 
-  shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size");
-  std::string constant_value_str = std::string("let constant_value = ") +
-                                   (is_float16_ ? "bitcast<vec2<f16>>(uniforms.constant_value)[0];\n" : "bitcast<output_value_t>(uniforms.constant_value);\n");
-  if (dim_value_zero_) {
-    // Only Constant mode needs fill output if the one dim value or mores dims' values of input are zero.
-    shader.MainFunctionBody() << constant_value_str
-                              << "output[global_idx] = constant_value;\n";
-    return Status::OK();
-  }
-
-  shader.MainFunctionBody() << "  let output_indices = " << output.OffsetToIndices("global_idx") << ";\n"
-                            << "  var input_index = u32(0);\n"
-                            << "  var use_pad_value = false;\n"
-                            << "  var in_coord = i32(0);\n";
-
-  const int rank = output.Rank();
-  std::string output_indices_str = "i32(" + GetElementAt("output_indices", "dim", rank) + ")";
-  std::string lower_pads_str = GetElementAt("uniforms.lower_pads", "dim", rank);
-  std::string data_shape_str = "i32(" + GetElementAt("uniforms.data_shape", "dim", rank) + ")";
-  std::string data_stride_str = rank == 1 ? "" : " * " + GetElementAt("uniforms.data_stride", "dim", rank - 1);
-  SS(axis_body_ss, 1024);
-  switch (mode_) {
-    case Mode::Constant:
-      axis_body_ss << "    if (" << output_indices_str << " < " << lower_pads_str << " || " << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n"
-                   << "      use_pad_value = true;\n";
-      break;
-    case Mode::Edge:
-      axis_body_ss << "    if (" << output_indices_str << " < " << lower_pads_str << ") {\n"
-                   << "      in_coord = 0;\n"
-                   << "    } else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n"
-                   << "      in_coord = " << data_shape_str + " - 1;\n";
-      break;
-    case Mode::Reflect:
-      axis_body_ss << "    if (" << output_indices_str << " < " << lower_pads_str << " || " << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n"
-                   << "      in_coord = " << output_indices_str << " - " << lower_pads_str << ";\n"
-                   << "      if (in_coord < 0) {\n"
-                   << "        in_coord = -in_coord;\n"
-                   << "       }\n"
-                   << "       {\n"
-                   << "         let _2n_1 = 2 * (" << data_shape_str << " - 1);\n"
-                   << "        in_coord = in_coord % _2n_1;\n"
-                   << "        if(in_coord >= " << data_shape_str << ") {\n"
-                   << "          in_coord = _2n_1 - in_coord;\n"
-                   << "        }\n"
-                   << "      }\n";
-      break;
-    case Mode::Wrap:
-      axis_body_ss << "    if (" << output_indices_str << " < " << lower_pads_str << ") {\n"
-                   << "      in_coord = " << data_shape_str << " + " << output_indices_str << " - " << lower_pads_str + ";\n"
-                   << "    } else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n"
-                   << "      in_coord = " << output_indices_str << " - " << lower_pads_str << " - " << data_shape_str << ";\n";
-      break;
-    default:
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported mode type: ", static_cast<int>(mode_));
-  }
-  axis_body_ss << "    } else {\n"
-               << "      " << "in_coord = " << output_indices_str << " - " << lower_pads_str << ";\n"
-               << "    }\n";
-
-  shader.MainFunctionBody() << "  for (var dim = 0; dim < " << rank << " && !use_pad_value; dim++) {\n"
-                            << SS_GET(axis_body_ss)
-                            << "    input_index += select(u32(in_coord)" << data_stride_str << ", u32(in_coord), dim == " << rank - 1 << ");\n"
-                            << "  }\n"
-                            << "  " << constant_value_str
-                            << "  " << output.SetByOffset("global_idx", "select(data[input_index], constant_value, use_pad_value)");
-
-  return Status::OK();
+  return WGSL_TEMPLATE_APPLY(shader, "tensor/pad.wgsl.template",
+                             WGSL_TEMPLATE_PARAMETER(dim_value_zero, dim_value_zero_),
+                             WGSL_TEMPLATE_PARAMETER(is_float16, is_float16_),
+                             WGSL_TEMPLATE_PARAMETER(pad_mode, mode_),
+                             WGSL_TEMPLATE_VARIABLE(output, output));
 }
 
 Status Pad::ComputeInternal(ComputeContext& context) const {

onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template

@@ -0,0 +1,74 @@
+#define PAD_MODE_CONSTANT 0
+#define PAD_MODE_REFLECT 1
+#define PAD_MODE_EDGE 2
+#define PAD_MODE_WRAP 3
+
+#use guardAgainstOutOfBoundsWorkgroupSizes
+#use getElementAt
+#use .offsetToIndices .setByOffset .rank
+
+#param dim_value_zero
+#param is_float16
+#param pad_mode
+
+$MAIN {
+  guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size);
+
+  let constant_value =
+#if is_float16
+      bitcast<vec2<f16>>(uniforms.constant_value)[0];
+#else
+      bitcast<output_value_t>(uniforms.constant_value);
+#endif
+
+#if dim_value_zero
+  output[global_idx] = constant_value;
+#else
+  let output_indices = output.offsetToIndices(global_idx);
+  var input_index = u32(0);
+  var use_pad_value = false;
+  var in_coord = i32(0);
+
+  for (var dim = 0; dim < output.rank && !use_pad_value; dim++) {
+    let output_index = i32(getElementAt(output_indices, dim, output.rank));
+    let lower_pads = getElementAt(uniforms.lower_pads, dim, output.rank);
+    let data_shape = i32(getElementAt(uniforms.data_shape, dim, output.rank));
+#if pad_mode == PAD_MODE_CONSTANT
+    if (output_index < lower_pads || output_index >= data_shape + lower_pads) {
+        use_pad_value = true;
+#elif pad_mode == PAD_MODE_EDGE
+    if (output_index < lower_pads) {
+      in_coord = 0;
+    } else if (output_index >= data_shape + lower_pads) {
+      in_coord = data_shape - 1;
+#elif pad_mode == PAD_MODE_REFLECT
+    if (output_index < lower_pads || output_index >= data_shape + lower_pads) {
+      in_coord = output_index - lower_pads;
+      if (in_coord < 0) {
+        in_coord = -in_coord;
+      }
+      let _2n_1 = 2 * (data_shape - 1);
+      in_coord = in_coord % _2n_1;
+      if (in_coord >= data_shape) {
+        in_coord = _2n_1 - in_coord;
+      }
+#else // PAD_MODE_WRAP
+    if (output_index < lower_pads) {
+      in_coord = data_shape + output_index - lower_pads;
+    } else if (output_index >= data_shape + lower_pads) {
+      in_coord = output_index - data_shape - lower_pads;
+#endif // pad_mode
+    } else {
+        in_coord = output_index - lower_pads;
+    }
+
+    input_index += select(u32(in_coord)
+#if output.rank > 1
+        * getElementAt(uniforms.data_stride, dim, output.rank - 1)
+#endif
+        , u32(in_coord), dim == output.rank - 1);
+  }
+
+  output.setByOffset(global_idx, select(data[input_index], constant_value, use_pad_value));
+#endif
+} // MAIN

onnxruntime/core/providers/webgpu/wgsl_templates/.gitignore

@@ -0,0 +1 @@
+node_modules/