diff --git a/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml index 15ab5bda2a8ec..e0fbce88ab4fa 100644 --- a/.github/workflows/windows_webgpu.yml +++ b/.github/workflows/windows_webgpu.yml @@ -197,7 +197,6 @@ jobs: --cmake_generator "Visual Studio 17 2022" ` --enable_onnx_tests ` --use_webgpu shared_lib ` - --wgsl_template static ` --use_vcpkg --use_vcpkg_ms_internal_asset_cache ` --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON onnxruntime_ENABLE_DAWN_BACKEND_D3D12=1 onnxruntime_ENABLE_DAWN_BACKEND_VULKAN=1 ` --disable_rtti ` diff --git a/.lintrunner.toml b/.lintrunner.toml index 7f6f61df59f3b..86d644389d29e 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -116,6 +116,8 @@ exclude_patterns = [ 'winml/lib/Api.Image/shaders/**', # Contains data chunks 'onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h', # Bool Switches hang Clang 'onnxruntime/core/providers/coreml/mlprogram_test_scripts/**', # test scripts only + 'tools/python/wgsl_template/test/testcases/**', # Generated golden fixtures; must match tool output byte-for-byte + 'tools/python/wgsl_template/test/in_tree_golden/**', # Generated golden snapshots; must match tool output byte-for-byte ] command = [ 'python', diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 9b46b9dc04e56..28d6aa83c5343 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -149,7 +149,6 @@ option(onnxruntime_USE_XNNPACK "Build with XNNPACK support. Provides an alternat 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_USE_EP_API_ADAPTERS "Build EP (e.g. WebGPU) as a plugin EP shared library using EP API adapters" 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_SHARED_LIBRARY "Build Dawn as a shared library" OFF) diff --git a/cmake/deps.txt b/cmake/deps.txt index d8cd7713bfe65..76e1ea99d36b3 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -60,5 +60,6 @@ kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.20.0.tar. # kleidiai-qmx is pinned to a specific commit as there are no tagged releases. When an appropriate tagged release becomes available, # this entry will be updated to use refs/tags/ instead of the raw commit hash. kleidiai-qmx;https://github.com/qualcomm/kleidiai/archive/2f10c9a8d32f81ffeeb6d4885a29cc35d2b0da87.zip;5e855730a2d69057a569f43dd7532db3b2d2a05c +# TODO(danielsongmicrosoft): Remove duktape once the dynamic WGSL generator is deleted; it is no longer used by any active build path. duktape;https://github.com/svaarala/duktape/releases/download/v2.7.0/duktape-2.7.0.tar.xz;8200c8e417dbab7adcc12c4dbdef7651cfc55794 vulkan_headers;https://codeload.github.com/KhronosGroup/Vulkan-Headers/tar.gz/refs/tags/v1.4.344;57bc528ef7c4a3f7bfbb59e64a187e3734bd29d8 diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake index f523fefdd153a..0828db442b4ca 100644 --- a/cmake/external/onnxruntime_external_deps.cmake +++ b/cmake/external/onnxruntime_external_deps.cmake @@ -828,6 +828,9 @@ if (onnxruntime_USE_WEBGPU) list(APPEND onnxruntime_EXTERNAL_LIBRARIES webgpu_glfw glfw) endif() + # TODO: Remove duktape once the dynamic WGSL generator is deleted. It is only + # needed by that generator, which is no longer wired up, so + # onnxruntime_WGSL_TEMPLATE is unset and this block never runs. if (NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND onnxruntime_WGSL_TEMPLATE STREQUAL "dynamic") if(onnxruntime_USE_VCPKG) find_package(unofficial-duktape CONFIG REQUIRED) diff --git a/cmake/onnxruntime_providers_webgpu.cmake b/cmake/onnxruntime_providers_webgpu.cmake index 6d2994744bcfd..d4b8ca2aba96b 100644 --- a/cmake/onnxruntime_providers_webgpu.cmake +++ b/cmake/onnxruntime_providers_webgpu.cmake @@ -10,17 +10,6 @@ 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" @@ -271,36 +260,21 @@ add_dependencies(onnxruntime_providers_webgpu onnx ${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") + if (onnxruntime_USE_WEBGPU) + # The Python wgsl-gen tool lives at the repo level under tools/python. + set(WGSL_GEN_PYTHON_DIR "${REPO_ROOT}/tools/python") set(WGSL_GENERATED_ROOT "${CMAKE_CURRENT_BINARY_DIR}/wgsl_generated") - # Include the Node.js helper for finding and validating Node.js and NPM - include(node_helper.cmake) + # The top-level find_package(Python ...) in cmake/CMakeLists.txt is gated + # on BUILD_SHARED_LIB OR ENABLE_PYTHON, so Python_EXECUTABLE is not always + # set in WebGPU-enabled builds (e.g. the WASM lane). Find Python ourselves + # so this branch works in every config. + find_package(Python 3.10 COMPONENTS Interpreter REQUIRED) - # Install npm dependencies - add_custom_command( - OUTPUT "${WGSL_TEMPLATES_DIR}/node_modules/.install_complete" - COMMAND ${NPM_CLI} 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() + set(WGSL_GENERATED_DIR "${WGSL_GENERATED_ROOT}/wgsl_template_gen") + # 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") # Ensure the output directory exists file(MAKE_DIRECTORY ${WGSL_GENERATED_DIR}) @@ -324,47 +298,42 @@ list(APPEND WGSL_GEN_OPTIONS "-i" "${ONNXRUNTIME_ROOT}/contrib_ops/webgpu") endif() - 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") + 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() # Generate WGSL templates add_custom_command( - OUTPUT ${WGSL_GENERATED_INDEX_H} ${WGSL_GENERATED_INDEX_IMPL_H} ${WGSL_GENERATED_TEMPLATES_JS} - COMMAND ${NPM_CLI} 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" + OUTPUT ${WGSL_GENERATED_INDEX_H} ${WGSL_GENERATED_INDEX_IMPL_H} + COMMAND ${Python_EXECUTABLE} "${WGSL_GEN_PYTHON_DIR}/wgsl_gen.py" ${WGSL_GEN_OPTIONS} + DEPENDS ${WGSL_TEMPLATE_FILES} + WORKING_DIRECTORY ${WGSL_GEN_PYTHON_DIR} + COMMENT "Generating WGSL templates from *.wgsl.template files (Python)" 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} + DEPENDS ${WGSL_GENERATED_INDEX_H} ${WGSL_GENERATED_INDEX_IMPL_H} 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") - target_link_libraries(onnxruntime_providers_webgpu PRIVATE duktape_static) - onnxruntime_add_include_to_target(onnxruntime_providers_webgpu duktape_static) - - # 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() + # Add the generated directory to include paths + target_include_directories(onnxruntime_providers_webgpu PRIVATE ${WGSL_GENERATED_ROOT}) # Make sure generation happens before building the provider add_dependencies(onnxruntime_providers_webgpu onnxruntime_webgpu_wgsl_generation) + + # Wire the Python wgsl_template test suite into ctest. + if (BUILD_TESTING) + add_test( + NAME wgsl_template_python_tests + COMMAND ${Python_EXECUTABLE} "${WGSL_GEN_PYTHON_DIR}/wgsl_template/test/run_tests.py" + WORKING_DIRECTORY ${WGSL_GEN_PYTHON_DIR} + ) + endif() endif() if (NOT onnxruntime_BUILD_SHARED_LIB) diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/.gitignore b/onnxruntime/core/providers/webgpu/wgsl_templates/.gitignore deleted file mode 100644 index c2658d7d1b318..0000000000000 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/.gitignore +++ /dev/null @@ -1 +0,0 @@ -node_modules/ diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/README.md b/onnxruntime/core/providers/webgpu/wgsl_templates/README.md index 6bd2f98cc5713..2f99daa425265 100644 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/README.md +++ b/onnxruntime/core/providers/webgpu/wgsl_templates/README.md @@ -1,34 +1,40 @@ # WebGPU WGSL Templates -This directory contains the infrastructure, scripts, and documentation for the WGSL template system used by the ONNX Runtime WebGPU Execution Provider (EP). The template system enables the generation of optimized WGSL shaders at build time or runtime, with parameterization and reusability across different operators. +This directory contains the infrastructure and documentation for the WGSL +template system used by the ONNX Runtime WebGPU Execution Provider (EP). The +template system generates optimized WGSL shaders at build time, with +parameterization and reusability across different operators. -For detailed information about the underlying template engine, see [wgsl-template](https://github.com/fs-eire/wgsl-template). +The template engine is implemented in Python and lives at +[`tools/python/wgsl_gen.py`](../../../../../tools/python/wgsl_gen.py) (with the +supporting package at +[`tools/python/wgsl_template/`](../../../../../tools/python/wgsl_template/)). +It requires only a Python 3.10+ interpreter. ## Overview -The WGSL template system provides a flexible framework for generating WebGPU shaders for ONNX Runtime operators. Instead of writing static shader code, developers can create parameterized templates that adapt to different input configurations, data types, and optimization requirements. +The WGSL template system provides a flexible framework for generating WebGPU +shaders for ONNX Runtime operators. Instead of writing static shader code, +developers can create parameterized templates that adapt to different input +configurations, data types, and optimization requirements. **Key Benefits:** - **Code Reusability**: Share common shader patterns across multiple operators -- **Type Safety**: Template parameters are validated at build time (static mode) +- **Type Safety**: Template parameters are validated at build time - **Performance**: Generated shaders are optimized for specific use cases - **Maintainability**: Centralized shader logic with clear parameterization ## Terms - **WGSL**: The **W**eb**G**PU **S**hader **L**anguage - the shading language for WebGPU -- **WGSL Template File**: A file with `.wgsl.template` extension containing WGSL shader code with template syntax and utilities -- **WGSL-Template**: A JavaScript library developed specifically for generating WGSL code from template files +- **WGSL Template File**: A file with the `.wgsl.template` extension containing WGSL shader code with template syntax and utilities - **Template Parameters**: Configuration objects that control shader generation (data types, dimensions, etc.) ## How It Works -The system supports two generation modes to balance performance and flexibility: - -### Static Generation Mode - -Templates are processed at **build time** to generate C++ header files embedded in the WebGPU EP binary. +Templates are processed at **build time** to generate C++ header files embedded +in the WebGPU EP binary. **Advantages:** @@ -37,31 +43,16 @@ Templates are processed at **build time** to generate C++ header files embedded - Type-safe template parameters validated at compile time - Optimal for production deployments -**Use Cases:** - -- Production builds - -### Dynamic Generation Mode - -Templates are processed at **runtime** using an embedded JavaScript engine to generate shaders on-demand. - -**Advantages:** - -- Faster development iteration (no rebuild required) -- Support for runtime-dependent parameters -- Easier debugging and experimentation -- Flexible for research and development - -**Use Cases:** - -- Development and debugging workflows -- Research environments with frequently changing parameters +CMake invokes the Python tool, which walks the `.wgsl.template` files and emits +`index.h` / `index_impl.h` (plus per-template headers) into the build directory. +The EP includes these generated headers via [`wgsl_gen.h`](wgsl_gen.h) / +[`wgsl_gen.cc`](wgsl_gen.cc). ## Development -This section includes instructions for how to use the template system in the development. +This section describes how to use the template system during development. -1. Create WGSL template files in `.wgsl.template` extension. +1. Create WGSL template files with the `.wgsl.template` extension. - [Reference: Template Syntax](https://github.com/fs-eire/wgsl-template?tab=readme-ov-file#template-syntax) - [Reference: Built-in Utilities](https://github.com/fs-eire/wgsl-template?tab=readme-ov-file#Utilities) @@ -73,48 +64,52 @@ This section includes instructions for how to use the template system in the dev 3. Build. - - Using static code generator - - Static code generator is enabled by default: - - ```sh - ./build.sh --use_webgpu - ``` - - Rebuild is needed when any C/C++ source file or WGSL template file is updated. - - - Using dynamic code generator - - Append `--wgsl_template dynamic` to the ORT build script: + The static code generator is always enabled when WebGPU is built: - ```sh - ./build.sh --use_webgpu --wgsl_template dynamic - ``` + ```sh + ./build.sh --use_webgpu + ``` - Rebuild is not needed when only WGSL template files are updated. In this case, you just need to compile the template files. See the next section below. + A rebuild is needed when any C/C++ source file or WGSL template file is updated. -4. Compile the template files. (for dynamic code generator only) +## Python tool reference - When you are using dynamic code generator, you don't need to rebuild ONNX Runtime when you made changes to the WGSL template files. +The build invokes [`tools/python/wgsl_gen.py`](../../../../../tools/python/wgsl_gen.py) +directly from CMake; you should not normally need to run it by hand. The CLI +surface is: - > Suppose `` is the build directory. It's usually something like `/build/Linux/Debug` if not explicitly specified. +``` +python tools/python/wgsl_gen.py \ + -i [-i ...] \ + --output \ + --generator {static-cpp|static-cpp-literal} \ + [-I ] \ + [--ext .wgsl.template] \ + [--preserve-code-ref] \ + [--clean] \ + [--verbose] +``` - There are 2 ways to compile the template files: +* `static-cpp` (Release): emits short `__str_N` identifiers backed by a `string_table.h` for shader-string deduplication. +* `static-cpp-literal` (Debug): inlines string literals; easier to read while debugging. - 1. Use a NPM script in the current folder: +### Running the test suite - ```sh - npm run gen -- -i ../ --preserve-code-ref --verbose --generator dynamic --output /wgsl_generated/dynamic - ``` +The Python tool ships with a unit + fixture test suite. CMake adds +`wgsl_template_python_tests` to `ctest`, so any standard ORT build with WGSL +templates enabled will run them: - This script will generate the file `/wgsl_generated/dynamic/templates.js` and exit. +``` +ctest -R wgsl_template_python_tests +``` - 2. Use the same script but in watch mode (recommended): +You can also run the suite directly from the source tree: - Use the same script as above but append `--watch` flag. This will launch a service monitoring the file system change and automatically compile the templates if any change occurred. +``` +python tools/python/wgsl_template/test/run_tests.py +``` - The typical development workflow is: - 1. Build ORT once with dynamic template mode - 2. Launch wgsl-gen in watch mode - 3. Run ORT to debug/validate the shader - 4. Make changes to the template files, and repeat step (c) +Tests cover the loader, parser, generator, build orchestrator, and a smoke test +against the in-tree templates (Pad, Transpose, im2col-matmul). The fixtures live +under +[`tools/python/wgsl_template/test/testcases/`](../../../../../tools/python/wgsl_template/test/testcases). diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/package-lock.json b/onnxruntime/core/providers/webgpu/wgsl_templates/package-lock.json deleted file mode 100644 index 73662e8ebf45f..0000000000000 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/package-lock.json +++ /dev/null @@ -1,37 +0,0 @@ -{ - "name": "ort_wgsl_template_gen", - "version": "1.0.0", - "lockfileVersion": 3, - "requires": true, - "packages": { - "": { - "name": "ort_wgsl_template_gen", - "version": "1.0.0", - "license": "MIT", - "dependencies": { - "@fs-eire/wgsl-template": "^0.1.15" - } - }, - "node_modules/@fs-eire/wgsl-template": { - "version": "0.1.15", - "resolved": "https://registry.npmjs.org/@fs-eire/wgsl-template/-/wgsl-template-0.1.15.tgz", - "integrity": "sha512-Sx+hUpef7hoEQ3Af+NGURcEukXDnvB/4TnkIuu5G+toGTwuw6pCdLpqoaf0fi3XHoZFMREtdq0vKjdrLqkPeqg==", - "license": "MIT", - "dependencies": { - "minimist": "^1.2.8" - }, - "bin": { - "wgsl-gen": "dist/bin/cli.js" - } - }, - "node_modules/minimist": { - "version": "1.2.8", - "resolved": "https://registry.npmjs.org/minimist/-/minimist-1.2.8.tgz", - "integrity": "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==", - "license": "MIT", - "funding": { - "url": "https://github.com/sponsors/ljharb" - } - } - } -} diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/package.json b/onnxruntime/core/providers/webgpu/wgsl_templates/package.json deleted file mode 100644 index bbeb58f6bcd52..0000000000000 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/package.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "name": "ort_wgsl_template_gen", - "version": "1.0.0", - "private": true, - "description": "", - "main": "index.js", - "scripts": { - "gen": "wgsl-gen" - }, - "author": "", - "license": "MIT", - "dependencies": { - "@fs-eire/wgsl-template": "^0.1.15" - } -} diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.cc b/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.cc index 3ded2865d51d5..59adafad60a70 100644 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.cc +++ b/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.cc @@ -56,6 +56,10 @@ namespace wgsl_gen { #else // Use dynamic generator +// TODO(danielsongmicrosoft): Remove the dynamic WGSL generator and its duktape +// dependency. The static Python generator is the only supported mode, so +// ORT_WGSL_TEMPLATE_DYNAMIC is never defined and everything below is unreachable. + #include "duktape.h" namespace { diff --git a/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.h b/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.h index 83c8074825d0b..9dd59336722e2 100644 --- a/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.h +++ b/onnxruntime/core/providers/webgpu/wgsl_templates/wgsl_gen.h @@ -78,6 +78,10 @@ onnxruntime::common::Status ApplyTemplate(ShaderHelper& shader_helper, TemplateP #else // Use dynamic generator +// TODO(danielsongmicrosoft): Remove the dynamic WGSL generator. The static +// Python generator is the only supported mode, so ORT_WGSL_TEMPLATE_DYNAMIC is +// never defined and this branch is unreachable. + #define WGSL_TEMPLATE_PARAMETER(name, value) \ onnxruntime::webgpu::wgsl_gen::TemplateParam(#name, static_cast(value)) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 310bf67928f46..231888f2204a8 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -263,8 +263,6 @@ def generate_vcpkg_install_options(build_dir, args): vcpkg_install_options.append("--x-feature=vsinpu-ep") if args.use_webgpu: vcpkg_install_options.append("--x-feature=webgpu-ep") - if args.wgsl_template == "dynamic": - vcpkg_install_options.append("--x-feature=webgpu-ep-wgsl-template-dynamic") if args.use_webnn: vcpkg_install_options.append("--x-feature=webnn-ep") if args.use_xnnpack: @@ -462,7 +460,6 @@ def generate_build_tree( "-Donnxruntime_USE_JSEP=" + ("ON" if args.use_jsep else "OFF"), "-Donnxruntime_USE_WEBGPU=" + ("ON" if args.use_webgpu else "OFF"), "-Donnxruntime_USE_EXTERNAL_DAWN=" + ("ON" if args.use_external_dawn else "OFF"), - "-Donnxruntime_WGSL_TEMPLATE=" + args.wgsl_template, # Training related flags "-Donnxruntime_ENABLE_NVTX_PROFILE=" + ("ON" if args.enable_nvtx_profile else "OFF"), "-Donnxruntime_ENABLE_TRAINING=" + ("ON" if args.enable_training else "OFF"), diff --git a/tools/ci_build/build_args.py b/tools/ci_build/build_args.py index e0ee8dddbac1b..abc9a8d91779c 100644 --- a/tools/ci_build/build_args.py +++ b/tools/ci_build/build_args.py @@ -830,9 +830,9 @@ def add_execution_provider_args(parser: argparse.ArgumentParser) -> None: ) webgpu_group.add_argument( "--wgsl_template", - choices=["static", "dynamic"], - default="static", # By default, use static WGSL template generation - help="Specify the generator for WebGPU WGSL template generation.", + choices=["static"], + default="static", + help="Deprecated no-op. WGSL template generation is always static; kept for backward compatibility.", ) # --- XNNPACK --- diff --git a/tools/python/wgsl_gen.py b/tools/python/wgsl_gen.py new file mode 100644 index 0000000000000..245b3fb0f9dd9 --- /dev/null +++ b/tools/python/wgsl_gen.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""CLI entry point for the WGSL template generator. + +Invoked from CMake during the WebGPU EP build to translate +``*.wgsl.template`` files into C++ headers. + +Usage:: + + python wgsl_gen.py \\ + -i [-i ...] \\ + --output \\ + --generator {static-cpp|static-cpp-literal} \\ + [-I ] \\ + [--ext .wgsl.template] \\ + [--preserve-code-ref] \\ + [--clean] \\ + [--verbose] +""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Sequence +from pathlib import Path + +# Make the sibling wgsl_template package importable when invoked from any cwd. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from wgsl_template import build + + +def _build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="wgsl_gen.py", + description="Generate C++ headers from WGSL template files.", + ) + + parser.add_argument( + "-i", + "--input", + action="append", + default=[], + metavar="DIR", + help="Source directory to scan for template files. May be repeated.", + ) + parser.add_argument( + "-o", + "--output", + required=False, + metavar="DIR", + help="Output directory for generated files.", + ) + parser.add_argument( + "-g", + "--generator", + choices=["static-cpp", "static-cpp-literal"], + default="static-cpp-literal", + help=( + "Code generator to use (default: static-cpp-literal). " + "'static-cpp' deduplicates shader strings into a shared string " + "table referenced by short __str_N ids (smaller binary; used by " + "release builds). 'static-cpp-literal' inlines each string as a " + "C++ literal (larger, but easier to read when debugging; used by " + "Debug builds)." + ), + ) + parser.add_argument( + "-I", + "--include-prefix", + default="", + metavar="PREFIX", + help=( + "Prefix prepended to the #include paths emitted in index_impl.h " + '(e.g. "wgsl_template_gen/"). Must match how the generated output ' + "directory is reachable from the C++ include search paths, or the " + "includes will not resolve. Defaults to empty (includes relative " + "to index_impl.h's directory)." + ), + ) + parser.add_argument( + "-e", + "--ext", + default=".wgsl.template", + metavar="EXT", + help="Template file extension to scan for (default: .wgsl.template).", + ) + parser.add_argument( + "--preserve-code-ref", + action="store_true", + help="Emit source line references as comments before each generated chunk.", + ) + parser.add_argument( + "--clean", + action="store_true", + help="Recursively delete the output directory before generating.", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Print verbose progress information.", + ) + parser.add_argument( + "--version", + action="version", + version="wgsl_gen.py (Python port) 0.1.0", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _build_arg_parser() + args = parser.parse_args(argv) + + # argparse can't express "this option is required only when no + # subcommand is given"; validate -i/--input and --output here. + errors: list[str] = [] + if not args.input: + errors.append("at least one -i/--input directory is required") + if not args.output: + errors.append("--output is required") + if errors: + for err in errors: + print(f"error: {err}", file=sys.stderr) + parser.print_usage(sys.stderr) + return 2 + + if args.verbose: + print("wgsl_gen.py invoked with:") + print(f" inputs: {args.input}") + print(f" output: {args.output}") + print(f" generator: {args.generator}") + print(f" include-prefix: {args.include_prefix!r}") + print(f" ext: {args.ext}") + print(f" preserve-code-ref: {args.preserve_code_ref}") + print(f" clean: {args.clean}") + + build( + source_dirs=[Path(p) for p in args.input], + out_dir=Path(args.output), + template_ext=args.ext, + generator=args.generator, + include_path_prefix=args.include_prefix, + preserve_code_reference=args.preserve_code_ref, + clean=args.clean, + verbose=args.verbose, + ) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/tools/python/wgsl_template/__init__.py b/tools/python/wgsl_template/__init__.py new file mode 100644 index 0000000000000..e98d0bed00f5f --- /dev/null +++ b/tools/python/wgsl_template/__init__.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""WGSL template engine. + +The public entry point is :func:`build`, which orchestrates the +load -> parse -> generate -> emit pipeline. +""" + +from __future__ import annotations + +import os +import shutil +from collections.abc import Sequence +from pathlib import Path + +from .code_generator import resolve_code_generator +from .errors import ( + WgslTemplateBuildError, + WgslTemplateError, + WgslTemplateGenerateError, + WgslTemplateLoadError, + WgslTemplateParseError, +) +from .generator import generate_directory +from .loader import load_from_directories +from .parser import parse +from .types import SourceDir + + +def build( + *, + source_dirs: Sequence[str | os.PathLike[str] | SourceDir], + out_dir: str | os.PathLike[str], + template_ext: str = ".wgsl.template", + generator: str = "static-cpp-literal", + include_path_prefix: str = "", + preserve_code_reference: bool = False, + clean: bool = False, + verbose: bool = False, +) -> list[str]: + """Run the full build pipeline: load → parse → generate → emit. + + Returns the list of relative output paths that were inspected (one + per emitted file). Files whose contents are unchanged are not + rewritten, which keeps CMake's mtime-based incremental rebuild + tracking honest. + """ + + if not source_dirs: + raise WgslTemplateBuildError( + "source_dirs must be provided and cannot be empty", + "invalid-options", + ) + + out_path = Path(out_dir).resolve() + + # --clean wipes the output directory before building. + if clean and out_path.exists(): + if verbose: + print(f"Cleaning output directory: {out_path}") + if out_path.is_file(): + raise WgslTemplateBuildError( + f'Output path "{out_path}" is an existing file; expected a directory', + "invalid-options", + ) + shutil.rmtree(out_path) + + if verbose: + print("Building WGSL templates...") + print(f" Sources: {list(source_dirs)}") + print(f" Output: {out_path}") + print(f" Generator: {generator}") + print(f" Template extension: {template_ext}") + if include_path_prefix: + print(f" Include prefix: {include_path_prefix}") + if preserve_code_reference: + print(" Preserve code references: enabled") + + # PASS0 - load. + pass0 = load_from_directories(source_dirs, ext=template_ext) + # PASS1 - parse. + pass1 = parse(pass0) + # PASS2 - generate. + code_generator = resolve_code_generator(generator) + pass2 = generate_directory( + pass1, + code_generator, + preserve_code_reference=preserve_code_reference, + ) + # Build (per-template + index*.h + string_table.h). + files = code_generator.build( + pass2, + template_ext=template_ext, + include_path_prefix=include_path_prefix, + ) + + # Write files. Idempotent: only overwrite when content differs. + out_path.mkdir(parents=True, exist_ok=True) + + written: list[str] = [] + for rel_path, content in files.templates.items(): + assert isinstance(content, str) + full_path = (out_path / rel_path).resolve() + # Security: refuse to write outside the output directory. + try: + full_path.relative_to(out_path) + except ValueError as e: + raise WgslTemplateBuildError( + f"Security violation: attempted to write file outside output directory: {rel_path}", + "path-security-violation", + file_path=str(rel_path), + ) from e + + full_path.parent.mkdir(parents=True, exist_ok=True) + + # Idempotent write: only rewrite when content differs. This + # preserves mtime so CMake's incremental builds stay clean. + # Force LF line endings unconditionally for stable output + # across Windows / Linux. + normalized = content.replace("\r\n", "\n") + new_bytes = normalized.encode("utf-8") + existing_bytes: bytes | None = None + if full_path.exists(): + try: + existing_bytes = full_path.read_bytes() + except OSError: + existing_bytes = None + if existing_bytes != new_bytes: + full_path.write_bytes(new_bytes) + written.append(rel_path) + + if verbose: + print(f"Build completed successfully! ({len(written)} file(s))") + + return written + + +__all__ = [ + "SourceDir", + "WgslTemplateBuildError", + "WgslTemplateError", + "WgslTemplateGenerateError", + "WgslTemplateLoadError", + "WgslTemplateParseError", + "build", +] diff --git a/tools/python/wgsl_template/code_generator/__init__.py b/tools/python/wgsl_template/code_generator/__init__.py new file mode 100644 index 0000000000000..59b7cd5173ad8 --- /dev/null +++ b/tools/python/wgsl_template/code_generator/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Code generators that emit C++ from PASS2 results. + +The single :class:`StaticCodeGenerator` covers both ``static-cpp`` and +``static-cpp-literal`` (the difference is whether a string table is +used). +""" + +from __future__ import annotations + +from ..errors import WgslTemplateGenerateError +from .static_cpp import ( + CodeSegment, + CodeSegmentArg, + StaticCodeGenerator, +) + + +def resolve_code_generator(name: str) -> StaticCodeGenerator: + """Resolve a generator name to a fresh code-generator instance. + + Only the static C++ generators are supported; the dynamic generator + is provided by the npm path and is not implemented here. + """ + if name == "static-cpp": + return StaticCodeGenerator(use_string_table=True) + if name == "static-cpp-literal": + return StaticCodeGenerator(use_string_table=False) + raise WgslTemplateGenerateError( + f"Unknown code generator: {name}", + "generator-not-found", + ) + + +__all__ = [ + "CodeSegment", + "CodeSegmentArg", + "StaticCodeGenerator", + "resolve_code_generator", +] diff --git a/tools/python/wgsl_template/code_generator/static_cpp.py b/tools/python/wgsl_template/code_generator/static_cpp.py new file mode 100644 index 0000000000000..f25257acc828f --- /dev/null +++ b/tools/python/wgsl_template/code_generator/static_cpp.py @@ -0,0 +1,355 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Static C++ code generator. + +The single :class:`StaticCodeGenerator` class implements both +``static-cpp`` (uses a string table) and ``static-cpp-literal`` +(inlines string literals); the only difference is the constructor +flag. +""" + +from __future__ import annotations + +import builtins +import hashlib +import json +from dataclasses import dataclass + +from ..types import TemplatePass2, TemplateRepository + +# ---------------------------------------------------------------------- +# Code segment data model +# ---------------------------------------------------------------------- + + +@dataclass +class CodeSegment: + type: str # "raw" | "code" | "expression" + content: str + + +@dataclass +class CodeSegmentArg: + type: str # "expression" | "string" | "auto" + code: list[CodeSegment] + + +# ---------------------------------------------------------------------- +# String literal renderer (JSON-style escaping) +# ---------------------------------------------------------------------- + + +def _js_stringify(s: str) -> str: + """Render ``s`` as a double-quoted, JSON-escaped C++ string literal. + + Non-ASCII characters are escaped via ``\\uXXXX`` so the output is + pure ASCII and stable across hosts. + """ + return json.dumps(s, ensure_ascii=True) + + +# ---------------------------------------------------------------------- +# Static code generator +# ---------------------------------------------------------------------- + + +class StaticCodeGenerator: + """Emits C++ from PASS2 :class:`CodeSegment` lists. + + With ``use_string_table=True`` (the ``static-cpp`` Release mode), + every WGSL string literal is replaced with a short ``__str_N`` + identifier and dedup'd across the whole build. With + ``use_string_table=False`` (Debug mode), strings appear inline as + ``"..."`` literals. + """ + + def __init__(self, use_string_table: bool = True) -> None: + self.use_string_table = use_string_table + # Maps unique string content to its assigned __str_N id. + # Insertion order is the assignment order. + self._string_table: dict[str, int] | None = {} if use_string_table else None + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _render_string(self, s: str) -> str: + if self._string_table is not None: + existing = self._string_table.get(s) + if existing is None: + existing = len(self._string_table) + self._string_table[s] = existing + return f"__str_{existing}" + return _js_stringify(s) + + def _render_arg(self, arg: CodeSegmentArg) -> str: + if not arg.code: + return "" + + if len(arg.code) == 1: + seg = arg.code[0] + if arg.type == "string" and seg.type == "expression": + return f"wgsl_detail::pass_as_string({seg.content})" + if arg.type != "expression" and seg.type == "code": + return self._render_string(seg.content) + return seg.content + + # Multi-segment argument. + if arg.type != "expression": + rendered = [] + for seg in arg.code: + if seg.type == "code": + rendered.append(self._render_string(seg.content)) + else: + rendered.append(seg.content) + return f"absl::StrCat({', '.join(rendered)})" + return "".join(seg.content for seg in arg.code) + + # ------------------------------------------------------------------ + # CodeGenerator interface (called from the PASS2 generator) + # ------------------------------------------------------------------ + + def emit(self, code: list[CodeSegment]) -> str: + out = [] + for seg in code: + if seg.type == "raw": + out.append(seg.content) + elif seg.type == "code": + out.append(f"ss << {self._render_string(seg.content)};\n") + elif seg.type == "expression": + out.append(f"ss << {seg.content};\n") + else: # pragma: no cover defensive + raise ValueError(f"Unknown segment type: {seg.type!r}") + return "".join(out) + + def param(self, name: str) -> str: + return f"__param_{name}" + + def variable(self, name: str) -> str: + return f"__var_{name}" + + def property(self, obj: str, property_name: str) -> str: + return f"__var_{obj}.{property_name}" + + def function(self, name: str, args: list[CodeSegmentArg]) -> str: + rendered = ", ".join(self._render_arg(a) for a in args) + return f"{name}({rendered})" + + def method(self, obj: str, method_name: str, args: list[CodeSegmentArg]) -> str: + rendered = ", ".join(self._render_arg(a) for a in args) + return f"__var_{obj}.{method_name}({rendered})" + + # ------------------------------------------------------------------ + # Read-only access to the string table (used by build() to emit + # string_table.h). + # ------------------------------------------------------------------ + + @builtins.property + def string_table(self) -> dict[str, int] | None: + return self._string_table + + # ------------------------------------------------------------------ + # build() — produce the per-template files plus index.h / + # index_impl.h / string_table.h. + # ------------------------------------------------------------------ + + def build( + self, + repo: TemplateRepository, + *, + template_ext: str, + include_path_prefix: str = "", + ) -> TemplateRepository: + result_files: dict[str, str] = {} + impl_hashes: dict[str, str] = {} + + # STEP 1 — per-template implementation files. + for name, template in repo.templates.items(): + assert isinstance(template, TemplatePass2) + if not name.endswith(template_ext): + raise ValueError(f'Template name "{name}" does not end with the expected extension "{template_ext}"') + base = name[: -len(template_ext)] + content = self._build_template_implementation(name, template) + result_files[f"generated/{base}.h"] = content + impl_hashes[name] = hashlib.sha256(content.encode("utf-8")).hexdigest() + + # STEP 2 — string_table.h if enabled. + if self._string_table is not None: + content = self._build_generate_string_table() + result_files["string_table.h"] = content + impl_hashes["string_table.h"] = hashlib.sha256(content.encode("utf-8")).hexdigest() + + # STEP 3 — index_impl.h. + result_files["index_impl.h"] = self._build_generate_index_impl( + repo, impl_hashes, include_path_prefix, template_ext + ) + + # STEP 4 — index.h. + result_files["index.h"] = self._build_generate_index(repo) + + return TemplateRepository(base_path=repo.base_path, templates=dict(result_files)) + + # ---- internal builders ------------------------------------------- + + def _build_generate_index(self, repo: TemplateRepository) -> str: + out: list[str] = [] + out.append("// This file is auto-generated by wgsl-gen. Do not edit manually.") + out.append("") + out.append("#ifndef INCLUDED_BY_WGSL_GEN_HEADER") + out.append('#error "This file is expected to be included by wgsl-gen header. Do not include it directly."') + out.append("#endif") + out.append("") + for name, template in repo.templates.items(): + assert isinstance(template, TemplatePass2) + out.append("//") + out.append(f"// Template: {name}") + out.append("//") + out.append("") + quoted = _js_stringify(name) + out.append("template <>") + out.append(f"struct TemplateParameter<{quoted}> {{") + out.append(" using type = struct {") + for param_name in template.generate_result.params: + out.append(f" int param_{param_name};") + for var_name in template.generate_result.variables: + out.append(f" const ShaderVariableHelper* var_{var_name};") + out.append(" };") + out.append("};") + out.append("") + out.append("template <>") + out.append( + f"Status ApplyTemplate<{quoted}>(ShaderHelper& shader_helper, " + f"TemplateParameter<{quoted}>::type params);" + ) + out.append("") + return "\n".join(out) + + def _build_generate_string_table(self) -> str: + if self._string_table is None: + raise RuntimeError("String table is not enabled") + + out: list[str] = [] + out.append("// This file is auto-generated by wgsl-gen. Do not edit manually.") + out.append("") + out.append("#pragma once") + out.append("#ifndef INCLUDED_BY_WGSL_GEN_IMPL") + out.append('#error "This file is expected to be included by wgsl-gen impl. Do not include it directly."') + out.append("#endif") + out.append("") + out.append("// String table constants") + + # Sort by id to ensure consistent output. + sorted_entries = sorted(self._string_table.items(), key=lambda kv: kv[1]) + for s, sid in sorted_entries: + out.append(f"constexpr const char* __str_{sid} = {_js_stringify(s)};") + + out.append("") + return "\n".join(out) + + def _build_generate_index_impl( + self, + repo: TemplateRepository, + impl_hashes: dict[str, str], + include_path_prefix: str, + template_ext: str, + ) -> str: + out: list[str] = [] + out.append( + "// This file is auto-generated by wgsl-gen. Do not edit manually.\n" + "\n" + "#pragma once\n" + "#ifndef INCLUDED_BY_WGSL_GEN_IMPL\n" + '#error "This file is expected to be included by wgsl-gen impl. ' + 'Do not include it directly."\n' + "#endif\n" + "\n" + "// Helper functions or macros\n" + "\n" + '#pragma push_macro("MainFunctionStart")\n' + "#undef MainFunctionStart\n" + "#define MainFunctionStart() { [[maybe_unused]] auto& ss = " + "shader_helper.MainFunctionBody();\n" + '#pragma push_macro("MainFunctionEnd")\n' + "#undef MainFunctionEnd\n" + "#define MainFunctionEnd() }\n" + "\n" + "// Helper templates\n" + "\n" + "namespace wgsl_detail {\n" + "template >>\n" + "std::string pass_as_string(T&& v) {\n" + " return std::to_string(std::forward(v));\n" + "}\n" + "template \n" + "std::string_view pass_as_string(std::string_view sv) {\n" + " return sv;\n" + "}\n" + "template \n" + "std::string pass_as_string(T&& v) {\n" + " return std::forward(v);\n" + "}\n" + "} // namespace wgsl_detail\n" + ) + + if self._string_table is not None: + sthash = impl_hashes.get("string_table.h", "") + out.append(f'#include "{include_path_prefix}string_table.h" // {sthash}') + out.append("") + out.append("// Include template implementations") + out.append("") + + for name in repo.templates: + if not name.endswith(template_ext): + raise ValueError(f'Template name "{name}" does not end with the expected extension "{template_ext}"') + base = name[: -len(template_ext)] + h = impl_hashes.get(name, "") + out.append(f'#include "{include_path_prefix}generated/{base}.h" // {h}') + + out.append("") + out.append('#pragma pop_macro("MainFunctionStart")') + out.append('#pragma pop_macro("MainFunctionEnd")') + + return "\n".join(out) + + def _build_template_implementation(self, file_path: str, template: TemplatePass2) -> str: + out: list[str] = [] + out.append("// This file is auto-generated by wgsl-gen. Do not edit manually.") + out.append("") + out.append("#pragma once") + out.append("") + out.append("// Template implementation") + out.append(f"// Source: {file_path}") + out.append("") + + gr = template.generate_result + params_unused = len(gr.params) == 0 and len(gr.variables) == 0 + quoted = _js_stringify(file_path) + + out.append("template <>") + params_arg = "" if params_unused else "params" + out.append( + f"Status ApplyTemplate<{quoted}>(ShaderHelper& shader_helper, " + f"TemplateParameter<{quoted}>::type {params_arg}) {{" + ) + out.append(" [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation();") + out.append("") + + if gr.params: + out.append(" // Extract parameters") + for param_name in gr.params: + out.append(f" auto& {self.param(param_name)} = params.param_{param_name};") + out.append("") + + if gr.variables: + out.append(" // Extract variables") + for var_name in gr.variables: + out.append(f" auto& {self.variable(var_name)} = *params.var_{var_name};") + out.append("") + + out.append(gr.code) + + out.append("") + out.append(" return Status::OK();") + out.append("}") + + return "\n".join(out) diff --git a/tools/python/wgsl_template/code_pattern.py b/tools/python/wgsl_template/code_pattern.py new file mode 100644 index 0000000000000..6e6892e555a97 --- /dev/null +++ b/tools/python/wgsl_template/code_pattern.py @@ -0,0 +1,232 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Built-in code patterns and indices-helper rewrite table. + +A :class:`CodePattern` is a single rewrite rule consumed by the PASS2 +generator. ``DEFAULT_PATTERNS`` are always active; the named entries +in :func:`lookup_pattern` are opt-in via ``#use NAME`` directives. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from re import Pattern + +# ---------------------------------------------------------------------- +# Pattern data class +# ---------------------------------------------------------------------- + + +@dataclass +class CodePattern: + """A single rewrite rule for the PASS2 generator. + + ``type`` selects how the matcher uses this pattern: + * "control" — punctuation tokens: ``(``, ``)``, ``,``, ``{``, + ``}``, ``$MAIN {``. + * "param" — parameter name introduced by ``#param``. + * "variable" — bare shader variable identifier. + * "function" — function call: name followed by ``(``. + * "method" — method call on a shader variable: ``var.name(...)``. + * "property" — property access on a shader variable: ``var.prop``. + """ + + type: str + pattern: str | Pattern[str] + + # If present, replaces matched group(s). For function/method/property + # patterns: index 0 replaces the receiver capture (group 1), index 1 + # replaces the name capture (group 2). A ``None`` element means + # "leave that group as-is". + replace: str | list[str | None] | None = None + + variable_type: str | None = None # "shader-variable" + param_type: str | None = None # "int" + arg_types: list[str] = field(default_factory=list) # "expression"|"string"|"auto" + + +# ---------------------------------------------------------------------- +# Default control-token patterns (always active) +# ---------------------------------------------------------------------- + + +DEFAULT_PATTERNS: list[CodePattern] = [ + CodePattern(type="control", pattern=re.compile(r"\(")), + CodePattern(type="control", pattern=re.compile(r"\)")), + CodePattern(type="control", pattern=re.compile(r",")), + CodePattern(type="control", pattern=re.compile(r"\{")), + CodePattern(type="control", pattern=re.compile(r"\}")), + CodePattern( + type="control", + pattern=re.compile(r"(? CodePattern | None: + """Look up a built-in pattern by ``#use`` name. Returns ``None`` if + the name isn't recognized.""" + return _PATTERN_MAP.get(name) + + +def create_param_pattern(name: str) -> CodePattern: + """Create a runtime ``#param`` pattern that matches the identifier + as a whole word. Raises ``ValueError`` if the name is not a valid + identifier.""" + if not _IDENTIFIER_RE.match(name): + raise ValueError( + f'Invalid parameter identifier: "{name}". Parameter names ' + f"must start with a letter or underscore and contain only " + f"letters, digits, and underscores." + ) + return CodePattern(type="param", pattern=re.compile(rf"\b{re.escape(name)}\b")) diff --git a/tools/python/wgsl_template/errors.py b/tools/python/wgsl_template/errors.py new file mode 100644 index 0000000000000..7403ac8d0821e --- /dev/null +++ b/tools/python/wgsl_template/errors.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Exception types for the WGSL template engine.""" + +from __future__ import annotations + + +class WgslTemplateError(Exception): + """Base class for all WGSL template errors.""" + + def __init__( + self, + message: str, + kind: str = "", + *, + file_path: str | None = None, + line_number: int | None = None, + ) -> None: + super().__init__(message) + self.kind = kind + self.file_path = file_path + self.line_number = line_number + + +class WgslTemplateLoadError(WgslTemplateError): + """Raised by the loader (PASS0).""" + + +class WgslTemplateParseError(WgslTemplateError): + """Raised by the parser (PASS1).""" + + +class WgslTemplateGenerateError(WgslTemplateError): + """Raised by the generator (PASS2).""" + + +class WgslTemplateBuildError(WgslTemplateError): + """Raised by the top-level build orchestrator or code generators.""" diff --git a/tools/python/wgsl_template/generator.py b/tools/python/wgsl_template/generator.py new file mode 100644 index 0000000000000..0f18d743a6d0d --- /dev/null +++ b/tools/python/wgsl_template/generator.py @@ -0,0 +1,782 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""PASS2: pattern matching, directives, and code generation. + +The state machine in :func:`_generate_impl` tracks parentheses, +brackets, a function-call stack, the ``$MAIN`` context, and a separate +output bucket for ``#if``/``#elif`` condition expressions. ``#if`` +expressions go through the same pattern matcher as the body but their +output is buffered and flushed between ``if (`` and ``) {\n``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +from .code_generator.static_cpp import ( + CodeSegment, + CodeSegmentArg, + StaticCodeGenerator, +) +from .code_pattern import ( + DEFAULT_PATTERNS, + CodePattern, + create_param_pattern, + lookup_pattern, +) +from .errors import WgslTemplateGenerateError +from .types import ( + GenerateResult, + ParsedLine, + TemplatePass1, + TemplatePass2, + TemplateRepository, +) + +# ---------------------------------------------------------------------- +# Generator state +# ---------------------------------------------------------------------- + + +@dataclass +class _FunctionCallState: + name: str + parentheses_state: int + params: list[list[CodeSegment]] = field(default_factory=list) + current_param: list[CodeSegment] = field(default_factory=list) + caller: str | None = None + arg_types: list[str] = field(default_factory=list) + + +# preprocessIfStack entry: (type, init_paren, init_bracket, prev_paren, prev_bracket) +_IfStackEntry = list[object] + + +@dataclass +class _GeneratorState: + repo: TemplateRepository + pass1: list[ParsedLine] + code_generator: StaticCodeGenerator + file_path: str + + current_line: int = 0 + current_column: int = 0 + preprocess_if_stack: list[_IfStackEntry] = field(default_factory=list) + patterns: list[CodePattern] = field(default_factory=lambda: list(DEFAULT_PATTERNS)) + current_function_call: list[_FunctionCallState] = field(default_factory=list) + current_parentheses_state: int = 0 + main_function: str = "not-started" # "not-started" | "started" | "ended" + current_bracket_state: int = 0 + result: list[CodeSegment] = field(default_factory=list) + used_params: dict[str, str] = field(default_factory=dict) # name -> param-type + used_variables: dict[str, str] = field(default_factory=dict) # name -> variable-type + + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +def _merge_adjacent_segments(segments: list[CodeSegment]) -> list[CodeSegment]: + """Merge runs of adjacent ``raw`` or ``code`` segments to keep the + final ``emit`` output compact.""" + out: list[CodeSegment] = [] + for seg in segments: + if out and out[-1].type == seg.type and seg.type in ("raw", "code"): + out[-1].content += seg.content + else: + out.append(CodeSegment(type=seg.type, content=seg.content)) + return out + + +def _match_next_pattern( + content: str, patterns: list[CodePattern] +) -> tuple[CodePattern, int, int, re.Match[str] | None] | None: + """Find the pattern that starts earliest in ``content``. + + Returns ``(pattern, start_index, length, match)`` or ``None``. + """ + earliest: tuple[CodePattern, int, int, re.Match] | None = None + for pattern in patterns: + regex = pattern.pattern + if isinstance(regex, str): + regex = re.compile(regex) + m = regex.search(content) + if m is None: + continue + idx = m.start() + if earliest is None or idx < earliest[1]: + earliest = (pattern, idx, len(m.group(0)), m) + return earliest + + +# ---------------------------------------------------------------------- +# Core generator +# ---------------------------------------------------------------------- + + +def _generate_impl(state: _GeneratorState, preserve_code_reference: bool) -> None: + pass1 = state.pass1 + cg = state.code_generator + if_stack = state.preprocess_if_stack + patterns = state.patterns + fn_call = state.current_function_call + + # current_pre_processor_expression: + # None - normal output flow + # list - buffering segments for an #if/#elif condition + pre_processor_expression: list[CodeSegment] | None = None + + def output(typ: str, content: str) -> None: + nonlocal pre_processor_expression + seg = CodeSegment(type=typ, content=content) + if pre_processor_expression is not None: + if typ == "raw": + raise WgslTemplateGenerateError( + f"Raw content inside preprocessor expression at line " + f"{state.current_line + 1}, column {state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + elif typ == "code" and content == "\n": + # End of expression — newline closes the #if line. + pass + else: + seg.type = "raw" # convert to raw inside the expression + pre_processor_expression.append(seg) + return + + if fn_call: + if typ == "raw": + raise WgslTemplateGenerateError( + f"Raw content inside function call at line {state.current_line + 1}, column {state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + fn_call[-1].current_param.append(seg) + else: + state.result.append(seg) + + def process_pattern_match(line: str, restline: str, next_match) -> None: + """Handle one pattern match. Mutates state in place.""" + nonlocal pre_processor_expression + pattern, start_idx, length, match = next_match + + matched = line[state.current_column + start_idx : state.current_column + start_idx + length] + + # Advance past the matched text. + state.current_column += start_idx + length + + caller: str | None = None + + ptype = pattern.type + + if ptype == "control": + if matched == "(": + state.current_parentheses_state += 1 + output("code", "(") + elif matched == ",": + if fn_call and fn_call[-1].parentheses_state + 1 == state.current_parentheses_state: + call = fn_call[-1] + if not call.current_param: + raise WgslTemplateGenerateError( + f"Empty parameter at line {state.current_line + 1}, column {state.current_column}", + "parameter-missing", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + call.params.append(call.current_param) + call.current_param = [] + else: + output("code", ",") + elif matched == ")": + state.current_parentheses_state -= 1 + if state.current_parentheses_state < 0: + raise WgslTemplateGenerateError( + f"Unmatched closing parenthesis at line " + f"{state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if fn_call and fn_call[-1].parentheses_state == state.current_parentheses_state: + call = fn_call.pop() + if call.current_param: + call.params.append(call.current_param) + params = [_merge_adjacent_segments(p) for p in call.params] + # Trim leading/trailing whitespace-only or whitespace + # text from each parameter. + for param in params: + if param and param[-1].type == "code": + if param[-1].content.strip() == "": + param.pop() + else: + param[-1].content = param[-1].content.rstrip() + if param and param[0].type == "code": + if param[0].content.strip() == "": + param.pop(0) + else: + param[0].content = param[0].content.lstrip() + + code_args = [ + CodeSegmentArg( + type=(call.arg_types[i] if i < len(call.arg_types) else "auto"), + code=p, + ) + for i, p in enumerate(params) + ] + + if call.caller: + emitted = cg.method(call.caller, call.name, code_args) + else: + emitted = cg.function(call.name, code_args) + output("expression", emitted) + else: + output("code", ")") + elif matched == "{": + state.current_bracket_state += 1 + output("code", "{") + elif matched == "}": + state.current_bracket_state -= 1 + if state.current_bracket_state < 0: + raise WgslTemplateGenerateError( + f"Unmatched closing bracket at line {state.current_line + 1}, column {state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if state.current_bracket_state == 0 and state.main_function == "started": + output("raw", "MainFunctionEnd();\n") + state.main_function = "ended" + else: + output("code", "}") + elif "$MAIN" in matched: + if state.main_function != "not-started": + raise WgslTemplateGenerateError( + f"Multiple main function start ($MAIN) detected at " + f"line {state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if fn_call: + raise WgslTemplateGenerateError( + f"$MAIN directive inside function call at line " + f"{state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if state.current_parentheses_state != 0: + raise WgslTemplateGenerateError( + f"$MAIN directive inside parentheses at line " + f"{state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if state.current_bracket_state != 0: + raise WgslTemplateGenerateError( + f"$MAIN directive inside brackets at line " + f"{state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if pre_processor_expression is not None: + raise WgslTemplateGenerateError( + f"$MAIN directive inside preprocessor expression " + f"at line {state.current_line + 1}, column " + f"{state.current_column}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + state.main_function = "started" + output("raw", "MainFunctionStart();\n") + state.current_bracket_state = 1 + + elif ptype == "param": + state.used_params[matched] = pattern.param_type or "int" + output("expression", cg.param(matched)) + + elif ptype == "variable": + variable_name = matched + if isinstance(pattern.replace, list) and pattern.replace and pattern.replace[0]: + variable_name = pattern.replace[0] + state.used_variables[variable_name] = pattern.variable_type or "shader-variable" + output("expression", cg.variable(variable_name)) + + elif ptype in ("method", "function"): + # Method-call patterns first extract their receiver, then + # fall through to the same logic as a plain function call. + if ptype == "method": + if isinstance(pattern.replace, list) and pattern.replace and pattern.replace[0]: + caller = pattern.replace[0] + else: + span = match.span(1) + caller = restline[span[0] : span[1]] + state.used_variables[caller] = pattern.variable_type or "shader-variable" + + # Resolve function/method name. + replace_index = 1 if caller else 0 + if ( + isinstance(pattern.replace, list) + and len(pattern.replace) > replace_index + and pattern.replace[replace_index] + ): + name = pattern.replace[replace_index] + else: + group_index = 2 if caller else 1 + span = match.span(group_index) + name = restline[span[0] : span[1]] + + fn_call.append( + _FunctionCallState( + name=name, + parentheses_state=state.current_parentheses_state, + caller=caller, + arg_types=list(pattern.arg_types) if pattern.arg_types else [], + ) + ) + state.current_parentheses_state += 1 + + elif ptype == "property": + if isinstance(pattern.replace, list) and pattern.replace and pattern.replace[0]: + caller = pattern.replace[0] + else: + span = match.span(1) + caller = restline[span[0] : span[1]] + state.used_variables[caller] = pattern.variable_type or "shader-variable" + if isinstance(pattern.replace, list) and len(pattern.replace) > 1 and pattern.replace[1]: + name = pattern.replace[1] + else: + span = match.span(2) + name = restline[span[0] : span[1]] + output("expression", cg.property(caller, name)) + + def process_current_line(line: str) -> None: + """Tokenize and emit ``line`` from ``current_column`` onward.""" + while state.current_column < len(line): + restline = line[state.current_column :] + nxt = _match_next_pattern(restline, patterns) + if nxt is None: + break + start_idx = nxt[1] + if start_idx > 0: + output( + "code", + line[state.current_column : state.current_column + start_idx], + ) + process_pattern_match(line, restline, nxt) + if state.current_column < len(line): + output("code", line[state.current_column :]) + output("code", "\n") + + previous_line_was_empty = True + + for i in range(len(pass1)): + line = pass1[i].line + state.current_line = i + state.current_column = 0 + + if preserve_code_reference: + max_line_number = len(pass1) + line_number_width = len(str(max_line_number)) + padded_line_number = str(state.current_line + 1).rjust(line_number_width) + source_path = pass1[i].code_reference.file_path + source_template = state.repo.templates[source_path] + assert isinstance(source_template, TemplatePass1) + source_line = source_template.raw[pass1[i].code_reference.line_number - 1] + output("raw", f"// {padded_line_number} | {source_line}\n") + + if line == "": + if i == len(pass1) - 1: + continue + if previous_line_was_empty: + continue + previous_line_was_empty = True + else: + previous_line_was_empty = False + + stripped = line.lstrip() + if stripped.startswith("#"): + if stripped.startswith("#use "): + uses = [p for p in stripped[5:].split(" ") if p.strip()] + for use in uses: + pat = lookup_pattern(use) + if pat is None: + raise WgslTemplateGenerateError( + f"Unknown use: {use} at line {state.current_line + 1}", + "code-pattern-not-found", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + patterns.append(pat) + elif stripped.startswith("#param "): + params = [p for p in stripped[7:].split(" ") if p.strip()] + if not params: + raise WgslTemplateGenerateError( + f"No parameters specified in #param directive at line {state.current_line + 1}", + "parameter-missing", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + for param in params: + pat = create_param_pattern(param) + # Allow duplicates if same; error on type mismatch. + pat_str = pat.pattern.pattern if hasattr(pat.pattern, "pattern") else str(pat.pattern) + existing = next( + ( + p + for p in patterns + if p.type == "param" + and (p.pattern.pattern if hasattr(p.pattern, "pattern") else str(p.pattern)) == pat_str + ), + None, + ) + if existing is None: + patterns.append(pat) + else: + existing_type = existing.param_type or "int" + new_type = pat.param_type or "int" + if existing_type != new_type: + raise WgslTemplateGenerateError( + f"Duplicate param with different type: {param} at line {state.current_line + 1}", + "parameter-type-mismatch", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + elif stripped.startswith("#if "): + if fn_call: + raise WgslTemplateGenerateError( + f"Preprocessor directive inside function call at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + condition = stripped[4:].strip() + if not condition: + raise WgslTemplateGenerateError( + f"Empty condition in #if directive at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + leading_ws = len(line) - len(stripped) + state.current_column = leading_ws + 4 + if_stack.append( + [ + "if", + state.current_parentheses_state, + state.current_bracket_state, + None, + None, + ] + ) + output("raw", "if (") + pre_processor_expression = [] + cached_paren = state.current_parentheses_state + state.current_parentheses_state = 0 + process_current_line(line) + if state.current_parentheses_state != 0: + raise WgslTemplateGenerateError( + f"Unmatched parentheses in preprocessor expression at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + state.current_parentheses_state = cached_paren + state.result.extend(pre_processor_expression) + pre_processor_expression = None + output("raw", ") {\n") + elif stripped.startswith("#elif "): + if not if_stack or if_stack[-1][0] not in ("if", "elif"): + raise WgslTemplateGenerateError( + f"#elif mismatch at line {state.current_line + 1}", + "code-generation-failed", + line_number=state.current_line + 1, + ) + if fn_call: + raise WgslTemplateGenerateError( + f"Preprocessor directive inside function call at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + prev_paren = if_stack[-1][3] + if prev_paren is not None and state.current_parentheses_state != prev_paren: + raise WgslTemplateGenerateError( + f"Parentheses state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_paren}, got {state.current_parentheses_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + prev_bracket = if_stack[-1][4] + if prev_bracket is not None and state.current_bracket_state != prev_bracket: + raise WgslTemplateGenerateError( + f"Bracket state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_bracket}, got {state.current_bracket_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if_stack[-1][3] = state.current_parentheses_state + if_stack[-1][4] = state.current_bracket_state + state.current_parentheses_state = if_stack[-1][1] + state.current_bracket_state = if_stack[-1][2] + + condition = stripped[6:].strip() + if not condition: + raise WgslTemplateGenerateError( + f"Empty condition in #elif directive at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + leading_ws = len(line) - len(stripped) + state.current_column = leading_ws + 6 + if_stack[-1][0] = "elif" + output("raw", "} else if (") + pre_processor_expression = [] + cached_paren = state.current_parentheses_state + state.current_parentheses_state = 0 + process_current_line(line) + if state.current_parentheses_state != 0: + raise WgslTemplateGenerateError( + f"Unmatched parentheses in preprocessor expression at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + state.current_parentheses_state = cached_paren + state.result.extend(pre_processor_expression) + pre_processor_expression = None + output("raw", ") {\n") + elif stripped.startswith("#else"): + if not if_stack or if_stack[-1][0] not in ("if", "elif"): + raise WgslTemplateGenerateError( + f"#else mismatch at line {state.current_line + 1}", + "code-generation-failed", + line_number=state.current_line + 1, + ) + if fn_call: + raise WgslTemplateGenerateError( + f"Preprocessor directive inside function call at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + prev_paren = if_stack[-1][3] + if prev_paren is not None and state.current_parentheses_state != prev_paren: + raise WgslTemplateGenerateError( + f"Parentheses state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_paren}, got {state.current_parentheses_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + prev_bracket = if_stack[-1][4] + if prev_bracket is not None and state.current_bracket_state != prev_bracket: + raise WgslTemplateGenerateError( + f"Bracket state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_bracket}, got {state.current_bracket_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if_stack[-1][3] = state.current_parentheses_state + if_stack[-1][4] = state.current_bracket_state + state.current_parentheses_state = if_stack[-1][1] + state.current_bracket_state = if_stack[-1][2] + + if stripped[5:].strip() != "": + raise WgslTemplateGenerateError( + f"Unexpected content after #else at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + if_stack[-1][0] = "else" + output("raw", "} else {\n") + elif stripped.startswith("#endif"): + if not if_stack: + raise WgslTemplateGenerateError( + f"#endif mismatch at line {state.current_line + 1}", + "code-generation-failed", + line_number=state.current_line + 1, + ) + if fn_call: + raise WgslTemplateGenerateError( + f"Preprocessor directive inside function call at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + prev_paren = if_stack[-1][3] + if prev_paren is not None and state.current_parentheses_state != prev_paren: + raise WgslTemplateGenerateError( + f"Parentheses state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_paren}, got {state.current_parentheses_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + prev_bracket = if_stack[-1][4] + if prev_bracket is not None and state.current_bracket_state != prev_bracket: + raise WgslTemplateGenerateError( + f"Bracket state mismatch in #elif directive at " + f"line {state.current_line + 1}, expected " + f"{prev_bracket}, got {state.current_bracket_state}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + + if stripped[6:].strip() != "": + raise WgslTemplateGenerateError( + f"Unexpected content after #endif at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + output("raw", "}\n") + if_stack.pop() + else: + if stripped in ("#use", "#param", "#if", "#elif"): + raise WgslTemplateGenerateError( + f"Missing content after preprocessor directive at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + raise WgslTemplateGenerateError( + f"Unknown preprocessor directive: {stripped} at line {state.current_line + 1}", + "code-generation-failed", + file_path=state.file_path, + line_number=state.current_line + 1, + ) + previous_line_was_empty = True + else: + process_current_line(line) + + +# ---------------------------------------------------------------------- +# Public API +# ---------------------------------------------------------------------- + + +def generate( + file_path: str, + repo: TemplateRepository, + code_generator: StaticCodeGenerator, + *, + preserve_code_reference: bool = False, +) -> GenerateResult: + """Run PASS2 on a single template, returning the per-template result.""" + + template = repo.templates.get(file_path) + if template is None or not isinstance(template, TemplatePass1): + raise WgslTemplateGenerateError( + f"Template not found for file: {file_path}", + "generator-not-found", + file_path=file_path, + ) + + state = _GeneratorState( + repo=repo, + pass1=list(template.pass1), + code_generator=code_generator, + file_path=file_path, + ) + + _generate_impl(state, preserve_code_reference) + + if state.preprocess_if_stack: + kinds = ", ".join(str(entry[0]) for entry in state.preprocess_if_stack) + raise WgslTemplateGenerateError( + f"Unclosed preprocessor directive: {kinds}", + "code-generation-failed", + ) + if state.current_function_call: + names = ", ".join(c.name for c in state.current_function_call) + raise WgslTemplateGenerateError( + f"Unclosed function call: {names}", + "code-generation-failed", + ) + if state.current_parentheses_state != 0: + raise WgslTemplateGenerateError( + "Unmatched parentheses at the end of processing", + "code-generation-failed", + ) + + if state.main_function == "started": + raise WgslTemplateGenerateError( + "Main function context started but not ended at the end of processing", + "code-generation-failed", + ) + if state.current_bracket_state != 0: + raise WgslTemplateGenerateError( + "Unmatched brackets at the end of processing", + "code-generation-failed", + ) + + state.result = _merge_adjacent_segments(state.result) + + sorted_params = {k: state.used_params[k] for k in sorted(state.used_params)} + sorted_variables = {k: state.used_variables[k] for k in sorted(state.used_variables)} + + return GenerateResult( + code=code_generator.emit(state.result), + params=sorted_params, + variables=sorted_variables, + has_main_function=(state.main_function == "ended"), + ) + + +def generate_directory( + repo: TemplateRepository, + code_generator: StaticCodeGenerator, + *, + preserve_code_reference: bool = False, +) -> TemplateRepository: + """Run PASS2 on every template in the repository.""" + + out_templates: dict[str, object] = {} + for file_path in repo.templates: + result = generate( + file_path, + repo, + code_generator, + preserve_code_reference=preserve_code_reference, + ) + template = repo.templates[file_path] + assert isinstance(template, TemplatePass1) + out_templates[file_path] = TemplatePass2( + file_path=template.file_path, + generate_result=result, + ) + + return TemplateRepository( + base_path=repo.base_path, + templates=out_templates, + ) diff --git a/tools/python/wgsl_template/loader.py b/tools/python/wgsl_template/loader.py new file mode 100644 index 0000000000000..d11d98aebce7b --- /dev/null +++ b/tools/python/wgsl_template/loader.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""PASS0: load template files from one or more source directories.""" + +from __future__ import annotations + +import os +import re +from collections.abc import Sequence +from pathlib import Path +from typing import Union + +from .errors import WgslTemplateLoadError +from .types import SourceDir, TemplatePass0, TemplateRepository + +DEFAULT_EXT = ".wgsl.template" + +_LINE_SPLIT_RE = re.compile(r"\r?\n") + + +SourceDirInput = Union[str, "os.PathLike[str]", SourceDir] + + +def _normalize_source_dir(item: SourceDirInput) -> SourceDir: + if isinstance(item, SourceDir): + return item + return SourceDir(path=os.fspath(item), alias=None) + + +def _resolve_within(base: Path, target: Path) -> bool: + """True iff ``target`` is the same as ``base`` or lives underneath it. + + Used as a defense against symlinks or crafted paths that would + escape the source directory. + """ + + try: + target.relative_to(base) + return True + except ValueError: + return False + + +def _load_directory( + source_dir: Path, + base_dir: Path, + ext: str, + alias: str | None, + templates: dict, +) -> None: + """Recursively scan ``source_dir`` for files ending in ``ext``.""" + + try: + entries = sorted(os.listdir(source_dir)) + except OSError as e: + raise WgslTemplateLoadError( + f"Error scanning directory {source_dir}: {e}", + "scan-directory", + ) from e + + for entry in entries: + full_path = source_dir / entry + + # Reject anything resolving outside base_dir. + try: + resolved = full_path.resolve() + except OSError: + # Broken symlink etc. + continue + if not _resolve_within(base_dir, resolved): + continue + + # Skip symlinks and special files; only follow real directories + # and read real files. + if full_path.is_symlink(): + continue + if full_path.is_dir(): + _load_directory(full_path, base_dir, ext, alias, templates) + elif full_path.is_file() and entry.endswith(ext): + _load_file(full_path, base_dir, alias, templates) + + +def _load_file( + file_path: Path, + base_dir: Path, + alias: str | None, + templates: dict, +) -> None: + resolved_file = file_path.resolve() + if not _resolve_within(base_dir, resolved_file): + raise WgslTemplateLoadError( + f"Security violation: attempted to read file outside base directory: {file_path}", + "read-file", + file_path=str(file_path), + ) + + try: + content = file_path.read_text(encoding="utf-8") + except OSError as e: + raise WgslTemplateLoadError( + f"Error loading template file {file_path}: {e}", + "read-file", + file_path=str(file_path), + ) from e + + lines = _LINE_SPLIT_RE.split(content) + + relative = file_path.relative_to(base_dir) + # Always use POSIX-style paths in template names. + template_name = relative.as_posix() + if alias: + template_name = f"{alias}/{template_name}" + + if template_name in templates: + raise WgslTemplateLoadError( + f"Template name conflict: {template_name} already exists", + "template-conflict", + file_path=str(file_path), + ) + + templates[template_name] = TemplatePass0( + file_path=str(resolved_file), + raw=lines, + ) + + +def load_from_directories( + directories: Sequence[SourceDirInput], + *, + ext: str = DEFAULT_EXT, +) -> TemplateRepository: + """Load all template files matching ``ext`` from each directory. + + Returns a :class:`TemplateRepository` whose ``templates`` map is + sorted by template name. + """ + + if not directories: + raise WgslTemplateLoadError( + "At least one source directory is required", + "scan-directory", + ) + + raw_templates: dict = {} + resolved_base_paths: list[Path] = [] + + for raw_item in directories: + spec = _normalize_source_dir(raw_item) + dir_path = Path(spec.path) + + # Resolve relative paths against the cwd. + if not dir_path.is_absolute(): + dir_path = Path.cwd() / dir_path + + if not dir_path.exists(): + raise WgslTemplateLoadError( + f"Cannot access directory {dir_path}: not found", + "scan-directory", + ) + if not dir_path.is_dir(): + raise WgslTemplateLoadError( + f"Path {dir_path} is not a directory", + "scan-directory", + ) + + resolved_base = dir_path.resolve() + _load_directory( + source_dir=resolved_base, + base_dir=resolved_base, + ext=ext, + alias=spec.alias, + templates=raw_templates, + ) + resolved_base_paths.append(resolved_base) + + # Sort by template name for stable, host-independent ordering. + sorted_templates = {name: raw_templates[name] for name in sorted(raw_templates)} + + # Use the first input directory as the repository's base path. + base_path = str(resolved_base_paths[0]) if resolved_base_paths else "" + + return TemplateRepository( + base_path=base_path, + templates=sorted_templates, + ) + + +def load_from_directory( + directory: SourceDirInput, + *, + ext: str = DEFAULT_EXT, +) -> TemplateRepository: + """Convenience wrapper for the single-directory case.""" + return load_from_directories([directory], ext=ext) diff --git a/tools/python/wgsl_template/parser.py b/tools/python/wgsl_template/parser.py new file mode 100644 index 0000000000000..3fe0ea800a934 --- /dev/null +++ b/tools/python/wgsl_template/parser.py @@ -0,0 +1,352 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""PASS1: comment stripping, #include expansion, #define substitution. + +Three sub-passes run in order on each top-level template: + +1. :func:`_strip_comments` — replace ``//`` and ``/* */`` regions with + empty text while preserving line count (every line is kept, even if + it becomes empty). Not string-aware. +2. :func:`_expand_includes` — flatten ``#include "name"`` directives, + detect cycles (incl. self-include), reject missing files. +3. :func:`_apply_macros` — process ``#define`` directives. Expansion + is eager at definition time, whole-word match, and rejects + self-reference / mutual circular / duplicate / empty value. +""" + +from __future__ import annotations + +import re + +from .errors import WgslTemplateParseError +from .types import ( + CodeReference, + ParsedLine, + TemplatePass0, + TemplatePass1, + TemplateRepository, +) + +_INCLUDE_RE = re.compile(r"^\s*#include\s+(.+)$") +_DEFINE_FULL_RE = re.compile(r"^\s*#define\s+([a-zA-Z_][a-zA-Z0-9_]*)\s+(.+)$") +_DEFINE_EMPTY_RE = re.compile(r"^\s*#define\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*$") +_DEFINE_INVALID_NAME_RE = re.compile(r"^\s*#define\s+(\S+)(?:\s+(.+))?$") + + +# ---------------------------------------------------------------------- +# Sub-pass 1: comment stripping +# ---------------------------------------------------------------------- + + +def _strip_comments(file_path: str, raw: list[str]) -> list[ParsedLine]: + """Strip ``//`` and ``/* */`` comments, returning one + :class:`ParsedLine` per input line. + + Each line's :class:`CodeReference` line number is bound to its index + in ``raw`` (the original source), so errors and ``--preserve-code-ref`` + output point at the true source line regardless of comment removal. + + NOT string-aware. WGSL has no string type, but template arguments + like ``getElementAt("a/*b", "c")`` would be mis-stripped. + """ + + out: list[ParsedLine] = [] + in_multi = False + + for source_index, line in enumerate(raw): + processed: list[str] = [] + i = 0 + n = len(line) + + while i < n: + if in_multi: + # Look for end of multi-line comment. + if i < n - 1 and line[i] == "*" and line[i + 1] == "/": + in_multi = False + i += 2 + else: + i += 1 + else: + # Single-line comment: rest of line is a comment. + if i < n - 1 and line[i] == "/" and line[i + 1] == "/": + break + # Multi-line comment start. + if i < n - 1 and line[i] == "/" and line[i + 1] == "*": + in_multi = True + i += 2 + else: + processed.append(line[i]) + i += 1 + + # Always append (even if empty) to preserve line count, and tag + # each line with its true source line number (1-based). + out.append( + ParsedLine( + line="".join(processed).rstrip(), + code_reference=CodeReference( + file_path=file_path, + line_number=source_index + 1, + ), + ) + ) + + if in_multi: + raise WgslTemplateParseError( + "Unterminated multi-line comment detected in template", + "comment-removal", + file_path=file_path, + ) + + return out + + +# ---------------------------------------------------------------------- +# Sub-pass 2: #include expansion +# ---------------------------------------------------------------------- + + +class _ParseEntry: + __slots__ = ("include_processed", "lines") + + def __init__(self, lines: list[ParsedLine]) -> None: + self.lines: list[ParsedLine] = lines + self.include_processed: bool = False + + +def _expand_includes( + parse_state: dict[str, _ParseEntry], + include_stack: list[str], +) -> None: + """Recursively expand ``#include`` directives in place. + + Mutates ``parse_state[include_stack[-1]].lines`` to contain the + flattened content. Sets ``include_processed`` so repeat includes + are no-ops. + """ + + current_file = include_stack[-1] + current = parse_state.get(current_file) + if current is None: + raise WgslTemplateParseError( + f'File "{current_file}" not found in parse state', + "include-resolution", + file_path=current_file, + ) + + if current.include_processed: + return + + flattened: list[ParsedLine] = [] + for line_index, parsed_line in enumerate(current.lines): + line = parsed_line.line + m = _INCLUDE_RE.match(line) + if m is None: + flattened.append(parsed_line) + continue + + include_param = m.group(1).strip() + if not (include_param.startswith('"') and include_param.endswith('"')): + raise WgslTemplateParseError( + f"Invalid #include directive in file {current_file} at line " + f"{line_index + 1}: file path must be enclosed in double quotes", + "syntax-error", + file_path=current_file, + line_number=line_index + 1, + ) + + include_path = include_param[1:-1] + + if include_path in include_stack: + raise WgslTemplateParseError( + f"Circular #include detected in file {current_file} at line " + f"{line_index + 1}: {include_path} is already included", + "include-circular", + file_path=current_file, + line_number=line_index + 1, + ) + if include_path not in parse_state: + raise WgslTemplateParseError( + f'File "{include_path}" not found in parse state for ' + f'#include directive in file "{current_file}" at line ' + f"{line_index + 1}", + "include-not-found", + file_path=current_file, + line_number=line_index + 1, + ) + + include_stack.append(include_path) + _expand_includes(parse_state, include_stack) + flattened.extend(parse_state[include_path].lines) + include_stack.pop() + + current.include_processed = True + current.lines = flattened + + +# ---------------------------------------------------------------------- +# Sub-pass 3: #define substitution +# ---------------------------------------------------------------------- + + +def _apply_macros(lines: list[ParsedLine], file_name: str) -> list[ParsedLine]: + """Process ``#define`` directives and apply substitutions. + + Rules: + * File-scoped per top-level template (after include flattening). + * Eager expansion at definition time: ``#define X Y`` stores the + already-expanded ``Y``. Forward references do not work. + * Whole-word match (``\\bNAME\\b``). + * Reject self-reference, mutual circular dependencies, + duplicate definitions, empty / whitespace-only values. + """ + + macros: dict[str, str] = {} + out: list[ParsedLine] = [] + + for line_index, parsed_line in enumerate(lines): + line = parsed_line.line + + if line.lstrip().startswith("#define "): + full = _DEFINE_FULL_RE.match(line) + if full is None: + # Try to give a more specific error message. + empty = _DEFINE_EMPTY_RE.match(line) + if empty is not None: + raise WgslTemplateParseError( + f"Invalid macro definition in file {file_name} at " + f'line {line_index + 1}: macro "{empty.group(1)}" ' + f"has no value", + "syntax-error", + file_path=file_name, + line_number=line_index + 1, + ) + + invalid = _DEFINE_INVALID_NAME_RE.match(line) + if invalid is not None: + raise WgslTemplateParseError( + f"Invalid macro definition in file {file_name} at " + f"line {line_index + 1}: invalid macro name " + f'"{invalid.group(1)}" (must start with letter or ' + f"underscore, contain only letters, numbers, and " + f"underscores)", + "syntax-error", + file_path=file_name, + line_number=line_index + 1, + ) + + raise WgslTemplateParseError( + f"Invalid macro definition in file {file_name} at line " + f"{line_index + 1}: malformed #define directive", + "syntax-error", + file_path=file_name, + line_number=line_index + 1, + ) + + macro_name = full.group(1) + macro_value = full.group(2).strip() + + if macro_value == "": + raise WgslTemplateParseError( + f"Invalid macro definition in file {file_name} at line " + f'{line_index + 1}: macro "{macro_name}" has no value', + "syntax-error", + file_path=file_name, + line_number=line_index + 1, + ) + + if macro_name in macros: + raise WgslTemplateParseError( + f"Duplicate macro definition in file {file_name} at line " + f'{line_index + 1}: "{macro_name}" is already defined', + "define-expansion", + file_path=file_name, + line_number=line_index + 1, + ) + + # Direct self-reference check. + self_ref = re.compile(rf"\b{re.escape(macro_name)}\b") + if self_ref.search(macro_value): + raise WgslTemplateParseError( + f"Circular macro reference in file {file_name} at line " + f'{line_index + 1}: macro "{macro_name}" references ' + f"itself", + "define-expansion", + file_path=file_name, + line_number=line_index + 1, + ) + + # Apply existing macros to the new value (eager expansion). + expanded_value = macro_value + for existing_name, existing_value in macros.items(): + regex = re.compile(rf"\b{re.escape(existing_name)}\b") + if regex.search(expanded_value): + # Mutual circular check: existing macro references + # the new one. + circ = re.compile(rf"\b{re.escape(macro_name)}\b") + if circ.search(existing_value): + raise WgslTemplateParseError( + f"Circular macro reference in file {file_name} " + f"at line {line_index + 1}: macro " + f'"{macro_name}" creates circular dependency ' + f'with "{existing_name}"', + "define-expansion", + file_path=file_name, + line_number=line_index + 1, + ) + expanded_value = regex.sub(existing_value, expanded_value) + + macros[macro_name] = expanded_value + line = "" # Clear the #define line in the output. + else: + # Apply macro substitutions to the current line. + for macro_name, macro_value in macros.items(): + regex = re.compile(rf"\b{re.escape(macro_name)}\b") + line = regex.sub(macro_value, line) + + out.append(ParsedLine(line=line, code_reference=parsed_line.code_reference)) + + return out + + +# ---------------------------------------------------------------------- +# Top-level entry point +# ---------------------------------------------------------------------- + + +def parse(repo: TemplateRepository) -> TemplateRepository: + """Run PASS1 on every template in the repository.""" + + # STEP 1 — strip comments. _strip_comments returns ParsedLine + # objects already tagged with their true source line numbers. + parse_state: dict[str, _ParseEntry] = {} + for template_key, template in repo.templates.items(): + assert isinstance(template, TemplatePass0) + parsed_lines = _strip_comments(template_key, template.raw) + parse_state[template_key] = _ParseEntry(parsed_lines) + + # STEP 2 — expand #include directives in every top-level template. + pass1_templates: dict[str, object] = {} + for template_key, template in repo.templates.items(): + assert isinstance(template, TemplatePass0) + _expand_includes(parse_state, [template_key]) + pass1_templates[template_key] = TemplatePass1( + file_path=template.file_path, + raw=template.raw, + pass1=parse_state[template_key].lines, + ) + + # STEP 3 — apply #define substitutions on the post-include result. + for template_key in list(pass1_templates): + template = pass1_templates[template_key] + assert isinstance(template, TemplatePass1) + processed = _apply_macros(list(template.pass1), template_key) + pass1_templates[template_key] = TemplatePass1( + file_path=template.file_path, + raw=template.raw, + pass1=processed, + ) + + return TemplateRepository( + base_path=repo.base_path, + templates=pass1_templates, + ) diff --git a/tools/python/wgsl_template/test/golden_compare.py b/tools/python/wgsl_template/test/golden_compare.py new file mode 100644 index 0000000000000..16297d9297334 --- /dev/null +++ b/tools/python/wgsl_template/test/golden_compare.py @@ -0,0 +1,61 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Shared helpers for comparing generated WGSL output against goldens. + +Some details of the generated file set are host-dependent and NOT part +of the equivalence contract (see DESIGN_wgsl_python_port.md, +"Determinism"): ``__str_N`` ids, the ordering of constants / +declarations / includes in the aggregate index/table files, and the +per-file ``// `` markers. Only the shader string *contents* and +the C++ structure are contractual. :func:`canonicalize` normalizes the +non-contractual details away so goldens stay stable across tools and +hosts, while still comparing shader strings exactly (via their resolved +``__str_N`` values). +""" + +from __future__ import annotations + +import os +import re +from pathlib import Path + +# Aggregate files whose line *order* is not semantically meaningful; +# compared as unordered line sets. Per-template generated/*.h files keep +# their in-order comparison. +_ORDER_INSENSITIVE_FILES = {"index.h", "index_impl.h", "string_table.h"} + +_STR_ID_RE = re.compile(r"__str_(\d+)") +_STR_TABLE_DEF_RE = re.compile(r"__str_(\d+)\s*=\s*(.*);") +_SHA256_COMMENT_RE = re.compile(r" *// [0-9a-f]{64}") + + +def read_tree(root: Path) -> dict[str, str]: + """Read every file under ``root`` into a ``{posix_relpath: text}`` + map, normalizing CRLF that a Windows checkout may have introduced.""" + out: dict[str, str] = {} + for dirpath, _dirs, files in os.walk(root): + for name in files: + full = Path(dirpath) / name + rel = full.relative_to(root).as_posix() + out[rel] = full.read_bytes().replace(b"\r\n", b"\n").decode("utf-8") + return out + + +def canonicalize(files: dict[str, str]) -> dict[str, str]: + """Reduce a generated file set to a host-independent canonical form. + + Resolves ``__str_N`` references to their string literals, drops the + derived sha256 markers, and sorts the order-insensitive aggregate + files. + """ + string_table = files.get("string_table.h", "") + id_to_literal = {m.group(1): m.group(2) for m in _STR_TABLE_DEF_RE.finditer(string_table)} + + canonical: dict[str, str] = {} + for rel, text in files.items(): + resolved = _STR_ID_RE.sub(lambda m: id_to_literal.get(m.group(1), m.group(0)), text) + resolved = _SHA256_COMMENT_RE.sub("", resolved) + if rel.rsplit("/", 1)[-1] in _ORDER_INSENSITIVE_FILES: + resolved = "\n".join(sorted(resolved.split("\n"))) + canonical[rel] = resolved + return canonical diff --git a/tools/python/wgsl_template/test/in_tree_golden/.gitattributes b/tools/python/wgsl_template/test/in_tree_golden/.gitattributes new file mode 100644 index 0000000000000..339c7de8b3519 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/.gitattributes @@ -0,0 +1,3 @@ +# Generated golden snapshots are written with LF; keep them LF on checkout +# so they match the tool's output and produce clean cross-platform diffs. +* text eol=lf diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/im2col_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/im2col_matmul.h new file mode 100644 index 0000000000000..739ca2b734954 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/im2col_matmul.h @@ -0,0 +1,323 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: nn/im2col_matmul.wgsl.template + +template <> +Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_has_bias = params.param_has_bias; + auto& __param_tile_m = params.param_tile_m; + auto& __param_tile_n = params.param_tile_n; + auto& __param_use_subgroup = params.param_use_subgroup; + auto& __param_vec_size = params.param_vec_size; + + // Extract variables + auto& __var_output = *params.var_output; + auto& __var_src = *params.var_src; + auto& __var_weight = *params.var_weight; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | #param has_bias +// 5 | #param tile_m +// 6 | #param tile_n +// 7 | #param use_subgroup +// 8 | #param vec_size +// 9 | +// 10 | #use .getByOffset .setByOffset +// 11 | +// 12 | // im2col access for src: [N, H_i, W_i, C_i / vec_size] +// 13 | // Conceptual Matrix Shape: N * (H_o * W_o) x (K_h * K_w * C_i / vec_size) +// 14 | fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t { +ss << "fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t {\n"; +// 15 | if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * vec_size >= uniforms.im2col_k) { +ss << " if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * "; +ss << __param_vec_size; +ss << " >= uniforms.im2col_k) {\n"; +// 16 | return src_value_t(); +ss << " return src_value_t();\n"; +// 17 | } +ss << " }\n"; +// 18 | +ss << "\n"; +// 19 | let channel_i_vec = uniforms.channel_i / vec_size; +ss << " let channel_i_vec = uniforms.channel_i / "; +ss << __param_vec_size; +ss << ";\n"; +// 20 | +ss << "\n"; +// 21 | // 1. Decompose M index (H_o * W_o) into (h_idx, w_idx) +// 22 | let h_idx = m / uniforms.output_w; // Output H index (H_o) +ss << " let h_idx = m / uniforms.output_w;\n"; +// 23 | let w_idx = m % uniforms.output_w; // Output W index (W_o) +ss << " let w_idx = m % uniforms.output_w;\n"; +// 24 | +ss << "\n"; +// 25 | // 2. Decompose K index into (k_h, k_w, c_i_vec_idx) +// 26 | let c_i_vec_idx = k_packed_idx % channel_i_vec; +ss << " let c_i_vec_idx = k_packed_idx % channel_i_vec;\n"; +// 27 | let k_h_w_idx = k_packed_idx / channel_i_vec; +ss << " let k_h_w_idx = k_packed_idx / channel_i_vec;\n"; +// 28 | let k_h = k_h_w_idx / uniforms.kernel_w; // Kernel Row +ss << " let k_h = k_h_w_idx / uniforms.kernel_w;\n"; +// 29 | let k_w = k_h_w_idx % uniforms.kernel_w; // Kernel Column +ss << " let k_w = k_h_w_idx % uniforms.kernel_w;\n"; +// 30 | +ss << "\n"; +// 31 | // 3. Calculate the coordinate in the padded input tensor +// 32 | let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x; +ss << " let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x;\n"; +// 33 | let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y; +ss << " let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y;\n"; +// 34 | +ss << "\n"; +// 35 | // 4. Calculate the coordinate in the original input tensor +// 36 | let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x); +ss << " let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x);\n"; +// 37 | let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y); +ss << " let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y);\n"; +// 38 | +ss << "\n"; +// 39 | // 5. Check for padding/out-of-bounds +// 40 | if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) || +ss << " if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) ||\n"; +// 41 | src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) { +ss << " src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) {\n"; +// 42 | return src_value_t(); +ss << " return src_value_t();\n"; +// 43 | } +ss << " }\n"; +// 44 | +ss << "\n"; +// 45 | // 6. Calculate final NHWC index +// 46 | let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec + +ss << " let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec +\n"; +// 47 | u32(src_h_coord) * uniforms.src_w * channel_i_vec + +ss << " u32(src_h_coord) * uniforms.src_w * channel_i_vec +\n"; +// 48 | u32(src_w_coord) * channel_i_vec + +ss << " u32(src_w_coord) * channel_i_vec +\n"; +// 49 | c_i_vec_idx; +ss << " c_i_vec_idx;\n"; +// 50 | return src.getByOffset(src_idx); +ss << " return "; +ss << __var_src.GetByOffset("src_idx"); +ss << ";\n"; +// 51 | } +ss << "}\n"; +// 52 | +ss << "\n"; +// 53 | // weight shape: [Co, K_h, K_w, C_i / vec_size] (CoHWCi) +// 54 | fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t { +ss << "fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t {\n"; +// 55 | if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / vec_size) { +ss << " if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / "; +ss << __param_vec_size; +ss << ") {\n"; +// 56 | let weight_idx = n * uniforms.im2col_k / vec_size + +ss << " let weight_idx = n * uniforms.im2col_k / "; +ss << __param_vec_size; +ss << " +\n"; +// 57 | k_packed_idx; +ss << " k_packed_idx;\n"; +// 58 | return weight.getByOffset(weight_idx); +ss << " return "; +ss << __var_weight.GetByOffset("weight_idx"); +ss << ";\n"; +// 59 | } +ss << " }\n"; +// 60 | return weight_value_t(); +ss << " return weight_value_t();\n"; +// 61 | } +ss << "}\n"; +// 62 | +ss << "\n"; +// 63 | fn load_bias(n : u32) -> output_element_t { +ss << "fn load_bias(n : u32) -> output_element_t {\n"; +// 64 | #if has_bias +if (__param_has_bias) { +// 65 | if (n < uniforms.im2col_n) { +ss << " if (n < uniforms.im2col_n) {\n"; +// 66 | return output_element_t(bias[n]); +ss << " return output_element_t(bias[n]);\n"; +// 67 | } +ss << " }\n"; +// 68 | #endif +} +// 69 | return output_element_t(); +ss << " return output_element_t();\n"; +// 70 | } +ss << "}\n"; +// 71 | +ss << "\n"; +// 72 | // output shape: [N, H_o, W_o, C_o] (NHWC) +// 73 | fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) { +ss << "fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) {\n"; +// 74 | if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) { +ss << " if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) {\n"; +// 75 | let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n + +ss << " let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n +\n"; +// 76 | m * uniforms.im2col_n + +ss << " m * uniforms.im2col_n +\n"; +// 77 | n; +ss << " n;\n"; +// 78 | output.setByOffset(output_idx, value); +ss << " "; +ss << __var_output.SetByOffset("output_idx", "value"); +ss << ";\n"; +// 79 | } +ss << " }\n"; +// 80 | } +ss << "}\n"; +// 81 | +ss << "\n"; +// 82 | const TILE_M_SIZE : u32 = tile_m; +ss << "const TILE_M_SIZE : u32 = "; +ss << __param_tile_m; +ss << ";\n"; +// 83 | const TILE_N_SIZE : u32 = tile_n; +ss << "const TILE_N_SIZE : u32 = "; +ss << __param_tile_n; +ss << ";\n"; +// 84 | // In dimension K, the tile consists of 16 scalars, requiring `16 / vec_size` vector loads. +ss << "\n"; +// 85 | const TILE_K_VEC_SIZE : u32 = 16 / vec_size; +ss << "const TILE_K_VEC_SIZE : u32 = 16 / "; +ss << __param_vec_size; +ss << ";\n"; +// 86 | // In dimensions M and N, since a workgroup has 64 threads, it advances by `64 / TILE_K_VEC_SIZE`. +ss << "\n"; +// 87 | const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE; +ss << "const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE;\n"; +// 88 | +ss << "\n"; +// 89 | var src_tile : array, TILE_K_VEC_SIZE>; +ss << "var src_tile : array, TILE_K_VEC_SIZE>;\n"; +// 90 | var weight_tile : array, TILE_K_VEC_SIZE>; +ss << "var weight_tile : array, TILE_K_VEC_SIZE>;\n"; +// 91 | +ss << "\n"; +// 92 | $MAIN { +MainFunctionStart(); +ss << "\n"; +// 93 | let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles); +ss << " let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles);\n"; +// 94 | let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE; +ss << " let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE;\n"; +// 95 | let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE; +ss << " let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE;\n"; +// 96 | +ss << "\n"; +// 97 | var results : array; +ss << " var results : array;\n"; +// 98 | for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) { +ss << " for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) {\n"; +// 99 | for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) { +ss << " for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) {\n"; +// 100 | // Loads a 64 vec of src into the workgroup memory. +ss << "\n"; +// 101 | let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE; +ss << " let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE;\n"; +// 102 | let load_src_k = local_idx % TILE_K_VEC_SIZE; +ss << " let load_src_k = local_idx % TILE_K_VEC_SIZE;\n"; +// 103 | +ss << "\n"; +// 104 | src_tile[load_src_k][load_src_m] = load_src(batch, +ss << " src_tile[load_src_k][load_src_m] = load_src(batch,\n"; +// 105 | m_global_base + load_src_m, +ss << " m_global_base + load_src_m,\n"; +// 106 | k_idx * TILE_K_VEC_SIZE + load_src_k); +ss << " k_idx * TILE_K_VEC_SIZE + load_src_k);\n"; +// 107 | } +ss << " }\n"; +// 108 | +ss << "\n"; +// 109 | for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) { +ss << " for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) {\n"; +// 110 | // Loads a 64 vec of weight into the workgroup memory. +ss << "\n"; +// 111 | let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE; +ss << " let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE;\n"; +// 112 | let load_weight_k = local_idx % TILE_K_VEC_SIZE; +ss << " let load_weight_k = local_idx % TILE_K_VEC_SIZE;\n"; +// 113 | +ss << "\n"; +// 114 | weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n, +ss << " weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n,\n"; +// 115 | k_idx * TILE_K_VEC_SIZE + load_weight_k); +ss << " k_idx * TILE_K_VEC_SIZE + load_weight_k);\n"; +// 116 | } +ss << " }\n"; +// 117 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 118 | +ss << "\n"; +// 119 | for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) { +ss << " for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) {\n"; +// 120 | let weight_data = weight_tile[inner_k_idx][local_idx]; +ss << " let weight_data = weight_tile[inner_k_idx][local_idx];\n"; +// 121 | #if use_subgroup +if (__param_use_subgroup) { +// 122 | let src_data = src_tile[inner_k_idx][sg_id]; +ss << " let src_data = src_tile[inner_k_idx][sg_id];\n"; +// 123 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; +// 124 | results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx))); +ss << " results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx)));\n"; +// 125 | } +ss << " }\n"; +// 126 | #else +} else { +// 127 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; +// 128 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 129 | results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]); +ss << " results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]);\n"; +// 130 | #else +} else { +// 131 | results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx])); +ss << " results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx]));\n"; +// 132 | #endif +} +// 133 | } +ss << " }\n"; +// 134 | #endif +} +// 135 | } +ss << " }\n"; +// 136 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 137 | } +ss << " }\n"; +// 138 | +ss << "\n"; +// 139 | let m_base = m_global_base; +ss << " let m_base = m_global_base;\n"; +// 140 | let n_base = n_global_base + local_idx; +ss << " let n_base = n_global_base + local_idx;\n"; +// 141 | +ss << "\n"; +// 142 | let bias = load_bias(n_base); +ss << " let bias = load_bias(n_base);\n"; +// 143 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; +// 144 | var output_data = results[m_idx] + bias; +ss << " var output_data = results[m_idx] + bias;\n"; +// 145 | write_output(batch, m_base + m_idx, n_base, output_data); +ss << " write_output(batch, m_base + m_idx, n_base, output_data);\n"; +// 146 | } +ss << " }\n"; +// 147 | } // MAIN +MainFunctionEnd(); +ss << "\n"; +// 148 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/oihw_to_ohwi.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/oihw_to_ohwi.h new file mode 100644 index 0000000000000..a9eb76724c328 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/oihw_to_ohwi.h @@ -0,0 +1,128 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: tensor/oihw_to_ohwi.wgsl.template + +template <> +Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract variables + auto& __var_output = *params.var_output; + auto& __var_src = *params.var_src; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | #use .getByOffset .setByOffset +// 5 | +// 6 | fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t { +ss << "fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t {\n"; +// 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { +ss << " if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) {\n"; +// 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + +ss << " let offset = co * uniforms.I * uniforms.H * uniforms.W +\n"; +// 9 | ci * uniforms.H * uniforms.W + +ss << " ci * uniforms.H * uniforms.W +\n"; +// 10 | h_w; +ss << " h_w;\n"; +// 11 | return src.getByOffset(offset); +ss << " return "; +ss << __var_src.GetByOffset("offset"); +ss << ";\n"; +// 12 | } +ss << " }\n"; +// 13 | return src_element_t(); +ss << " return src_element_t();\n"; +// 14 | } +ss << "}\n"; +// 15 | +ss << "\n"; +// 16 | fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) { +ss << "fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) {\n"; +// 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { +ss << " if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) {\n"; +// 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + +ss << " let offset = co * uniforms.H * uniforms.W * uniforms.I +\n"; +// 19 | h_w * uniforms.I + +ss << " h_w * uniforms.I +\n"; +// 20 | ci; +ss << " ci;\n"; +// 21 | output.setByOffset(offset, value); +ss << " "; +ss << __var_output.SetByOffset("offset", "value"); +ss << ";\n"; +// 22 | } +ss << " }\n"; +// 23 | } +ss << "}\n"; +// 24 | +ss << "\n"; +// 25 | var data_cache : array, 4>; +ss << "var data_cache : array, 4>;\n"; +// 26 | +ss << "\n"; +// 27 | $MAIN { +MainFunctionStart(); +ss << "\n"; +// 28 | let group_co : u32 = workgroup_idx / uniforms.Ci_tiles; +ss << " let group_co : u32 = workgroup_idx / uniforms.Ci_tiles;\n"; +// 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; +ss << " let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64;\n"; +// 30 | +ss << "\n"; +// 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { +ss << " if (group_co >= uniforms.O || group_ci >= uniforms.I) {\n"; +// 32 | return; +ss << " return;\n"; +// 33 | } +ss << " }\n"; +// 34 | +ss << "\n"; +// 35 | for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) { +ss << " for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) {\n"; +// 36 | // load +ss << "\n"; +// 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { +ss << " for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) {\n"; +// 38 | let load_ci_idx = ci_idx + local_idx / 4; +ss << " let load_ci_idx = ci_idx + local_idx / 4;\n"; +// 39 | let load_h_w_idx = local_idx % 4; +ss << " let load_h_w_idx = local_idx % 4;\n"; +// 40 | +ss << "\n"; +// 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, +ss << " data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co,\n"; +// 42 | group_ci + load_ci_idx, +ss << " group_ci + load_ci_idx,\n"; +// 43 | h_w_idx * 4 + load_h_w_idx); +ss << " h_w_idx * 4 + load_h_w_idx);\n"; +// 44 | } +ss << " }\n"; +// 45 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 46 | +ss << "\n"; +// 47 | // store +// 48 | for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) { +ss << " for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) {\n"; +// 49 | let output_data = data_cache[local_h_w_idx][local_idx]; +ss << " let output_data = data_cache[local_h_w_idx][local_idx];\n"; +// 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); +ss << " write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data);\n"; +// 51 | } +ss << " }\n"; +// 52 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 53 | } +ss << " }\n"; +// 54 | } // MAIN +MainFunctionEnd(); +ss << "\n"; +// 55 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/pad.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/pad.h new file mode 100644 index 0000000000000..081654e6385c7 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/tensor/pad.h @@ -0,0 +1,178 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: tensor/pad.wgsl.template + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_dim_value_zero = params.param_dim_value_zero; + auto& __param_is_float16 = params.param_is_float16; + auto& __param_pad_mode = params.param_pad_mode; + + // Extract variables + auto& __var_output = *params.var_output; + +// 1 | #define PAD_MODE_CONSTANT 0 +// 2 | #define PAD_MODE_REFLECT 1 +// 3 | #define PAD_MODE_EDGE 2 +// 4 | #define PAD_MODE_WRAP 3 +// 5 | +// 6 | #use guardAgainstOutOfBoundsWorkgroupSizes +// 7 | #use getElementAt +// 8 | #use .offsetToIndices .setByOffset .rank +// 9 | +// 10 | #param dim_value_zero +// 11 | #param is_float16 +// 12 | #param pad_mode +// 13 | +// 14 | $MAIN { +MainFunctionStart(); +ss << "\n"; +// 15 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); +ss << " "; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size"); +ss << ";\n"; +// 16 | +ss << "\n"; +// 17 | let constant_value = +ss << " let constant_value =\n"; +// 18 | #if is_float16 +if (__param_is_float16) { +// 19 | bitcast>(uniforms.constant_value)[0]; +ss << " bitcast>(uniforms.constant_value)[0];\n"; +// 20 | #else +} else { +// 21 | bitcast(uniforms.constant_value); +ss << " bitcast(uniforms.constant_value);\n"; +// 22 | #endif +} +// 23 | +// 24 | #if dim_value_zero +if (__param_dim_value_zero) { +// 25 | output[global_idx] = constant_value; +ss << " output[global_idx] = constant_value;\n"; +// 26 | #else +} else { +// 27 | let output_indices = output.offsetToIndices(global_idx); +ss << " let output_indices = "; +ss << __var_output.OffsetToIndices("global_idx"); +ss << ";\n"; +// 28 | var input_index = u32(0); +ss << " var input_index = u32(0);\n"; +// 29 | var use_pad_value = false; +ss << " var use_pad_value = false;\n"; +// 30 | var in_coord = i32(0); +ss << " var in_coord = i32(0);\n"; +// 31 | +ss << "\n"; +// 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { +ss << " for (var dim = 0; dim < "; +ss << __var_output.Rank(); +ss << " && !use_pad_value; dim++) {\n"; +// 33 | let output_index = i32(getElementAt(output_indices, dim, output.rank)); +ss << " let output_index = i32("; +ss << GetElementAt("output_indices", "dim", __var_output.Rank()); +ss << ");\n"; +// 34 | let lower_pads = getElementAt(uniforms.lower_pads, dim, output.rank); +ss << " let lower_pads = "; +ss << GetElementAt("uniforms.lower_pads", "dim", __var_output.Rank()); +ss << ";\n"; +// 35 | let data_shape = i32(getElementAt(uniforms.data_shape, dim, output.rank)); +ss << " let data_shape = i32("; +ss << GetElementAt("uniforms.data_shape", "dim", __var_output.Rank()); +ss << ");\n"; +// 36 | #if pad_mode == PAD_MODE_CONSTANT +if (__param_pad_mode == 0) { +// 37 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n"; +// 38 | use_pad_value = true; +ss << " use_pad_value = true;\n"; +// 39 | #elif pad_mode == PAD_MODE_EDGE +} else if (__param_pad_mode == 2) { +// 40 | if (output_index < lower_pads) { +ss << " if (output_index < lower_pads) {\n"; +// 41 | in_coord = 0; +ss << " in_coord = 0;\n"; +// 42 | } else if (output_index >= data_shape + lower_pads) { +ss << " } else if (output_index >= data_shape + lower_pads) {\n"; +// 43 | in_coord = data_shape - 1; +ss << " in_coord = data_shape - 1;\n"; +// 44 | #elif pad_mode == PAD_MODE_REFLECT +} else if (__param_pad_mode == 1) { +// 45 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n"; +// 46 | in_coord = output_index - lower_pads; +ss << " in_coord = output_index - lower_pads;\n"; +// 47 | if (in_coord < 0) { +ss << " if (in_coord < 0) {\n"; +// 48 | in_coord = -in_coord; +ss << " in_coord = -in_coord;\n"; +// 49 | } +ss << " }\n"; +// 50 | let _2n_1 = 2 * (data_shape - 1); +ss << " let _2n_1 = 2 * (data_shape - 1);\n"; +// 51 | in_coord = in_coord % _2n_1; +ss << " in_coord = in_coord % _2n_1;\n"; +// 52 | if (in_coord >= data_shape) { +ss << " if (in_coord >= data_shape) {\n"; +// 53 | in_coord = _2n_1 - in_coord; +ss << " in_coord = _2n_1 - in_coord;\n"; +// 54 | } +ss << " }\n"; +// 55 | #else // PAD_MODE_WRAP +} else { +// 56 | if (output_index < lower_pads) { +ss << " if (output_index < lower_pads) {\n"; +// 57 | in_coord = data_shape + output_index - lower_pads; +ss << " in_coord = data_shape + output_index - lower_pads;\n"; +// 58 | } else if (output_index >= data_shape + lower_pads) { +ss << " } else if (output_index >= data_shape + lower_pads) {\n"; +// 59 | in_coord = output_index - data_shape - lower_pads; +ss << " in_coord = output_index - data_shape - lower_pads;\n"; +// 60 | #endif // pad_mode +} +// 61 | } else { +ss << " } else {\n"; +// 62 | in_coord = output_index - lower_pads; +ss << " in_coord = output_index - lower_pads;\n"; +// 63 | } +ss << " }\n"; +// 64 | +ss << "\n"; +// 65 | input_index += select(u32(in_coord) +ss << " input_index += select(u32(in_coord)\n"; +// 66 | #if output.rank > 1 +if (__var_output.Rank() > 1) { +// 67 | * getElementAt(uniforms.data_stride, dim, output.rank - 1) +ss << " * "; +ss << GetElementAt("uniforms.data_stride", "dim", __var_output.Rank() - 1); +ss << "\n"; +// 68 | #endif +} +// 69 | , u32(in_coord), dim == output.rank - 1); +ss << " , u32(in_coord), dim == "; +ss << __var_output.Rank(); +ss << " - 1);\n"; +// 70 | } +ss << " }\n"; +// 71 | +ss << "\n"; +// 72 | output.setByOffset(global_idx, select(data[input_index], constant_value, use_pad_value)); +ss << " "; +ss << __var_output.SetByOffset("global_idx", "select(data[input_index], constant_value, use_pad_value)"); +ss << ";\n"; +// 73 | #endif +} +// 74 | } // MAIN +MainFunctionEnd(); +ss << "\n"; +// 75 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h new file mode 100644 index 0000000000000..76f61dd62a030 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h @@ -0,0 +1,58 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#ifndef INCLUDED_BY_WGSL_GEN_HEADER +#error "This file is expected to be included by wgsl-gen header. Do not include it directly." +#endif + +// +// Template: nn/im2col_matmul.wgsl.template +// + +template <> +struct TemplateParameter<"nn/im2col_matmul.wgsl.template"> { + using type = struct { + int param_has_bias; + int param_tile_m; + int param_tile_n; + int param_use_subgroup; + int param_vec_size; + const ShaderVariableHelper* var_output; + const ShaderVariableHelper* var_src; + const ShaderVariableHelper* var_weight; + }; +}; + +template <> +Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params); + +// +// Template: tensor/oihw_to_ohwi.wgsl.template +// + +template <> +struct TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template"> { + using type = struct { + const ShaderVariableHelper* var_output; + const ShaderVariableHelper* var_src; + }; +}; + +template <> +Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template">::type params); + +// +// Template: tensor/pad.wgsl.template +// + +template <> +struct TemplateParameter<"tensor/pad.wgsl.template"> { + using type = struct { + int param_dim_value_zero; + int param_is_float16; + int param_pad_mode; + const ShaderVariableHelper* var_output; + }; +}; + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h new file mode 100644 index 0000000000000..932eb9d9f61f4 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h @@ -0,0 +1,42 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// Helper functions or macros + +#pragma push_macro("MainFunctionStart") +#undef MainFunctionStart +#define MainFunctionStart() { [[maybe_unused]] auto& ss = shader_helper.MainFunctionBody(); +#pragma push_macro("MainFunctionEnd") +#undef MainFunctionEnd +#define MainFunctionEnd() } + +// Helper templates + +namespace wgsl_detail { +template >> +std::string pass_as_string(T&& v) { + return std::to_string(std::forward(v)); +} +template +std::string_view pass_as_string(std::string_view sv) { + return sv; +} +template +std::string pass_as_string(T&& v) { + return std::forward(v); +} +} // namespace wgsl_detail + + +// Include template implementations + +#include "wgsl_template_gen/generated/nn/im2col_matmul.h" // d8d68023c1442e7366ed4a28cb4dd402ba4a6da8629fa4b926f4042b7ba6d70c +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 35487692058e08b027768dcb1ab30ef93772da9a4a664702c34351a74d2568dc +#include "wgsl_template_gen/generated/tensor/pad.h" // 3b9bc9a6a50b8642d1220b2354b8e88f516f664ed9b4bbe8b5d9e7fe4f4b97b0 + +#pragma pop_macro("MainFunctionStart") +#pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h new file mode 100644 index 0000000000000..d7084e97b509e --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h @@ -0,0 +1,323 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: nn/im2col_matmul.wgsl.template + +template <> +Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_has_bias = params.param_has_bias; + auto& __param_tile_m = params.param_tile_m; + auto& __param_tile_n = params.param_tile_n; + auto& __param_use_subgroup = params.param_use_subgroup; + auto& __param_vec_size = params.param_vec_size; + + // Extract variables + auto& __var_output = *params.var_output; + auto& __var_src = *params.var_src; + auto& __var_weight = *params.var_weight; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | #param has_bias +// 5 | #param tile_m +// 6 | #param tile_n +// 7 | #param use_subgroup +// 8 | #param vec_size +// 9 | +// 10 | #use .getByOffset .setByOffset +// 11 | +// 12 | // im2col access for src: [N, H_i, W_i, C_i / vec_size] +// 13 | // Conceptual Matrix Shape: N * (H_o * W_o) x (K_h * K_w * C_i / vec_size) +// 14 | fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t { +ss << __str_4; +// 15 | if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * vec_size >= uniforms.im2col_k) { +ss << __str_5; +ss << __param_vec_size; +ss << __str_6; +// 16 | return src_value_t(); +ss << __str_7; +// 17 | } +ss << __str_8; +// 18 | +ss << __str_9; +// 19 | let channel_i_vec = uniforms.channel_i / vec_size; +ss << __str_10; +ss << __param_vec_size; +ss << __str_11; +// 20 | +ss << __str_9; +// 21 | // 1. Decompose M index (H_o * W_o) into (h_idx, w_idx) +// 22 | let h_idx = m / uniforms.output_w; // Output H index (H_o) +ss << __str_12; +// 23 | let w_idx = m % uniforms.output_w; // Output W index (W_o) +ss << __str_13; +// 24 | +ss << __str_9; +// 25 | // 2. Decompose K index into (k_h, k_w, c_i_vec_idx) +// 26 | let c_i_vec_idx = k_packed_idx % channel_i_vec; +ss << __str_14; +// 27 | let k_h_w_idx = k_packed_idx / channel_i_vec; +ss << __str_15; +// 28 | let k_h = k_h_w_idx / uniforms.kernel_w; // Kernel Row +ss << __str_16; +// 29 | let k_w = k_h_w_idx % uniforms.kernel_w; // Kernel Column +ss << __str_17; +// 30 | +ss << __str_9; +// 31 | // 3. Calculate the coordinate in the padded input tensor +// 32 | let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x; +ss << __str_18; +// 33 | let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y; +ss << __str_19; +// 34 | +ss << __str_9; +// 35 | // 4. Calculate the coordinate in the original input tensor +// 36 | let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x); +ss << __str_20; +// 37 | let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y); +ss << __str_21; +// 38 | +ss << __str_9; +// 39 | // 5. Check for padding/out-of-bounds +// 40 | if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) || +ss << __str_22; +// 41 | src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) { +ss << __str_23; +// 42 | return src_value_t(); +ss << __str_7; +// 43 | } +ss << __str_8; +// 44 | +ss << __str_9; +// 45 | // 6. Calculate final NHWC index +// 46 | let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec + +ss << __str_24; +// 47 | u32(src_h_coord) * uniforms.src_w * channel_i_vec + +ss << __str_25; +// 48 | u32(src_w_coord) * channel_i_vec + +ss << __str_26; +// 49 | c_i_vec_idx; +ss << __str_27; +// 50 | return src.getByOffset(src_idx); +ss << __str_28; +ss << __var_src.GetByOffset(__str_0); +ss << __str_11; +// 51 | } +ss << __str_29; +// 52 | +ss << __str_9; +// 53 | // weight shape: [Co, K_h, K_w, C_i / vec_size] (CoHWCi) +// 54 | fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t { +ss << __str_30; +// 55 | if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / vec_size) { +ss << __str_31; +ss << __param_vec_size; +ss << __str_32; +// 56 | let weight_idx = n * uniforms.im2col_k / vec_size + +ss << __str_33; +ss << __param_vec_size; +ss << __str_34; +// 57 | k_packed_idx; +ss << __str_35; +// 58 | return weight.getByOffset(weight_idx); +ss << __str_36; +ss << __var_weight.GetByOffset(__str_1); +ss << __str_11; +// 59 | } +ss << __str_8; +// 60 | return weight_value_t(); +ss << __str_37; +// 61 | } +ss << __str_29; +// 62 | +ss << __str_9; +// 63 | fn load_bias(n : u32) -> output_element_t { +ss << __str_38; +// 64 | #if has_bias +if (__param_has_bias) { +// 65 | if (n < uniforms.im2col_n) { +ss << __str_39; +// 66 | return output_element_t(bias[n]); +ss << __str_40; +// 67 | } +ss << __str_8; +// 68 | #endif +} +// 69 | return output_element_t(); +ss << __str_41; +// 70 | } +ss << __str_29; +// 71 | +ss << __str_9; +// 72 | // output shape: [N, H_o, W_o, C_o] (NHWC) +// 73 | fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) { +ss << __str_42; +// 74 | if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) { +ss << __str_43; +// 75 | let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n + +ss << __str_44; +// 76 | m * uniforms.im2col_n + +ss << __str_45; +// 77 | n; +ss << __str_46; +// 78 | output.setByOffset(output_idx, value); +ss << __str_47; +ss << __var_output.SetByOffset(__str_2, __str_3); +ss << __str_11; +// 79 | } +ss << __str_8; +// 80 | } +ss << __str_29; +// 81 | +ss << __str_9; +// 82 | const TILE_M_SIZE : u32 = tile_m; +ss << __str_48; +ss << __param_tile_m; +ss << __str_11; +// 83 | const TILE_N_SIZE : u32 = tile_n; +ss << __str_49; +ss << __param_tile_n; +ss << __str_11; +// 84 | // In dimension K, the tile consists of 16 scalars, requiring `16 / vec_size` vector loads. +ss << __str_9; +// 85 | const TILE_K_VEC_SIZE : u32 = 16 / vec_size; +ss << __str_50; +ss << __param_vec_size; +ss << __str_11; +// 86 | // In dimensions M and N, since a workgroup has 64 threads, it advances by `64 / TILE_K_VEC_SIZE`. +ss << __str_9; +// 87 | const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE; +ss << __str_51; +// 88 | +ss << __str_9; +// 89 | var src_tile : array, TILE_K_VEC_SIZE>; +ss << __str_52; +// 90 | var weight_tile : array, TILE_K_VEC_SIZE>; +ss << __str_53; +// 91 | +ss << __str_9; +// 92 | $MAIN { +MainFunctionStart(); +ss << __str_9; +// 93 | let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles); +ss << __str_54; +// 94 | let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE; +ss << __str_55; +// 95 | let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE; +ss << __str_56; +// 96 | +ss << __str_9; +// 97 | var results : array; +ss << __str_57; +// 98 | for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) { +ss << __str_58; +// 99 | for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) { +ss << __str_59; +// 100 | // Loads a 64 vec of src into the workgroup memory. +ss << __str_9; +// 101 | let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE; +ss << __str_60; +// 102 | let load_src_k = local_idx % TILE_K_VEC_SIZE; +ss << __str_61; +// 103 | +ss << __str_9; +// 104 | src_tile[load_src_k][load_src_m] = load_src(batch, +ss << __str_62; +// 105 | m_global_base + load_src_m, +ss << __str_63; +// 106 | k_idx * TILE_K_VEC_SIZE + load_src_k); +ss << __str_64; +// 107 | } +ss << __str_65; +// 108 | +ss << __str_9; +// 109 | for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) { +ss << __str_66; +// 110 | // Loads a 64 vec of weight into the workgroup memory. +ss << __str_9; +// 111 | let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE; +ss << __str_67; +// 112 | let load_weight_k = local_idx % TILE_K_VEC_SIZE; +ss << __str_68; +// 113 | +ss << __str_9; +// 114 | weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n, +ss << __str_69; +// 115 | k_idx * TILE_K_VEC_SIZE + load_weight_k); +ss << __str_70; +// 116 | } +ss << __str_65; +// 117 | workgroupBarrier(); +ss << __str_71; +// 118 | +ss << __str_9; +// 119 | for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) { +ss << __str_72; +// 120 | let weight_data = weight_tile[inner_k_idx][local_idx]; +ss << __str_73; +// 121 | #if use_subgroup +if (__param_use_subgroup) { +// 122 | let src_data = src_tile[inner_k_idx][sg_id]; +ss << __str_74; +// 123 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << __str_75; +// 124 | results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx))); +ss << __str_76; +// 125 | } +ss << __str_77; +// 126 | #else +} else { +// 127 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << __str_75; +// 128 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 129 | results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]); +ss << __str_78; +// 130 | #else +} else { +// 131 | results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx])); +ss << __str_79; +// 132 | #endif +} +// 133 | } +ss << __str_77; +// 134 | #endif +} +// 135 | } +ss << __str_65; +// 136 | workgroupBarrier(); +ss << __str_71; +// 137 | } +ss << __str_8; +// 138 | +ss << __str_9; +// 139 | let m_base = m_global_base; +ss << __str_80; +// 140 | let n_base = n_global_base + local_idx; +ss << __str_81; +// 141 | +ss << __str_9; +// 142 | let bias = load_bias(n_base); +ss << __str_82; +// 143 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +ss << __str_83; +// 144 | var output_data = results[m_idx] + bias; +ss << __str_84; +// 145 | write_output(batch, m_base + m_idx, n_base, output_data); +ss << __str_85; +// 146 | } +ss << __str_8; +// 147 | } // MAIN +MainFunctionEnd(); +ss << __str_9; +// 148 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h new file mode 100644 index 0000000000000..db0d782cd9813 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h @@ -0,0 +1,128 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: tensor/oihw_to_ohwi.wgsl.template + +template <> +Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract variables + auto& __var_output = *params.var_output; + auto& __var_src = *params.var_src; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | #use .getByOffset .setByOffset +// 5 | +// 6 | fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t { +ss << __str_87; +// 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { +ss << __str_88; +// 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + +ss << __str_89; +// 9 | ci * uniforms.H * uniforms.W + +ss << __str_90; +// 10 | h_w; +ss << __str_91; +// 11 | return src.getByOffset(offset); +ss << __str_36; +ss << __var_src.GetByOffset(__str_86); +ss << __str_11; +// 12 | } +ss << __str_8; +// 13 | return src_element_t(); +ss << __str_92; +// 14 | } +ss << __str_29; +// 15 | +ss << __str_9; +// 16 | fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) { +ss << __str_93; +// 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { +ss << __str_88; +// 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + +ss << __str_94; +// 19 | h_w * uniforms.I + +ss << __str_95; +// 20 | ci; +ss << __str_96; +// 21 | output.setByOffset(offset, value); +ss << __str_47; +ss << __var_output.SetByOffset(__str_86, __str_3); +ss << __str_11; +// 22 | } +ss << __str_8; +// 23 | } +ss << __str_29; +// 24 | +ss << __str_9; +// 25 | var data_cache : array, 4>; +ss << __str_97; +// 26 | +ss << __str_9; +// 27 | $MAIN { +MainFunctionStart(); +ss << __str_9; +// 28 | let group_co : u32 = workgroup_idx / uniforms.Ci_tiles; +ss << __str_98; +// 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; +ss << __str_99; +// 30 | +ss << __str_9; +// 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { +ss << __str_100; +// 32 | return; +ss << __str_101; +// 33 | } +ss << __str_8; +// 34 | +ss << __str_9; +// 35 | for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) { +ss << __str_102; +// 36 | // load +ss << __str_9; +// 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { +ss << __str_103; +// 38 | let load_ci_idx = ci_idx + local_idx / 4; +ss << __str_104; +// 39 | let load_h_w_idx = local_idx % 4; +ss << __str_105; +// 40 | +ss << __str_9; +// 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, +ss << __str_106; +// 42 | group_ci + load_ci_idx, +ss << __str_107; +// 43 | h_w_idx * 4 + load_h_w_idx); +ss << __str_108; +// 44 | } +ss << __str_65; +// 45 | workgroupBarrier(); +ss << __str_71; +// 46 | +ss << __str_9; +// 47 | // store +// 48 | for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) { +ss << __str_109; +// 49 | let output_data = data_cache[local_h_w_idx][local_idx]; +ss << __str_110; +// 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); +ss << __str_111; +// 51 | } +ss << __str_65; +// 52 | workgroupBarrier(); +ss << __str_71; +// 53 | } +ss << __str_8; +// 54 | } // MAIN +MainFunctionEnd(); +ss << __str_9; +// 55 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h new file mode 100644 index 0000000000000..8b6dfdf13dece --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h @@ -0,0 +1,178 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: tensor/pad.wgsl.template + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_dim_value_zero = params.param_dim_value_zero; + auto& __param_is_float16 = params.param_is_float16; + auto& __param_pad_mode = params.param_pad_mode; + + // Extract variables + auto& __var_output = *params.var_output; + +// 1 | #define PAD_MODE_CONSTANT 0 +// 2 | #define PAD_MODE_REFLECT 1 +// 3 | #define PAD_MODE_EDGE 2 +// 4 | #define PAD_MODE_WRAP 3 +// 5 | +// 6 | #use guardAgainstOutOfBoundsWorkgroupSizes +// 7 | #use getElementAt +// 8 | #use .offsetToIndices .setByOffset .rank +// 9 | +// 10 | #param dim_value_zero +// 11 | #param is_float16 +// 12 | #param pad_mode +// 13 | +// 14 | $MAIN { +MainFunctionStart(); +ss << __str_9; +// 15 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); +ss << __str_120; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_112); +ss << __str_11; +// 16 | +ss << __str_9; +// 17 | let constant_value = +ss << __str_121; +// 18 | #if is_float16 +if (__param_is_float16) { +// 19 | bitcast>(uniforms.constant_value)[0]; +ss << __str_122; +// 20 | #else +} else { +// 21 | bitcast(uniforms.constant_value); +ss << __str_123; +// 22 | #endif +} +// 23 | +// 24 | #if dim_value_zero +if (__param_dim_value_zero) { +// 25 | output[global_idx] = constant_value; +ss << __str_124; +// 26 | #else +} else { +// 27 | let output_indices = output.offsetToIndices(global_idx); +ss << __str_125; +ss << __var_output.OffsetToIndices(__str_113); +ss << __str_11; +// 28 | var input_index = u32(0); +ss << __str_126; +// 29 | var use_pad_value = false; +ss << __str_127; +// 30 | var in_coord = i32(0); +ss << __str_128; +// 31 | +ss << __str_9; +// 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { +ss << __str_129; +ss << __var_output.Rank(); +ss << __str_130; +// 33 | let output_index = i32(getElementAt(output_indices, dim, output.rank)); +ss << __str_131; +ss << GetElementAt(__str_114, __str_115, __var_output.Rank()); +ss << __str_132; +// 34 | let lower_pads = getElementAt(uniforms.lower_pads, dim, output.rank); +ss << __str_133; +ss << GetElementAt(__str_116, __str_115, __var_output.Rank()); +ss << __str_11; +// 35 | let data_shape = i32(getElementAt(uniforms.data_shape, dim, output.rank)); +ss << __str_134; +ss << GetElementAt(__str_117, __str_115, __var_output.Rank()); +ss << __str_132; +// 36 | #if pad_mode == PAD_MODE_CONSTANT +if (__param_pad_mode == 0) { +// 37 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { +ss << __str_135; +// 38 | use_pad_value = true; +ss << __str_136; +// 39 | #elif pad_mode == PAD_MODE_EDGE +} else if (__param_pad_mode == 2) { +// 40 | if (output_index < lower_pads) { +ss << __str_137; +// 41 | in_coord = 0; +ss << __str_138; +// 42 | } else if (output_index >= data_shape + lower_pads) { +ss << __str_139; +// 43 | in_coord = data_shape - 1; +ss << __str_140; +// 44 | #elif pad_mode == PAD_MODE_REFLECT +} else if (__param_pad_mode == 1) { +// 45 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { +ss << __str_135; +// 46 | in_coord = output_index - lower_pads; +ss << __str_141; +// 47 | if (in_coord < 0) { +ss << __str_142; +// 48 | in_coord = -in_coord; +ss << __str_143; +// 49 | } +ss << __str_77; +// 50 | let _2n_1 = 2 * (data_shape - 1); +ss << __str_144; +// 51 | in_coord = in_coord % _2n_1; +ss << __str_145; +// 52 | if (in_coord >= data_shape) { +ss << __str_146; +// 53 | in_coord = _2n_1 - in_coord; +ss << __str_147; +// 54 | } +ss << __str_77; +// 55 | #else // PAD_MODE_WRAP +} else { +// 56 | if (output_index < lower_pads) { +ss << __str_137; +// 57 | in_coord = data_shape + output_index - lower_pads; +ss << __str_148; +// 58 | } else if (output_index >= data_shape + lower_pads) { +ss << __str_139; +// 59 | in_coord = output_index - data_shape - lower_pads; +ss << __str_149; +// 60 | #endif // pad_mode +} +// 61 | } else { +ss << __str_150; +// 62 | in_coord = output_index - lower_pads; +ss << __str_151; +// 63 | } +ss << __str_65; +// 64 | +ss << __str_9; +// 65 | input_index += select(u32(in_coord) +ss << __str_152; +// 66 | #if output.rank > 1 +if (__var_output.Rank() > 1) { +// 67 | * getElementAt(uniforms.data_stride, dim, output.rank - 1) +ss << __str_153; +ss << GetElementAt(__str_118, __str_115, __var_output.Rank() - 1); +ss << __str_9; +// 68 | #endif +} +// 69 | , u32(in_coord), dim == output.rank - 1); +ss << __str_154; +ss << __var_output.Rank(); +ss << __str_155; +// 70 | } +ss << __str_8; +// 71 | +ss << __str_9; +// 72 | output.setByOffset(global_idx, select(data[input_index], constant_value, use_pad_value)); +ss << __str_120; +ss << __var_output.SetByOffset(__str_113, __str_119); +ss << __str_11; +// 73 | #endif +} +// 74 | } // MAIN +MainFunctionEnd(); +ss << __str_9; +// 75 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h new file mode 100644 index 0000000000000..76f61dd62a030 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h @@ -0,0 +1,58 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#ifndef INCLUDED_BY_WGSL_GEN_HEADER +#error "This file is expected to be included by wgsl-gen header. Do not include it directly." +#endif + +// +// Template: nn/im2col_matmul.wgsl.template +// + +template <> +struct TemplateParameter<"nn/im2col_matmul.wgsl.template"> { + using type = struct { + int param_has_bias; + int param_tile_m; + int param_tile_n; + int param_use_subgroup; + int param_vec_size; + const ShaderVariableHelper* var_output; + const ShaderVariableHelper* var_src; + const ShaderVariableHelper* var_weight; + }; +}; + +template <> +Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params); + +// +// Template: tensor/oihw_to_ohwi.wgsl.template +// + +template <> +struct TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template"> { + using type = struct { + const ShaderVariableHelper* var_output; + const ShaderVariableHelper* var_src; + }; +}; + +template <> +Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/oihw_to_ohwi.wgsl.template">::type params); + +// +// Template: tensor/pad.wgsl.template +// + +template <> +struct TemplateParameter<"tensor/pad.wgsl.template"> { + using type = struct { + int param_dim_value_zero; + int param_is_float16; + int param_pad_mode; + const ShaderVariableHelper* var_output; + }; +}; + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h new file mode 100644 index 0000000000000..3c1acbf81a6bd --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h @@ -0,0 +1,43 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// Helper functions or macros + +#pragma push_macro("MainFunctionStart") +#undef MainFunctionStart +#define MainFunctionStart() { [[maybe_unused]] auto& ss = shader_helper.MainFunctionBody(); +#pragma push_macro("MainFunctionEnd") +#undef MainFunctionEnd +#define MainFunctionEnd() } + +// Helper templates + +namespace wgsl_detail { +template >> +std::string pass_as_string(T&& v) { + return std::to_string(std::forward(v)); +} +template +std::string_view pass_as_string(std::string_view sv) { + return sv; +} +template +std::string pass_as_string(T&& v) { + return std::forward(v); +} +} // namespace wgsl_detail + +#include "wgsl_template_gen/string_table.h" // fd21a84e13681c97e246eabaf5ca83c6687ba9da332980cda8aa246eee7390fe + +// Include template implementations + +#include "wgsl_template_gen/generated/nn/im2col_matmul.h" // 2c4e959f73738f0507ad0df24eee5cc62c674bb45409f3e8eee738af6573bb48 +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // ce8ad040ed278e6db07d7eefd7920831c05ee5ac8828f080fa0ecf67e53eb737 +#include "wgsl_template_gen/generated/tensor/pad.h" // 6b454a17c508a7ef0b1df42c549a304719779242d6a76f63e9a3f5eceb19cb78 + +#pragma pop_macro("MainFunctionStart") +#pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h new file mode 100644 index 0000000000000..da97b7d44ed56 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h @@ -0,0 +1,164 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// String table constants +constexpr const char* __str_0 = "src_idx"; +constexpr const char* __str_1 = "weight_idx"; +constexpr const char* __str_2 = "output_idx"; +constexpr const char* __str_3 = "value"; +constexpr const char* __str_4 = "fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t {\n"; +constexpr const char* __str_5 = " if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * "; +constexpr const char* __str_6 = " >= uniforms.im2col_k) {\n"; +constexpr const char* __str_7 = " return src_value_t();\n"; +constexpr const char* __str_8 = " }\n"; +constexpr const char* __str_9 = "\n"; +constexpr const char* __str_10 = " let channel_i_vec = uniforms.channel_i / "; +constexpr const char* __str_11 = ";\n"; +constexpr const char* __str_12 = " let h_idx = m / uniforms.output_w;\n"; +constexpr const char* __str_13 = " let w_idx = m % uniforms.output_w;\n"; +constexpr const char* __str_14 = " let c_i_vec_idx = k_packed_idx % channel_i_vec;\n"; +constexpr const char* __str_15 = " let k_h_w_idx = k_packed_idx / channel_i_vec;\n"; +constexpr const char* __str_16 = " let k_h = k_h_w_idx / uniforms.kernel_w;\n"; +constexpr const char* __str_17 = " let k_w = k_h_w_idx % uniforms.kernel_w;\n"; +constexpr const char* __str_18 = " let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x;\n"; +constexpr const char* __str_19 = " let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y;\n"; +constexpr const char* __str_20 = " let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x);\n"; +constexpr const char* __str_21 = " let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y);\n"; +constexpr const char* __str_22 = " if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) ||\n"; +constexpr const char* __str_23 = " src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) {\n"; +constexpr const char* __str_24 = " let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec +\n"; +constexpr const char* __str_25 = " u32(src_h_coord) * uniforms.src_w * channel_i_vec +\n"; +constexpr const char* __str_26 = " u32(src_w_coord) * channel_i_vec +\n"; +constexpr const char* __str_27 = " c_i_vec_idx;\n"; +constexpr const char* __str_28 = " return "; +constexpr const char* __str_29 = "}\n"; +constexpr const char* __str_30 = "fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t {\n"; +constexpr const char* __str_31 = " if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / "; +constexpr const char* __str_32 = ") {\n"; +constexpr const char* __str_33 = " let weight_idx = n * uniforms.im2col_k / "; +constexpr const char* __str_34 = " +\n"; +constexpr const char* __str_35 = " k_packed_idx;\n"; +constexpr const char* __str_36 = " return "; +constexpr const char* __str_37 = " return weight_value_t();\n"; +constexpr const char* __str_38 = "fn load_bias(n : u32) -> output_element_t {\n"; +constexpr const char* __str_39 = " if (n < uniforms.im2col_n) {\n"; +constexpr const char* __str_40 = " return output_element_t(bias[n]);\n"; +constexpr const char* __str_41 = " return output_element_t();\n"; +constexpr const char* __str_42 = "fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) {\n"; +constexpr const char* __str_43 = " if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) {\n"; +constexpr const char* __str_44 = " let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n +\n"; +constexpr const char* __str_45 = " m * uniforms.im2col_n +\n"; +constexpr const char* __str_46 = " n;\n"; +constexpr const char* __str_47 = " "; +constexpr const char* __str_48 = "const TILE_M_SIZE : u32 = "; +constexpr const char* __str_49 = "const TILE_N_SIZE : u32 = "; +constexpr const char* __str_50 = "const TILE_K_VEC_SIZE : u32 = 16 / "; +constexpr const char* __str_51 = "const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE;\n"; +constexpr const char* __str_52 = "var src_tile : array, TILE_K_VEC_SIZE>;\n"; +constexpr const char* __str_53 = "var weight_tile : array, TILE_K_VEC_SIZE>;\n"; +constexpr const char* __str_54 = " let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles);\n"; +constexpr const char* __str_55 = " let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE;\n"; +constexpr const char* __str_56 = " let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE;\n"; +constexpr const char* __str_57 = " var results : array;\n"; +constexpr const char* __str_58 = " for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) {\n"; +constexpr const char* __str_59 = " for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) {\n"; +constexpr const char* __str_60 = " let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE;\n"; +constexpr const char* __str_61 = " let load_src_k = local_idx % TILE_K_VEC_SIZE;\n"; +constexpr const char* __str_62 = " src_tile[load_src_k][load_src_m] = load_src(batch,\n"; +constexpr const char* __str_63 = " m_global_base + load_src_m,\n"; +constexpr const char* __str_64 = " k_idx * TILE_K_VEC_SIZE + load_src_k);\n"; +constexpr const char* __str_65 = " }\n"; +constexpr const char* __str_66 = " for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) {\n"; +constexpr const char* __str_67 = " let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE;\n"; +constexpr const char* __str_68 = " let load_weight_k = local_idx % TILE_K_VEC_SIZE;\n"; +constexpr const char* __str_69 = " weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n,\n"; +constexpr const char* __str_70 = " k_idx * TILE_K_VEC_SIZE + load_weight_k);\n"; +constexpr const char* __str_71 = " workgroupBarrier();\n"; +constexpr const char* __str_72 = " for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) {\n"; +constexpr const char* __str_73 = " let weight_data = weight_tile[inner_k_idx][local_idx];\n"; +constexpr const char* __str_74 = " let src_data = src_tile[inner_k_idx][sg_id];\n"; +constexpr const char* __str_75 = " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; +constexpr const char* __str_76 = " results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx)));\n"; +constexpr const char* __str_77 = " }\n"; +constexpr const char* __str_78 = " results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]);\n"; +constexpr const char* __str_79 = " results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx]));\n"; +constexpr const char* __str_80 = " let m_base = m_global_base;\n"; +constexpr const char* __str_81 = " let n_base = n_global_base + local_idx;\n"; +constexpr const char* __str_82 = " let bias = load_bias(n_base);\n"; +constexpr const char* __str_83 = " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; +constexpr const char* __str_84 = " var output_data = results[m_idx] + bias;\n"; +constexpr const char* __str_85 = " write_output(batch, m_base + m_idx, n_base, output_data);\n"; +constexpr const char* __str_86 = "offset"; +constexpr const char* __str_87 = "fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t {\n"; +constexpr const char* __str_88 = " if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) {\n"; +constexpr const char* __str_89 = " let offset = co * uniforms.I * uniforms.H * uniforms.W +\n"; +constexpr const char* __str_90 = " ci * uniforms.H * uniforms.W +\n"; +constexpr const char* __str_91 = " h_w;\n"; +constexpr const char* __str_92 = " return src_element_t();\n"; +constexpr const char* __str_93 = "fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) {\n"; +constexpr const char* __str_94 = " let offset = co * uniforms.H * uniforms.W * uniforms.I +\n"; +constexpr const char* __str_95 = " h_w * uniforms.I +\n"; +constexpr const char* __str_96 = " ci;\n"; +constexpr const char* __str_97 = "var data_cache : array, 4>;\n"; +constexpr const char* __str_98 = " let group_co : u32 = workgroup_idx / uniforms.Ci_tiles;\n"; +constexpr const char* __str_99 = " let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64;\n"; +constexpr const char* __str_100 = " if (group_co >= uniforms.O || group_ci >= uniforms.I) {\n"; +constexpr const char* __str_101 = " return;\n"; +constexpr const char* __str_102 = " for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) {\n"; +constexpr const char* __str_103 = " for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) {\n"; +constexpr const char* __str_104 = " let load_ci_idx = ci_idx + local_idx / 4;\n"; +constexpr const char* __str_105 = " let load_h_w_idx = local_idx % 4;\n"; +constexpr const char* __str_106 = " data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co,\n"; +constexpr const char* __str_107 = " group_ci + load_ci_idx,\n"; +constexpr const char* __str_108 = " h_w_idx * 4 + load_h_w_idx);\n"; +constexpr const char* __str_109 = " for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) {\n"; +constexpr const char* __str_110 = " let output_data = data_cache[local_h_w_idx][local_idx];\n"; +constexpr const char* __str_111 = " write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data);\n"; +constexpr const char* __str_112 = "uniforms.output_size"; +constexpr const char* __str_113 = "global_idx"; +constexpr const char* __str_114 = "output_indices"; +constexpr const char* __str_115 = "dim"; +constexpr const char* __str_116 = "uniforms.lower_pads"; +constexpr const char* __str_117 = "uniforms.data_shape"; +constexpr const char* __str_118 = "uniforms.data_stride"; +constexpr const char* __str_119 = "select(data[input_index], constant_value, use_pad_value)"; +constexpr const char* __str_120 = " "; +constexpr const char* __str_121 = " let constant_value =\n"; +constexpr const char* __str_122 = " bitcast>(uniforms.constant_value)[0];\n"; +constexpr const char* __str_123 = " bitcast(uniforms.constant_value);\n"; +constexpr const char* __str_124 = " output[global_idx] = constant_value;\n"; +constexpr const char* __str_125 = " let output_indices = "; +constexpr const char* __str_126 = " var input_index = u32(0);\n"; +constexpr const char* __str_127 = " var use_pad_value = false;\n"; +constexpr const char* __str_128 = " var in_coord = i32(0);\n"; +constexpr const char* __str_129 = " for (var dim = 0; dim < "; +constexpr const char* __str_130 = " && !use_pad_value; dim++) {\n"; +constexpr const char* __str_131 = " let output_index = i32("; +constexpr const char* __str_132 = ");\n"; +constexpr const char* __str_133 = " let lower_pads = "; +constexpr const char* __str_134 = " let data_shape = i32("; +constexpr const char* __str_135 = " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n"; +constexpr const char* __str_136 = " use_pad_value = true;\n"; +constexpr const char* __str_137 = " if (output_index < lower_pads) {\n"; +constexpr const char* __str_138 = " in_coord = 0;\n"; +constexpr const char* __str_139 = " } else if (output_index >= data_shape + lower_pads) {\n"; +constexpr const char* __str_140 = " in_coord = data_shape - 1;\n"; +constexpr const char* __str_141 = " in_coord = output_index - lower_pads;\n"; +constexpr const char* __str_142 = " if (in_coord < 0) {\n"; +constexpr const char* __str_143 = " in_coord = -in_coord;\n"; +constexpr const char* __str_144 = " let _2n_1 = 2 * (data_shape - 1);\n"; +constexpr const char* __str_145 = " in_coord = in_coord % _2n_1;\n"; +constexpr const char* __str_146 = " if (in_coord >= data_shape) {\n"; +constexpr const char* __str_147 = " in_coord = _2n_1 - in_coord;\n"; +constexpr const char* __str_148 = " in_coord = data_shape + output_index - lower_pads;\n"; +constexpr const char* __str_149 = " in_coord = output_index - data_shape - lower_pads;\n"; +constexpr const char* __str_150 = " } else {\n"; +constexpr const char* __str_151 = " in_coord = output_index - lower_pads;\n"; +constexpr const char* __str_152 = " input_index += select(u32(in_coord)\n"; +constexpr const char* __str_153 = " * "; +constexpr const char* __str_154 = " , u32(in_coord), dim == "; +constexpr const char* __str_155 = " - 1);\n"; diff --git a/tools/python/wgsl_template/test/run_tests.py b/tools/python/wgsl_template/test/run_tests.py new file mode 100644 index 0000000000000..b67b106ce6b75 --- /dev/null +++ b/tools/python/wgsl_template/test/run_tests.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Test runner for the WGSL template Python port. + +Discovers ``test_*.py`` siblings and aggregates them into one suite. +Invoked manually or via ``ctest`` (see CMake's ``add_test`` wiring). +""" + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +# Make the wgsl_template package importable regardless of cwd. +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + + +def load_tests(loader, standard_tests, pattern): + discovered = loader.discover( + start_dir=str(_THIS_DIR), + pattern="test_*.py", + top_level_dir=str(_THIS_DIR), + ) + standard_tests.addTests(discovered) + return standard_tests + + +if __name__ == "__main__": # pragma: no cover + unittest.main(module=__name__, verbosity=2) diff --git a/tools/python/wgsl_template/test/test_build.py b/tools/python/wgsl_template/test/test_build.py new file mode 100644 index 0000000000000..78995d344fd01 --- /dev/null +++ b/tools/python/wgsl_template/test/test_build.py @@ -0,0 +1,261 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""End-to-end tests for the top-level build() orchestrator. + +Combines focused unit tests with a fixture-driven runner over the +``build-*`` cases in ``testcases/``. Generated files are compared +against the ``expected//`` golden tree using +``canonicalize`` (content-equivalence, not byte-for-byte), so +host-dependent ``__str_N`` numbering / ordering / sha256 markers don't +cause spurious failures. See ``golden_compare.py``. +""" + +from __future__ import annotations + +import json +import os +import sys +import tempfile +import unittest +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + +from golden_compare import canonicalize, read_tree # noqa: E402 +from wgsl_template import build # noqa: E402 +from wgsl_template.errors import WgslTemplateError # noqa: E402 +from wgsl_template.types import SourceDir # noqa: E402 + +_TESTCASES_DIR = _THIS_DIR / "testcases" + +_SUPPORTED_GENERATORS = {"static-cpp", "static-cpp-literal"} + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content.encode("utf-8")) + + +def _walk_files(root: Path) -> list[str]: + """Return every file under ``root`` as a sorted POSIX path + relative to ``root``.""" + out: list[str] = [] + for dirpath, _dirs, files in os.walk(root): + for f in files: + full = Path(dirpath) / f + out.append(full.relative_to(root).as_posix()) + out.sort() + return out + + +# ---------------------------------------------------------------------- +# Focused unit tests +# ---------------------------------------------------------------------- + + +class BuildIdempotentWriteTest(unittest.TestCase): + def test_unchanged_file_is_not_rewritten(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + src = root / "src" + out = root / "out" + _write(src / "a.wgsl.template", "a\n") + + # First build. + build( + source_dirs=[src], + out_dir=out, + generator="static-cpp-literal", + ) + mtimes_before = {p: (out / p).stat().st_mtime_ns for p in _walk_files(out)} + self.assertTrue(mtimes_before, "expected files in output dir") + + # Second build with identical sources — files must not be rewritten. + build( + source_dirs=[src], + out_dir=out, + generator="static-cpp-literal", + ) + mtimes_after = {p: (out / p).stat().st_mtime_ns for p in _walk_files(out)} + self.assertEqual(mtimes_before, mtimes_after) + + def test_changed_file_is_rewritten(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + src = root / "src" + out = root / "out" + _write(src / "a.wgsl.template", "a\n") + + build(source_dirs=[src], out_dir=out, generator="static-cpp-literal") + content_before = (out / "generated/a.h").read_bytes() + + _write(src / "a.wgsl.template", "completely different\n") + build(source_dirs=[src], out_dir=out, generator="static-cpp-literal") + content_after = (out / "generated/a.h").read_bytes() + + self.assertNotEqual(content_before, content_after) + + +class BuildCleanFlagTest(unittest.TestCase): + def test_clean_removes_existing_files(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + src = root / "src" + out = root / "out" + _write(src / "a.wgsl.template", "a\n") + _write(out / "stale.h", "STALE\n") + + build(source_dirs=[src], out_dir=out, generator="static-cpp-literal", clean=True) + self.assertFalse((out / "stale.h").exists()) + self.assertTrue((out / "generated/a.h").exists()) + + +class BuildErrorsTest(unittest.TestCase): + def test_empty_source_dirs_raises(self) -> None: + with self.assertRaises(WgslTemplateError): + build(source_dirs=[], out_dir="/tmp/anywhere") + + +# ---------------------------------------------------------------------- +# Fixture-driven tests against the build-* cases +# ---------------------------------------------------------------------- + + +def _build_fixture_suite() -> unittest.TestSuite: + suite = unittest.TestSuite() + if not _TESTCASES_DIR.is_dir(): + return suite + + for entry in sorted(os.listdir(_TESTCASES_DIR)): + case_dir = _TESTCASES_DIR / entry + if not case_dir.is_dir(): + continue + if not entry.startswith("build-"): + continue + config_path = case_dir / "test-config.json" + if not config_path.is_file(): + continue + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + if config.get("disabled"): + continue + if config.get("type") != "build": + continue + + # Filter generators down to ones we support. + gen_cfg = config.get("generators") or {} + applicable = {name: cfg for name, cfg in gen_cfg.items() if name in _SUPPORTED_GENERATORS} + if not applicable: + continue + + case_class = _make_build_case(case_dir, config, applicable) + suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(case_class)) + return suite + + +def _make_build_case(case_dir: Path, config: dict, applicable: dict) -> type: + case_name = case_dir.name + template_ext = config.get("templateExt") or ".wgsl.template" + source_dirs_cfg = config.get("sourceDirs") + + class _Case(unittest.TestCase): + def __init__(self, methodName: str = "runTest") -> None: + super().__init__(methodName) + self._dir = case_dir + self._template_ext = template_ext + self._source_dirs_cfg = source_dirs_cfg + self._applicable = applicable + + def __str__(self) -> str: + return f"fixture[{case_name}]" + + def runTest(self) -> None: # noqa: N802 + src_dir = self._dir / "src" + expected_dir = self._dir / "expected" + self.assertTrue(src_dir.is_dir(), f"missing src dir: {src_dir}") + self.assertTrue(expected_dir.is_dir(), f"missing expected dir: {expected_dir}") + + for gen_name, gen_cfg in self._applicable.items(): + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) + + if self._source_dirs_cfg: + sources: list = [] + for d in self._source_dirs_cfg: + if isinstance(d, str): + sources.append(src_dir / d) + else: + sources.append( + SourceDir( + path=str(src_dir / d["path"]), + alias=d.get("alias"), + ) + ) + else: + sources = [src_dir] + + expects_error = gen_cfg.get("expectsError") + try: + build( + source_dirs=sources, + out_dir=out_dir, + template_ext=self._template_ext, + generator=gen_name, + ) + except Exception as e: + if expects_error: + if isinstance(expects_error, str): + self.assertIn(expects_error, str(e)) + return + raise + + if expects_error: + self.fail(f"{gen_name}: expected error containing {expects_error!r} but build succeeded") + + expected_gen = expected_dir / gen_name + self.assertTrue( + expected_gen.is_dir(), + f"missing expected/{gen_name} dir for {case_name}", + ) + + actual_files = _walk_files(out_dir) + expected_files = _walk_files(expected_gen) + + self.assertEqual( + actual_files, + expected_files, + f"{case_name} ({gen_name}): file set mismatch\n" + f"actual: {actual_files}\n" + f"expected: {expected_files}", + ) + + # Compare in canonical form (see golden_compare). + actual_canon = canonicalize(read_tree(out_dir)) + expected_canon = canonicalize(read_tree(expected_gen)) + + for rel in expected_files: + self.assertEqual( + actual_canon[rel], + expected_canon[rel], + f"{case_name} ({gen_name}): {rel} differs (after canonicalization)\n" + f"actual:\n{actual_canon[rel]!r}\n" + f"expected:\n{expected_canon[rel]!r}", + ) + + _Case.__name__ = f"BuildFixture_{case_name.replace('-', '_')}" + _Case.__qualname__ = _Case.__name__ + return _Case + + +def load_tests(loader, standard_tests, pattern): + standard_tests.addTests(_build_fixture_suite()) + return standard_tests + + +if __name__ == "__main__": # pragma: no cover + unittest.main(verbosity=2) diff --git a/tools/python/wgsl_template/test/test_generator.py b/tools/python/wgsl_template/test/test_generator.py new file mode 100644 index 0000000000000..2ee2e33c69a32 --- /dev/null +++ b/tools/python/wgsl_template/test/test_generator.py @@ -0,0 +1,297 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Unit tests for the PASS2 generator. + +Combines focused unit tests with a fixture-driven runner that pulls +``generator-*`` cases from the ``testcases/`` fixtures next to this +file. We only fixture-test cases where: + +* The fixture directory contains a ``*.static-cpp-literal.gen`` + golden, OR +* The case expects an error (``expectsError`` in test-config.json) + with at least one of its generators being one we support. +""" + +from __future__ import annotations + +import json +import os +import re +import sys +import tempfile +import unittest +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + +from wgsl_template.code_generator import resolve_code_generator # noqa: E402 +from wgsl_template.errors import WgslTemplateGenerateError # noqa: E402 +from wgsl_template.generator import generate # noqa: E402 +from wgsl_template.loader import load_from_directory # noqa: E402 +from wgsl_template.parser import parse # noqa: E402 + +# Fixtures live next to this file. Resolve relative to __file__ so the +# suite runs on any platform and from any working directory. +_TESTCASES_DIR = _THIS_DIR / "testcases" + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content.encode("utf-8")) + + +def _gen(text: str, *, generator: str = "static-cpp-literal", preserve: bool = False) -> str: + """Run a single template through PASS0+PASS1+PASS2, returning the + emitted C++ string.""" + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "test.wgsl.template", text) + repo = load_from_directory(root) + parsed = parse(repo) + cg = resolve_code_generator(generator) + return generate( + "test.wgsl.template", + parsed, + cg, + preserve_code_reference=preserve, + ).code + + +# ---------------------------------------------------------------------- +# Focused unit tests +# ---------------------------------------------------------------------- + + +class GeneratorBasicTest(unittest.TestCase): + def test_simple_passthrough(self) -> None: + # No directives, no patterns — every line becomes a code segment. + out = _gen("a\nb\n") + self.assertEqual(out, 'ss << "a\\nb\\n";\n') + + def test_param_declaration(self) -> None: + out = _gen("#param X\nlet a = X;\n") + # X expands to __param_X expression. + self.assertIn("__param_X", out) + + def test_main_block(self) -> None: + out = _gen("$MAIN {\n a;\n}\n") + self.assertIn("MainFunctionStart();", out) + self.assertIn("MainFunctionEnd();", out) + + +class GeneratorIfTest(unittest.TestCase): + def test_if_else(self) -> None: + out = _gen("#param X\n#if X\nyes\n#else\nno\n#endif\n") + self.assertIn("if (__param_X) {", out) + self.assertIn("} else {", out) + self.assertIn("}\n", out) + + def test_if_elif(self) -> None: + out = _gen("#param X\n#if X\na\n#elif X\nb\n#endif\n") + self.assertIn("if (__param_X) {", out) + self.assertIn("} else if (__param_X) {", out) + + def test_orphan_endif(self) -> None: + with self.assertRaises(WgslTemplateGenerateError): + _gen("#endif\n") + + def test_orphan_else(self) -> None: + with self.assertRaises(WgslTemplateGenerateError): + _gen("#else\n") + + def test_unclosed_if(self) -> None: + with self.assertRaises(WgslTemplateGenerateError): + _gen("#param X\n#if X\nbody\n") + + +class GeneratorPropertyTest(unittest.TestCase): + def test_rank_property(self) -> None: + out = _gen("#use .rank\nlet n = output.rank;\n") + self.assertIn("__var_output.Rank()", out) + + def test_method_call(self) -> None: + out = _gen("#use .offsetToIndices\nlet i = output.offsetToIndices(j);\n") + self.assertIn("__var_output.OffsetToIndices", out) + + +class GeneratorFunctionTest(unittest.TestCase): + def test_get_element_at(self) -> None: + out = _gen("#use getElementAt\nlet x = getElementAt(a, b, c);\n") + self.assertIn("GetElementAt", out) + + +class GeneratorParamErrorsTest(unittest.TestCase): + def test_invalid_param_name_with_dot(self) -> None: + with self.assertRaises(ValueError): + _gen("#param a.b\n") + + def test_invalid_param_starts_with_number(self) -> None: + with self.assertRaises(ValueError): + _gen("#param 9bad\n") + + +# ---------------------------------------------------------------------- +# Fixture-driven tests against the generator-* testcases +# ---------------------------------------------------------------------- + + +_SUPPORTED_GENERATORS = {"static-cpp", "static-cpp-literal"} + +# No fixtures are currently skipped. +_FIXTURE_SKIPS: set[str] = set() + + +def _normalize_lines(text: str) -> list: + """Split on newlines, drop blank lines, trim whitespace per line.""" + return [line.strip() for line in re.split(r"\r?\n", text) if line.strip() != ""] + + +def _build_fixture_suite() -> unittest.TestSuite: + suite = unittest.TestSuite() + if not _TESTCASES_DIR.is_dir(): + return suite + + for entry in sorted(os.listdir(_TESTCASES_DIR)): + case_dir = _TESTCASES_DIR / entry + if not case_dir.is_dir(): + continue + if not entry.startswith("generator-"): + continue + if entry in _FIXTURE_SKIPS: + continue + config_path = case_dir / "test-config.json" + if not config_path.is_file(): + continue + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + if config.get("disabled"): + continue + if config.get("type") != "generator": + continue + + entries = config.get("entries", {}) + applicable_entries = [] + for template_path, entry_cfg in entries.items(): + generators_cfg = entry_cfg.get("generators", {}) or {} + for gen_name, gen_cfg in generators_cfg.items(): + if gen_name not in _SUPPORTED_GENERATORS: + continue + # Only require .gen golden if no expectsError. + expects_error = gen_cfg.get("expectsError") or config.get("expectsError") + golden_path = case_dir / f"{template_path}.{gen_name}.gen" + if not expects_error and not golden_path.is_file(): + continue + applicable_entries.append( + { + "template_path": template_path, + "generator": gen_name, + "expects_error": expects_error, + "expected_params": gen_cfg.get("expectedParams"), + "expected_variables": gen_cfg.get("expectedVariables"), + "golden_path": golden_path, + } + ) + + if not applicable_entries: + continue + + case_class = _make_generator_case(case_dir, config, applicable_entries) + suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(case_class)) + + return suite + + +def _make_generator_case(case_dir: Path, config: dict, applicable_entries: list) -> type: + case_name = case_dir.name + preserve_code_reference = bool(config.get("preserveCodeReference")) + + class _Case(unittest.TestCase): + def __init__(self, methodName: str = "runTest") -> None: + super().__init__(methodName) + self._dir = case_dir + self._entries = applicable_entries + self._preserve = preserve_code_reference + + def __str__(self) -> str: + return f"fixture[{case_name}]" + + def runTest(self) -> None: # noqa: N802 + repo = load_from_directory(self._dir) + parsed = parse(repo) + for entry in self._entries: + self._run_entry(parsed, entry) + + def _run_entry(self, parsed, entry: dict) -> None: + cg = resolve_code_generator(entry["generator"]) + try: + result = generate( + entry["template_path"], + parsed, + cg, + preserve_code_reference=self._preserve, + ) + except Exception as e: + if entry["expects_error"]: + if isinstance(entry["expects_error"], str): + self.assertIn(entry["expects_error"], str(e)) + return + raise + + if entry["expects_error"]: + self.fail( + f"{entry['template_path']} ({entry['generator']}): " + f"expected error containing " + f"{entry['expects_error']!r} but generation succeeded" + ) + + golden_text = entry["golden_path"].read_text(encoding="utf-8") + actual_lines = _normalize_lines(result.code) + expected_lines = _normalize_lines(golden_text) + self.assertEqual( + len(actual_lines), + len(expected_lines), + f"{entry['template_path']} ({entry['generator']}): " + f"line count mismatch; expected {len(expected_lines)} " + f"got {len(actual_lines)}\n" + f"--- expected ---\n" + "\n".join(expected_lines) + "\n" + "--- actual ---\n" + "\n".join(actual_lines), + ) + for i, (a, b) in enumerate(zip(actual_lines, expected_lines, strict=True)): + self.assertEqual( + a, + b, + f"{entry['template_path']} ({entry['generator']}): " + f"line {i + 1} mismatch\n" + f" expected: {b!r}\n" + f" actual: {a!r}", + ) + + if entry["expected_params"] is not None: + self.assertEqual( + sorted(result.params.keys()), + sorted(entry["expected_params"]), + ) + if entry["expected_variables"] is not None: + self.assertEqual( + sorted(result.variables.keys()), + sorted(entry["expected_variables"]), + ) + + _Case.__name__ = f"Fixture_{case_name.replace('-', '_')}" + _Case.__qualname__ = _Case.__name__ + return _Case + + +def load_tests(loader, standard_tests, pattern): + standard_tests.addTests(_build_fixture_suite()) + return standard_tests + + +if __name__ == "__main__": # pragma: no cover + unittest.main(verbosity=2) diff --git a/tools/python/wgsl_template/test/test_in_tree_smoke.py b/tools/python/wgsl_template/test/test_in_tree_smoke.py new file mode 100644 index 0000000000000..38e33aa1e2be2 --- /dev/null +++ b/tools/python/wgsl_template/test/test_in_tree_smoke.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Golden tests against the real in-tree wgsl.template files. + +These templates back shipped WebGPU operators (Pad, Transpose, +im2col-matmul, and others under core/providers/webgpu). The generated +C++ is compared against a committed golden snapshot so generator +regressions are caught automatically rather than relying on a manual +GPU run. The comparison uses ``canonicalize`` (see golden_compare). + +When an in-tree template changes (or one is added), regenerate the +goldens and commit them:: + + UPDATE_WGSL_GOLDEN=1 python wgsl_template/test/run_tests.py +""" + +from __future__ import annotations + +import os +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + +from golden_compare import canonicalize, read_tree # noqa: E402 +from wgsl_template import build # noqa: E402 + +# Repo root: tools/python/wgsl_template/test/ -> ../../../.. +_REPO_ROOT = _THIS_DIR.parent.parent.parent.parent +# WebGPU EP source root, where the real templates live. +_WEBGPU_ROOT = _REPO_ROOT / "onnxruntime" / "core" / "providers" / "webgpu" +# Committed golden snapshots of the generated output. +_GOLDEN_DIR = _THIS_DIR / "in_tree_golden" + +# Set UPDATE_WGSL_GOLDEN=1 to (re)write the golden snapshots instead of +# comparing against them. +_UPDATE_GOLDEN = os.environ.get("UPDATE_WGSL_GOLDEN") == "1" + +_REGEN_HINT = "UPDATE_WGSL_GOLDEN=1 python wgsl_template/test/run_tests.py" + + +def _generate(generator: str) -> Path: + out = Path(tempfile.mkdtemp(prefix="wgsl_smoke_")) + try: + build( + source_dirs=[_WEBGPU_ROOT], + out_dir=out, + generator=generator, + include_path_prefix="wgsl_template_gen/", + preserve_code_reference=True, + ) + except Exception: + shutil.rmtree(out, ignore_errors=True) + raise + return out + + +def _write_golden(generator: str, out_dir: Path) -> None: + dest = _GOLDEN_DIR / generator + if dest.exists(): + shutil.rmtree(dest) + shutil.copytree(out_dir, dest) + + +class InTreeTemplatesGoldenTest(unittest.TestCase): + def _check(self, generator: str) -> None: + out = _generate(generator) + try: + if _UPDATE_GOLDEN: + _write_golden(generator, out) + self.skipTest(f"updated golden for {generator}") + return + + golden = _GOLDEN_DIR / generator + self.assertTrue( + golden.is_dir(), + f"missing golden for {generator}; run `{_REGEN_HINT}` to create it", + ) + + actual = canonicalize(read_tree(out)) + expected = canonicalize(read_tree(golden)) + + self.assertEqual( + sorted(actual), + sorted(expected), + f"{generator}: generated file set differs from golden. If an in-tree " + f"template was added/removed, regenerate with `{_REGEN_HINT}`.", + ) + for rel in sorted(expected): + self.assertEqual( + actual[rel], + expected[rel], + f"{generator}: {rel} differs from golden. If the in-tree template " + f"changed intentionally, regenerate with `{_REGEN_HINT}`.", + ) + finally: + shutil.rmtree(out, ignore_errors=True) + + def test_static_cpp_matches_golden(self) -> None: + self._check("static-cpp") + + def test_static_cpp_literal_matches_golden(self) -> None: + self._check("static-cpp-literal") diff --git a/tools/python/wgsl_template/test/test_loader.py b/tools/python/wgsl_template/test/test_loader.py new file mode 100644 index 0000000000000..dbf7bfbb4ec29 --- /dev/null +++ b/tools/python/wgsl_template/test/test_loader.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Unit tests for the PASS0 loader. + +Uses temporary directories with inline fixtures. +""" + +from __future__ import annotations + +import sys +import tempfile +import unittest +from pathlib import Path + +# Make the wgsl_template package importable regardless of cwd. +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + +from wgsl_template.errors import WgslTemplateLoadError # noqa: E402 +from wgsl_template.loader import ( # noqa: E402 + load_from_directories, + load_from_directory, +) +from wgsl_template.types import SourceDir, TemplatePass0 # noqa: E402 + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + # Use binary write so callers control line endings exactly. + path.write_bytes(content.encode("utf-8")) + + +class LoaderBasicTest(unittest.TestCase): + def test_loads_single_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write( + root / "basic.wgsl.template", + "fn a() {\n var x = 5;\n var y = 10;\n}\n", + ) + + repo = load_from_directory(root) + + self.assertEqual(set(repo.templates), {"basic.wgsl.template"}) + tpl = repo.templates["basic.wgsl.template"] + assert isinstance(tpl, TemplatePass0) + # Trailing \n means split produces an empty trailing entry. + self.assertEqual( + tpl.raw, + ["fn a() {", " var x = 5;", " var y = 10;", "}", ""], + ) + + def test_empty_directory_returns_empty_repo(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + repo = load_from_directory(tmp) + self.assertEqual(repo.templates, {}) + + def test_recurses_into_subdirectories(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "main.wgsl.template", "main\n") + _write(root / "shaders" / "vertex.wgsl.template", "vertex\n") + _write(root / "shaders" / "compute" / "matrix.wgsl.template", "matrix\n") + + repo = load_from_directory(root) + + self.assertEqual( + set(repo.templates), + { + "main.wgsl.template", + "shaders/vertex.wgsl.template", + "shaders/compute/matrix.wgsl.template", + }, + ) + + def test_template_keys_are_sorted(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + for name in ["zeta.wgsl.template", "alpha.wgsl.template", "mu.wgsl.template"]: + _write(root / name, "x\n") + + repo = load_from_directory(root) + + self.assertEqual( + list(repo.templates), + ["alpha.wgsl.template", "mu.wgsl.template", "zeta.wgsl.template"], + ) + + def test_skips_files_with_other_extensions(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "keep.wgsl.template", "keep\n") + _write(root / "drop.txt", "drop\n") + _write(root / "drop.wgsl", "drop\n") + + repo = load_from_directory(root) + + self.assertEqual(set(repo.templates), {"keep.wgsl.template"}) + + def test_uses_posix_paths_in_keys(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "a" / "b" / "c.wgsl.template", "x\n") + + repo = load_from_directory(root) + + self.assertEqual(set(repo.templates), {"a/b/c.wgsl.template"}) + + +class LoaderLineEndingsTest(unittest.TestCase): + def test_lf(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "lf.wgsl.template", "line1\nline2\n") + + repo = load_from_directory(root) + self.assertEqual(repo.templates["lf.wgsl.template"].raw, ["line1", "line2", ""]) + + def test_crlf(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "crlf.wgsl.template", "line1\r\nline2\r\n") + + repo = load_from_directory(root) + self.assertEqual(repo.templates["crlf.wgsl.template"].raw, ["line1", "line2", ""]) + + def test_mixed(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "mixed.wgsl.template", "a\r\nb\nc\r\n") + + repo = load_from_directory(root) + self.assertEqual(repo.templates["mixed.wgsl.template"].raw, ["a", "b", "c", ""]) + + +class LoaderExtensionTest(unittest.TestCase): + def test_custom_extension(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "shader.wgsl.custom", "x\n") + _write(root / "ignored.wgsl.template", "x\n") + + repo = load_from_directory(root, ext=".wgsl.custom") + self.assertEqual(set(repo.templates), {"shader.wgsl.custom"}) + + +class LoaderMultipleDirectoriesTest(unittest.TestCase): + def test_loads_from_multiple_dirs_without_aliases(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + d1, d2 = root / "d1", root / "d2" + _write(d1 / "a.wgsl.template", "a\n") + _write(d2 / "b.wgsl.template", "b\n") + + repo = load_from_directories([d1, d2]) + self.assertEqual(set(repo.templates), {"a.wgsl.template", "b.wgsl.template"}) + + def test_loads_with_aliases(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + d1, d2 = root / "templates", root / "shaders" + _write(d1 / "base.wgsl.template", "base\n") + _write(d2 / "vertex.wgsl.template", "vertex\n") + + repo = load_from_directories( + [ + SourceDir(path=str(d1), alias="@templates"), + SourceDir(path=str(d2), alias="@shaders"), + ] + ) + + self.assertEqual( + set(repo.templates), + { + "@templates/base.wgsl.template", + "@shaders/vertex.wgsl.template", + }, + ) + + def test_template_name_conflict_raises(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + d1, d2 = root / "d1", root / "d2" + _write(d1 / "shared.wgsl.template", "v1\n") + _write(d2 / "shared.wgsl.template", "v2\n") + + with self.assertRaises(WgslTemplateLoadError) as ctx: + load_from_directories([d1, d2]) + self.assertIn("Template name conflict", str(ctx.exception)) + + +class LoaderErrorsTest(unittest.TestCase): + def test_missing_directory_raises(self) -> None: + with self.assertRaises(WgslTemplateLoadError) as ctx: + load_from_directory("/no/such/path/does/not/exist/abc123") + self.assertIn("Cannot access directory", str(ctx.exception)) + + def test_path_is_a_file_raises(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + file_path = root / "notadir" + file_path.write_text("x") + with self.assertRaises(WgslTemplateLoadError) as ctx: + load_from_directory(file_path) + self.assertIn("not a directory", str(ctx.exception)) + + def test_no_directories_raises(self) -> None: + with self.assertRaises(WgslTemplateLoadError): + load_from_directories([]) + + +if __name__ == "__main__": # pragma: no cover + unittest.main(verbosity=2) diff --git a/tools/python/wgsl_template/test/test_parser.py b/tools/python/wgsl_template/test/test_parser.py new file mode 100644 index 0000000000000..1ebc2d7d47a11 --- /dev/null +++ b/tools/python/wgsl_template/test/test_parser.py @@ -0,0 +1,313 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Unit tests for the PASS1 parser. + +Mixes inline temp-dir tests for tightly scoped behaviors with a +fixture-driven runner that pulls ``parser-*`` and ``loader-*`` cases +from the ``testcases/`` fixtures next to this file. +""" + +from __future__ import annotations + +import json +import os +import sys +import tempfile +import unittest +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_PARENT_DIR = _THIS_DIR.parent.parent +if str(_PARENT_DIR) not in sys.path: + sys.path.insert(0, str(_PARENT_DIR)) + +from wgsl_template.errors import WgslTemplateParseError # noqa: E402 +from wgsl_template.loader import load_from_directory # noqa: E402 +from wgsl_template.parser import parse # noqa: E402 +from wgsl_template.types import TemplatePass1 # noqa: E402 + +# Fixtures live next to this file. Resolve relative to __file__ so the +# suite runs on any platform and from any working directory. +_TESTCASES_DIR = _THIS_DIR / "testcases" + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content.encode("utf-8")) + + +def _parse_text(text: str) -> list[str]: + """Helper: feed a single template through PASS0+PASS1, return the + resulting line strings.""" + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write(root / "test.wgsl.template", text) + repo = load_from_directory(root) + parsed = parse(repo) + tpl = parsed.templates["test.wgsl.template"] + assert isinstance(tpl, TemplatePass1) + return [line.line for line in tpl.pass1] + + +# ---------------------------------------------------------------------- +# Comment stripping (focused unit tests) +# ---------------------------------------------------------------------- + + +class ParserCommentsTest(unittest.TestCase): + def test_strips_single_line_comments(self) -> None: + self.assertEqual( + _parse_text("a // comment\nb\n"), + ["a", "b", ""], + ) + + def test_strips_multiline_comment_inline(self) -> None: + self.assertEqual( + _parse_text("a /* x */ b\n"), + ["a b", ""], + ) + + def test_strips_multiline_comment_spanning_lines(self) -> None: + # Leading whitespace on the post-`*/` portion is preserved + # (only trailing whitespace gets trimmed). + self.assertEqual( + _parse_text("a /* x\ny\nz */ b\n"), + ["a", "", " b", ""], + ) + + def test_preserves_line_count(self) -> None: + # Three input lines should still be three output lines. + self.assertEqual( + _parse_text("// a\n// b\n// c\n"), + ["", "", "", ""], + ) + + def test_unterminated_multiline_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError): + _parse_text("a /* never ends\n") + + +# ---------------------------------------------------------------------- +# #include expansion (focused unit tests) +# ---------------------------------------------------------------------- + + +class ParserIncludeTest(unittest.TestCase): + def _run(self, files: dict, top: str) -> list[str]: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + for name, content in files.items(): + _write(root / name, content) + repo = load_from_directory(root) + parsed = parse(repo) + tpl = parsed.templates[top] + assert isinstance(tpl, TemplatePass1) + return [line.line for line in tpl.pass1] + + def test_simple_include(self) -> None: + result = self._run( + { + "main.wgsl.template": '#include "utils.wgsl.template"\nmain\n', + "utils.wgsl.template": "u1\nu2\n", + }, + "main.wgsl.template", + ) + # Trailing empty line from the LF terminator is preserved. + self.assertEqual(result, ["u1", "u2", "", "main", ""]) + + def test_circular_include_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + self._run( + { + "a.wgsl.template": '#include "b.wgsl.template"\n', + "b.wgsl.template": '#include "a.wgsl.template"\n', + }, + "a.wgsl.template", + ) + self.assertIn("Circular #include", str(ctx.exception)) + + def test_self_include_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + self._run( + {"a.wgsl.template": '#include "a.wgsl.template"\n'}, + "a.wgsl.template", + ) + self.assertIn("Circular #include", str(ctx.exception)) + + def test_missing_include_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + self._run( + {"a.wgsl.template": '#include "missing.wgsl.template"\n'}, + "a.wgsl.template", + ) + self.assertIn("not found", str(ctx.exception)) + + def test_unquoted_include_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + self._run( + {"a.wgsl.template": "#include foo.wgsl.template\n"}, + "a.wgsl.template", + ) + self.assertIn("double quotes", str(ctx.exception)) + + +# ---------------------------------------------------------------------- +# #define substitution (focused unit tests) +# ---------------------------------------------------------------------- + + +class ParserMacroTest(unittest.TestCase): + def test_basic_substitution(self) -> None: + self.assertEqual( + _parse_text("#define X 42\nlet a = X;\n"), + ["", "let a = 42;", ""], + ) + + def test_whole_word_match(self) -> None: + # Should NOT replace inside identifiers. + self.assertEqual( + _parse_text("#define X 42\nlet Xy = 1;\nlet X = 2;\n"), + ["", "let Xy = 1;", "let 42 = 2;", ""], + ) + + def test_eager_expansion(self) -> None: + # Y is defined first as 10. Then X uses Y, gets stored as 10. + # Subsequent code uses X, which expands to 10. + self.assertEqual( + _parse_text("#define Y 10\n#define X Y\nlet a = X;\n"), + ["", "", "let a = 10;", ""], + ) + + def test_self_reference_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + _parse_text("#define X X+1\n") + self.assertIn("references itself", str(ctx.exception)) + + def test_duplicate_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + _parse_text("#define X 1\n#define X 2\n") + self.assertIn("Duplicate macro", str(ctx.exception)) + + def test_empty_value_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + _parse_text("#define X\n") + self.assertIn("no value", str(ctx.exception)) + + def test_invalid_name_raises(self) -> None: + with self.assertRaises(WgslTemplateParseError) as ctx: + _parse_text("#define 9bad value\n") + self.assertIn("invalid macro name", str(ctx.exception)) + + +# ---------------------------------------------------------------------- +# Fixture-driven tests against the parser-* and loader-* testcases +# ---------------------------------------------------------------------- + + +def _build_fixture_suite() -> unittest.TestSuite: + """Discover ``parser-*`` and ``loader-*`` fixture directories with + ``.pass1`` golden files and build a unittest suite. + + Returns an empty suite if the fixtures aren't present, so the run + stays green even in a trimmed checkout. + """ + + suite = unittest.TestSuite() + if not _TESTCASES_DIR.is_dir(): + return suite + + for entry in sorted(os.listdir(_TESTCASES_DIR)): + case_dir = _TESTCASES_DIR / entry + if not case_dir.is_dir(): + continue + if not entry.startswith(("parser-", "loader-")): + continue + config_path = case_dir / "test-config.json" + if not config_path.is_file(): + continue + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + if config.get("disabled"): + continue + # Skip loader-directories cases — they don't exercise PASS1 + # and need a different harness. + if config.get("type") != "parser": + continue + + case_class = _make_case(case_dir, config) + loader = unittest.defaultTestLoader + suite.addTests(loader.loadTestsFromTestCase(case_class)) + + return suite + + +def _make_case(case_dir: Path, config: dict) -> type: + """Build a one-off TestCase class for a single fixture directory.""" + + case_name = case_dir.name + expects_error = config.get("expectsError") + + class _Case(unittest.TestCase): + # Shorter id so `-v` output is readable. + def __init__(self, methodName: str = "runTest") -> None: + super().__init__(methodName) + self._dir = case_dir + self._expects_error = expects_error + + def __str__(self) -> str: + return f"fixture[{case_name}]" + + def runTest(self) -> None: # noqa: N802 + try: + repo = load_from_directory(self._dir) + parsed = parse(repo) + except Exception as e: + if self._expects_error: + if isinstance(self._expects_error, str): + self.assertIn(self._expects_error, str(e)) + return + raise + + if self._expects_error: + self.fail(f"Expected error containing {self._expects_error!r} but parse() succeeded") + + # Compare each template's pass1 lines against its .pass1 golden. + for template_key, template in parsed.templates.items(): + assert isinstance(template, TemplatePass1) + golden_path = self._dir / f"{template_key}.pass1" + if not golden_path.is_file(): + self.fail(f"Missing golden file for {template_key} at {golden_path}") + expected_text = golden_path.read_text(encoding="utf-8") + # Normalize line endings to match the loader's split. + expected_lines = expected_text.replace("\r\n", "\n").split("\n") + actual_lines = [line.line for line in template.pass1] + self.assertEqual( + len(actual_lines), + len(expected_lines), + f"Line count mismatch for {template_key}: expected " + f"{len(expected_lines)} got {len(actual_lines)}\n" + f"actual: {actual_lines!r}\n" + f"expected: {expected_lines!r}", + ) + for i, (a, b) in enumerate(zip(actual_lines, expected_lines, strict=True)): + self.assertEqual( + a, + b, + f"Line {i + 1} mismatch in {template_key}: actual={a!r} expected={b!r}", + ) + + _Case.__name__ = f"Fixture_{case_name.replace('-', '_')}" + _Case.__qualname__ = _Case.__name__ + return _Case + + +def load_tests(loader, standard_tests, pattern): + standard_tests.addTests(_build_fixture_suite()) + return standard_tests + + +if __name__ == "__main__": # pragma: no cover + unittest.main(verbosity=2) diff --git a/tools/python/wgsl_template/test/testcases/.gitattributes b/tools/python/wgsl_template/test/testcases/.gitattributes new file mode 100644 index 0000000000000..a6a813b4c3f2a --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/.gitattributes @@ -0,0 +1,4 @@ +# Preserve line endings for test expectation files +**/expected/** text eol=lf +*.gen text eol=lf +*.pass1 text eol=lf diff --git a/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/generated/shader/triangle.h b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/generated/shader/triangle.h new file mode 100644 index 0000000000000..d41bcf4bb5e20 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/generated/shader/triangle.h @@ -0,0 +1,16 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: shader/triangle.wgsl.template + +template <> +Status ApplyTemplate<"shader/triangle.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"shader/triangle.wgsl.template">::type ) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + +ss << __str_0; + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index.h b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index.h new file mode 100644 index 0000000000000..296c8e3f8ac4e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index.h @@ -0,0 +1,18 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#ifndef INCLUDED_BY_WGSL_GEN_HEADER +#error "This file is expected to be included by wgsl-gen header. Do not include it directly." +#endif + +// +// Template: shader/triangle.wgsl.template +// + +template <> +struct TemplateParameter<"shader/triangle.wgsl.template"> { + using type = struct { + }; +}; + +template <> +Status ApplyTemplate<"shader/triangle.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"shader/triangle.wgsl.template">::type params); diff --git a/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index_impl.h b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index_impl.h new file mode 100644 index 0000000000000..fdf862f4442ff --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/index_impl.h @@ -0,0 +1,41 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// Helper functions or macros + +#pragma push_macro("MainFunctionStart") +#undef MainFunctionStart +#define MainFunctionStart() { [[maybe_unused]] auto& ss = shader_helper.MainFunctionBody(); +#pragma push_macro("MainFunctionEnd") +#undef MainFunctionEnd +#define MainFunctionEnd() } + +// Helper templates + +namespace wgsl_detail { +template >> +std::string pass_as_string(T&& v) { + return std::to_string(std::forward(v)); +} +template +std::string_view pass_as_string(std::string_view sv) { + return sv; +} +template +std::string pass_as_string(T&& v) { + return std::forward(v); +} +} // namespace wgsl_detail + +#include "string_table.h" // 16d2563fc58e81b41c9b3be4a1c9a9e0a4458c7ff6849b4359f77a6ec7f3d23c + +// Include template implementations + +#include "generated/shader/triangle.h" // d61822d677fc73d87adbe0b7b8e81d556ca515fef270bbe01dcdcb10596c350b + +#pragma pop_macro("MainFunctionStart") +#pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/string_table.h b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/string_table.h new file mode 100644 index 0000000000000..6b24cc2962f07 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/expected/static-cpp/string_table.h @@ -0,0 +1,9 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// String table constants +constexpr const char* __str_0 = "@vertex\nfn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f {\n let pos = array(\n vec2f(-0.5, -0.5),\n vec2f( 0.5, -0.5),\n vec2f( 0.0, 0.5)\n );\n return vec4f(pos[vertex_index], 0.0, 1.0);\n}\n\n@fragment\nfn fs_main() -> @location(0) vec4f {\n return vec4f(1.0, 0.0, 0.0, 1.0);\n}\n"; diff --git a/tools/python/wgsl_template/test/testcases/build-basic/src/shader/triangle.wgsl.template b/tools/python/wgsl_template/test/testcases/build-basic/src/shader/triangle.wgsl.template new file mode 100644 index 0000000000000..65ed8fc9aa0d6 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/src/shader/triangle.wgsl.template @@ -0,0 +1,15 @@ +// Basic WGSL shader template +@vertex +fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f { + let pos = array( + vec2f(-0.5, -0.5), + vec2f( 0.5, -0.5), + vec2f( 0.0, 0.5) + ); + return vec4f(pos[vertex_index], 0.0, 1.0); +} + +@fragment +fn fs_main() -> @location(0) vec4f { + return vec4f(1.0, 0.0, 0.0, 1.0); // Red color +} diff --git a/tools/python/wgsl_template/test/testcases/build-basic/test-config.json b/tools/python/wgsl_template/test/testcases/build-basic/test-config.json new file mode 100644 index 0000000000000..b46ad50feb468 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-basic/test-config.json @@ -0,0 +1,7 @@ +{ + "type": "build", + "description": "Basic build test with simple template", + "generators": { + "static-cpp": {} + } +} diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/blur.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/blur.h new file mode 100644 index 0000000000000..6e4f73a175e9f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/blur.h @@ -0,0 +1,16 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: @effects/blur.wgsl.template + +template <> +Status ApplyTemplate<"@effects/blur.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"@effects/blur.wgsl.template">::type ) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + +ss << __str_1; + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/lighting.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/lighting.h new file mode 100644 index 0000000000000..293dc55a44ac9 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/@effects/lighting.h @@ -0,0 +1,16 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: @effects/lighting.wgsl.template + +template <> +Status ApplyTemplate<"@effects/lighting.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"@effects/lighting.wgsl.template">::type ) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + +ss << __str_2; + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/utils.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/utils.h new file mode 100644 index 0000000000000..6dc6b64bb0b12 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/generated/utils.h @@ -0,0 +1,16 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: utils.wgsl.template + +template <> +Status ApplyTemplate<"utils.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"utils.wgsl.template">::type ) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + +ss << __str_0; + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index.h new file mode 100644 index 0000000000000..6f0a4df72e1b3 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index.h @@ -0,0 +1,44 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#ifndef INCLUDED_BY_WGSL_GEN_HEADER +#error "This file is expected to be included by wgsl-gen header. Do not include it directly." +#endif + +// +// Template: utils.wgsl.template +// + +template <> +struct TemplateParameter<"utils.wgsl.template"> { + using type = struct { + }; +}; + +template <> +Status ApplyTemplate<"utils.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"utils.wgsl.template">::type params); + +// +// Template: @effects/blur.wgsl.template +// + +template <> +struct TemplateParameter<"@effects/blur.wgsl.template"> { + using type = struct { + }; +}; + +template <> +Status ApplyTemplate<"@effects/blur.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"@effects/blur.wgsl.template">::type params); + +// +// Template: @effects/lighting.wgsl.template +// + +template <> +struct TemplateParameter<"@effects/lighting.wgsl.template"> { + using type = struct { + }; +}; + +template <> +Status ApplyTemplate<"@effects/lighting.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"@effects/lighting.wgsl.template">::type params); diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index_impl.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index_impl.h new file mode 100644 index 0000000000000..0b428d5f8275c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/index_impl.h @@ -0,0 +1,43 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// Helper functions or macros + +#pragma push_macro("MainFunctionStart") +#undef MainFunctionStart +#define MainFunctionStart() { [[maybe_unused]] auto& ss = shader_helper.MainFunctionBody(); +#pragma push_macro("MainFunctionEnd") +#undef MainFunctionEnd +#define MainFunctionEnd() } + +// Helper templates + +namespace wgsl_detail { +template >> +std::string pass_as_string(T&& v) { + return std::to_string(std::forward(v)); +} +template +std::string_view pass_as_string(std::string_view sv) { + return sv; +} +template +std::string pass_as_string(T&& v) { + return std::forward(v); +} +} // namespace wgsl_detail + +#include "string_table.h" // 12c86e5439e35b74fef09f3dfdefe7340abe236c2eb1741f2d430ec1e5dcd66f + +// Include template implementations + +#include "generated/utils.h" // 3b5cffd7f2c5668e3a3d4c1115088fa15bdda6685f135e6a6dfafcbfad783f47 +#include "generated/@effects/blur.h" // 0c77005d684e0f007f7e0f68c00a9e7e18f30fbb99178f30db6fe17ed8853e20 +#include "generated/@effects/lighting.h" // 00427800cc55774fbb69509dc4ea99dfbfb40fa6fa0035bb2fa60d1462299ed9 + +#pragma pop_macro("MainFunctionStart") +#pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/string_table.h b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/string_table.h new file mode 100644 index 0000000000000..90a4a4cc71a81 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/expected/static-cpp/string_table.h @@ -0,0 +1,11 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// String table constants +constexpr const char* __str_0 = "@vertex\nfn vs_main() -> @builtin(position) vec4f {\n return vec4f(0.0, 0.0, 0.0, 1.0);\n}\n\n"; +constexpr const char* __str_1 = "fn blur_sample(tex: texture_2d, uv: vec2f) -> vec4f {\n return textureSample(tex, samp, uv);\n}\n"; +constexpr const char* __str_2 = "@fragment\nfn fs_main() -> @location(0) vec4f {\n return vec4f(1.0, 0.0, 0.0, 1.0);\n}\n\n"; diff --git a/tools/python/wgsl_template/test/testcases/build-directories/src/common/utils.wgsl.template b/tools/python/wgsl_template/test/testcases/build-directories/src/common/utils.wgsl.template new file mode 100644 index 0000000000000..7f7415770b5cb --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/src/common/utils.wgsl.template @@ -0,0 +1,7 @@ +// Common utility functions +/*{{ #main vertex }}*/ +@vertex +fn vs_main() -> @builtin(position) vec4f { + return vec4f(0.0, 0.0, 0.0, 1.0); +} +/*{{ /main }}*/ diff --git a/tools/python/wgsl_template/test/testcases/build-directories/src/effects/blur.wgsl.template b/tools/python/wgsl_template/test/testcases/build-directories/src/effects/blur.wgsl.template new file mode 100644 index 0000000000000..d9a8c145b69c6 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/src/effects/blur.wgsl.template @@ -0,0 +1,4 @@ +// Blur effect functions +fn blur_sample(tex: texture_2d, uv: vec2f) -> vec4f { + return textureSample(tex, samp, uv); +} diff --git a/tools/python/wgsl_template/test/testcases/build-directories/src/effects/lighting.wgsl.template b/tools/python/wgsl_template/test/testcases/build-directories/src/effects/lighting.wgsl.template new file mode 100644 index 0000000000000..714d8848c5cb3 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/src/effects/lighting.wgsl.template @@ -0,0 +1,9 @@ +// Effect shader +/*{{ #include "@effects/blur.wgsl.template" }}*/ + +/*{{ #main fragment }}*/ +@fragment +fn fs_main() -> @location(0) vec4f { + return vec4f(1.0, 0.0, 0.0, 1.0); +} +/*{{ /main }}*/ diff --git a/tools/python/wgsl_template/test/testcases/build-directories/test-config.json b/tools/python/wgsl_template/test/testcases/build-directories/test-config.json new file mode 100644 index 0000000000000..e8a164bbd7302 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-directories/test-config.json @@ -0,0 +1,15 @@ +{ + "type": "build", + "description": "Test building templates from multiple directories with aliases", + "templateExt": ".wgsl.template", + "generators": { + "static-cpp": {} + }, + "sourceDirs": [ + "common", + { + "path": "effects", + "alias": "@effects" + } + ] +} diff --git a/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/generated/tensor/pad.h b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/generated/tensor/pad.h new file mode 100644 index 0000000000000..32d03afa8ef67 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/generated/tensor/pad.h @@ -0,0 +1,69 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: tensor/pad.wgsl.template + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_dim_value_zero = params.param_dim_value_zero; + auto& __param_is_float16 = params.param_is_float16; + auto& __param_pad_mode = params.param_pad_mode; + + // Extract variables + auto& __var_output = *params.var_output; + +MainFunctionStart(); +ss << "\n "; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size"); +ss << ";\n\n let constant_value =\n"; +if (__param_is_float16) { +ss << " bitcast>(uniforms.constant_value)[0];\n"; +} else { +ss << " bitcast(uniforms.constant_value);\n"; +} +if (__param_dim_value_zero) { +ss << " output[global_idx] = constant_value;\n"; +} else { +ss << " let output_indices = "; +ss << __var_output.OffsetToIndices("global_idx"); +ss << ";\n var input_index = u32(0);\n var use_pad_value = false;\n var in_coord = i32(0);\n\n for (var dim = 0; dim < "; +ss << __var_output.Rank(); +ss << " && !use_pad_value; dim++) {\n let output_index = i32("; +ss << GetElementAt("output_indices", "dim", __var_output.Rank()); +ss << ");\n let lower_pads = "; +ss << GetElementAt("uniforms.lower_pads", "dim", __var_output.Rank()); +ss << ";\n let data_shape = i32("; +ss << GetElementAt("uniforms.data_shape", "dim", __var_output.Rank()); +ss << ");\n"; +if (__param_pad_mode == 0) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n use_pad_value = true;\n"; +} else if (__param_pad_mode == 2) { +ss << " if (output_index < lower_pads) {\n in_coord = 0;\n } else if (output_index >= data_shape + lower_pads) {\n in_coord = data_shape - 1;\n"; +} else if (__param_pad_mode == 1) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n in_coord = output_index - lower_pads;\n if (in_coord < 0) {\n in_coord = -in_coord;\n }\n let _2n_1 = 2 * (data_shape - 1);\n in_coord = in_coord % _2n_1;\n if (in_coord >= data_shape) {\n in_coord = _2n_1 - in_coord;\n }\n"; +} else { +ss << " if (output_index < lower_pads) {\n in_coord = data_shape + output_index - lower_pads;\n } else if (output_index >= data_shape + lower_pads) {\n in_coord = output_index - data_shape - lower_pads;\n"; +} +ss << " } else {\n in_coord = output_index - lower_pads;\n }\n\n input_index += select(u32(in_coord)\n"; +if (__var_output.Rank() > 1) { +ss << " * "; +ss << GetElementAt("uniforms.data_stride", "dim", __var_output.Rank() - 1); +ss << "\n"; +} +ss << " , u32(in_coord), dim == "; +ss << __var_output.Rank(); +ss << " - 1);\n }\n\n "; +ss << __var_output.SetByOffset("global_idx", "select(data[input_index], constant_value, use_pad_value)"); +ss << ";\n"; +} +MainFunctionEnd(); +ss << "\n"; + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index.h b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index.h new file mode 100644 index 0000000000000..33f5b70504d5f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index.h @@ -0,0 +1,22 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#ifndef INCLUDED_BY_WGSL_GEN_HEADER +#error "This file is expected to be included by wgsl-gen header. Do not include it directly." +#endif + +// +// Template: tensor/pad.wgsl.template +// + +template <> +struct TemplateParameter<"tensor/pad.wgsl.template"> { + using type = struct { + int param_dim_value_zero; + int param_is_float16; + int param_pad_mode; + const ShaderVariableHelper* var_output; + }; +}; + +template <> +Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"tensor/pad.wgsl.template">::type params); diff --git a/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index_impl.h b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index_impl.h new file mode 100644 index 0000000000000..338f13bb9c22b --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-example-pad/expected/static-cpp-literal/index_impl.h @@ -0,0 +1,40 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once +#ifndef INCLUDED_BY_WGSL_GEN_IMPL +#error "This file is expected to be included by wgsl-gen impl. Do not include it directly." +#endif + +// Helper functions or macros + +#pragma push_macro("MainFunctionStart") +#undef MainFunctionStart +#define MainFunctionStart() { [[maybe_unused]] auto& ss = shader_helper.MainFunctionBody(); +#pragma push_macro("MainFunctionEnd") +#undef MainFunctionEnd +#define MainFunctionEnd() } + +// Helper templates + +namespace wgsl_detail { +template >> +std::string pass_as_string(T&& v) { + return std::to_string(std::forward(v)); +} +template +std::string_view pass_as_string(std::string_view sv) { + return sv; +} +template +std::string pass_as_string(T&& v) { + return std::forward(v); +} +} // namespace wgsl_detail + + +// Include template implementations + +#include "generated/tensor/pad.h" // 6903e8c7560b2507fffd2da3b327d9c6a354da219640c3131c09a92781356c9f + +#pragma pop_macro("MainFunctionStart") +#pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/build-example-pad/src/tensor/pad.wgsl.template b/tools/python/wgsl_template/test/testcases/build-example-pad/src/tensor/pad.wgsl.template new file mode 100644 index 0000000000000..6a7480c2197ee --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-example-pad/src/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>(uniforms.constant_value)[0]; +#else + bitcast(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 diff --git a/tools/python/wgsl_template/test/testcases/build-example-pad/test-config.json b/tools/python/wgsl_template/test/testcases/build-example-pad/test-config.json new file mode 100644 index 0000000000000..fb8df20f95143 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/build-example-pad/test-config.json @@ -0,0 +1,7 @@ +{ + "type": "build", + "description": "Build test for complex padding operation template with conditional parameters", + "generators": { + "static-cpp-literal": {} + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-basic/test-config.json b/tools/python/wgsl_template/test/testcases/generator-basic/test-config.json new file mode 100644 index 0000000000000..3c776b036abcd --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-basic/test-config.json @@ -0,0 +1,11 @@ +{ + "type": "generator", + "description": "Basic generator test, using the test generator", + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": {} + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template new file mode 100644 index 0000000000000..14e6f4a4dc2c3 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template @@ -0,0 +1,5 @@ +@compute @workgroup_size(128) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let buffer_count = 4; + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..a627802a64dce --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-basic/test.wgsl.template.static-cpp-literal.gen @@ -0,0 +1 @@ +ss << "@compute @workgroup_size(128)\nfn main(@builtin(global_invocation_id) global_id: vec3) {\n let buffer_count = 4;\n let index = global_id.x;\n}\n"; diff --git a/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template new file mode 100644 index 0000000000000..25ed3eaf2740d --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template @@ -0,0 +1,74 @@ +#define PAD_MODE_CONSTANT 0 +#define PAD_MODE_EDGE 1 +#define PAD_MODE_REFLECT 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>(uniforms.constant_value)[0]; +#else + bitcast(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 \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..4a7da6667d088 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-example-pad/pad.wgsl.template.static-cpp-literal.gen @@ -0,0 +1,46 @@ +MainFunctionStart(); +ss << "\n "; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size"); +ss << ";\n\n let constant_value =\n"; +if (__param_is_float16) { +ss << " bitcast>(uniforms.constant_value)[0];\n"; +} else { +ss << " bitcast(uniforms.constant_value);\n"; +} +if (__param_dim_value_zero) { +ss << " output[global_idx] = constant_value;\n"; +} else { +ss << " let output_indices = "; +ss << __var_output.OffsetToIndices("global_idx"); +ss << ";\n var input_index = u32(0);\n var use_pad_value = false;\n var in_coord = i32(0);\n\n for (var dim = 0; dim < "; +ss << __var_output.Rank(); +ss << " && !use_pad_value; dim++) {\n let output_index = i32("; +ss << GetElementAt("output_indices", "dim", __var_output.Rank()); +ss << ");\n let lower_pads = "; +ss << GetElementAt("uniforms.lower_pads", "dim", __var_output.Rank()); +ss << ";\n let data_shape = i32("; +ss << GetElementAt("uniforms.data_shape", "dim", __var_output.Rank()); +ss << ");\n"; +if (__param_pad_mode == 0) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n use_pad_value = true;\n"; +} else if (__param_pad_mode == 1) { +ss << " if (output_index < lower_pads) {\n in_coord = 0;\n } else if (output_index >= data_shape + lower_pads) {\n in_coord = data_shape - 1;\n"; +} else if (__param_pad_mode == 2) { +ss << " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n in_coord = output_index - lower_pads;\n if (in_coord < 0) {\n in_coord = -in_coord;\n }\n let _2n_1 = 2 * (data_shape - 1);\n in_coord = in_coord % _2n_1;\n if (in_coord >= data_shape) {\n in_coord = _2n_1 - in_coord;\n }\n"; +} else { +ss << " if (output_index < lower_pads) {\n in_coord = data_shape + output_index - lower_pads;\n } else if (output_index >= data_shape + lower_pads) {\n in_coord = output_index - data_shape - lower_pads;\n"; +} +ss << " } else {\n in_coord = output_index - lower_pads;\n }\n\n input_index += select(u32(in_coord)\n"; +if (__var_output.Rank() > 1) { +ss << " * "; +ss << GetElementAt("uniforms.data_stride", "dim", __var_output.Rank() - 1); +ss << "\n"; +} +ss << " , u32(in_coord), dim == "; +ss << __var_output.Rank(); +ss << " - 1);\n }\n\n "; +ss << __var_output.SetByOffset("global_idx", "select(data[input_index], constant_value, use_pad_value)"); +ss << ";\n"; +} +MainFunctionEnd(); +ss << "\n"; \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/generator-example-pad/test-config.json b/tools/python/wgsl_template/test/testcases/generator-example-pad/test-config.json new file mode 100644 index 0000000000000..b62dacab15c44 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-example-pad/test-config.json @@ -0,0 +1,20 @@ +{ + "type": "generator", + "description": "Test complex padding operation template with multiple conditional parameters and pad modes", + "entries": { + "pad.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectedParams": [ + "dim_value_zero", + "is_float16", + "pad_mode" + ], + "expectedVariables": [ + "output" + ] + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test-config.json b/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test-config.json new file mode 100644 index 0000000000000..bb4cb5d3d1b4c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test-config.json @@ -0,0 +1,13 @@ +{ + "type": "generator", + "description": "Test error handling for multiple $MAIN blocks", + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectsError": "Multiple main function start ($MAIN) detected" + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test.wgsl.template new file mode 100644 index 0000000000000..251b1cb7c9a04 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-multiple/test.wgsl.template @@ -0,0 +1,10 @@ +$MAIN { + let index = global_id.x; + output[index] = input[index] * 2.0; +} + +$MAIN { + // Second MAIN block - should cause an error + let other_index = global_id.y; + result[other_index] = data[other_index]; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test-config.json b/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test-config.json new file mode 100644 index 0000000000000..81b5ceed9b37e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test-config.json @@ -0,0 +1,13 @@ +{ + "type": "generator", + "description": "Test error handling for nested $MAIN blocks", + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectsError": "Multiple main function start ($MAIN) detected" + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test.wgsl.template new file mode 100644 index 0000000000000..efd9b19723f79 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-nested/test.wgsl.template @@ -0,0 +1,11 @@ +$MAIN { + let outer_index = global_id.x; + + $MAIN { + // Nested MAIN block - should cause an error + let inner_index = global_id.y; + result[inner_index] = data[inner_index]; + } + + output[outer_index] = input[outer_index]; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test-config.json b/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test-config.json new file mode 100644 index 0000000000000..a85babdbc5b47 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test-config.json @@ -0,0 +1,13 @@ +{ + "type": "generator", + "description": "Test error handling for $MAIN with invalid syntax (missing opening brace)", + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectsError": "Unmatched closing bracket" + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test.wgsl.template new file mode 100644 index 0000000000000..abff2602d1654 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-syntax/test.wgsl.template @@ -0,0 +1,5 @@ +$MAIN + // Missing opening brace - should cause an error + let index = global_id.x; + output[index] = input[index] * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test-config.json b/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test-config.json new file mode 100644 index 0000000000000..04d4698afab5b --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test-config.json @@ -0,0 +1,13 @@ +{ + "type": "generator", + "description": "Test error handling for unclosed $MAIN block", + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectsError": "Main function context started but not ended at the end of processing" + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test.wgsl.template new file mode 100644 index 0000000000000..66851df7bdd57 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main-invalid-unclosed/test.wgsl.template @@ -0,0 +1,4 @@ +$MAIN { + let index = global_id.x; + output[index] = input[index] * 2.0; + // Missing closing brace - should cause an error diff --git a/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template new file mode 100644 index 0000000000000..e509d3fa6fc26 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template @@ -0,0 +1,5 @@ +$MAIN { + let index = global_id.x; + // Main compute logic goes here + output[index] = input[index] * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..a34864f67d9f2 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main/main.wgsl.template.static-cpp-literal.gen @@ -0,0 +1,4 @@ +MainFunctionStart(); +ss << "\n let index = global_id.x;\n\n output[index] = input[index] * 2.0;\n"; +MainFunctionEnd(); +ss << "\n"; \ No newline at end of file diff --git a/tools/python/wgsl_template/test/testcases/generator-main/test-config.json b/tools/python/wgsl_template/test/testcases/generator-main/test-config.json new file mode 100644 index 0000000000000..68671878df777 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-main/test-config.json @@ -0,0 +1,13 @@ +{ + "type": "generator", + "description": "Test MAIN variable substitution in generator", + "entries": { + "main.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectedVariables": [] + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template new file mode 100644 index 0000000000000..8b5785d2aa72e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template @@ -0,0 +1 @@ +let size = BUFFER_SIZE; diff --git a/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..02190ecfd5099 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-no-param.wgsl.template.static-cpp-literal.gen @@ -0,0 +1 @@ +ss << "let size = BUFFER_SIZE;\n"; diff --git a/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template new file mode 100644 index 0000000000000..c9675cc48201f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template @@ -0,0 +1,2 @@ +#param BUFFER_SIZE +let size = BUFFER_SIZE; diff --git a/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..67280a99ca166 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-param-basic/compute-param.wgsl.template.static-cpp-literal.gen @@ -0,0 +1,3 @@ +ss << "let size = "; +ss << __param_BUFFER_SIZE; +ss << ";\n"; diff --git a/tools/python/wgsl_template/test/testcases/generator-param-basic/test-config.json b/tools/python/wgsl_template/test/testcases/generator-param-basic/test-config.json new file mode 100644 index 0000000000000..a45c51e9728c2 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-param-basic/test-config.json @@ -0,0 +1,22 @@ +{ + "type": "generator", + "description": "Test basic parameter substitution in template generation", + "entries": { + "compute-param.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectedParams": [ + "BUFFER_SIZE" + ] + } + } + }, + "compute-no-param.wgsl.template": { + "generators": { + "static-cpp-literal": { + "expectedParams": [] + } + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test-config.json b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test-config.json new file mode 100644 index 0000000000000..7ca4e6b0f54d2 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test-config.json @@ -0,0 +1,12 @@ +{ + "type": "generator", + "description": "Basic generator test with preserveCodeReference enabled", + "preserveCodeReference": true, + "entries": { + "test.wgsl.template": { + "generators": { + "static-cpp-literal": {} + } + } + } +} diff --git a/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template new file mode 100644 index 0000000000000..7dc773164d5d1 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template @@ -0,0 +1,11 @@ +// This is a one line comment + +@compute @workgroup_size(128) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let buffer_count = 4; /* This + is + a + multi-line + comment */ + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template.static-cpp-literal.gen b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template.static-cpp-literal.gen new file mode 100644 index 0000000000000..81546aef7a8b7 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/generator-preserve-code-ref/test.wgsl.template.static-cpp-literal.gen @@ -0,0 +1,18 @@ +// 1 | // This is a one line comment +// 2 | +// 3 | @compute @workgroup_size(128) +ss << "@compute @workgroup_size(128)\n"; +// 4 | fn main(@builtin(global_invocation_id) global_id: vec3) { +ss << "fn main(@builtin(global_invocation_id) global_id: vec3) {\n"; +// 5 | let buffer_count = 4; /* This +ss << " let buffer_count = 4;\n"; +// 6 | is +ss << "\n"; +// 7 | a +// 8 | multi-line +// 9 | comment */ +// 10 | let index = global_id.x; +ss << " let index = global_id.x;\n"; +// 11 | } +ss << "}\n"; +// 12 | diff --git a/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template new file mode 100644 index 0000000000000..be0737d0e8c98 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template @@ -0,0 +1,7 @@ +@group(0) @binding(0) var output: array; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + output[index] = f32(index) * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template.pass1 new file mode 100644 index 0000000000000..be0737d0e8c98 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-basic/basic.wgsl.template.pass1 @@ -0,0 +1,7 @@ +@group(0) @binding(0) var output: array; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + output[index] = f32(index) * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-basic/test-config.json b/tools/python/wgsl_template/test/testcases/parser-basic/test-config.json new file mode 100644 index 0000000000000..d02d7281ee4a2 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-basic/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test basic parser functionality on a simple template with no comments or preprocessor directives" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template new file mode 100644 index 0000000000000..82080411b741c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template @@ -0,0 +1,14 @@ +// Single line comment at the beginning +@group(0) @binding(0) var output: array; // End of line comment + +/* Multi-line comment + spanning multiple lines + should be removed */ +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + // Another single line comment + let index = global_id.x; + output[index] = f32(index) * 2.0 /* inline comment */ + 1.0; + /* Another + multi-line comment */ +} diff --git a/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template.pass1 new file mode 100644 index 0000000000000..bd908d4562284 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-comment-removal/comments.wgsl.template.pass1 @@ -0,0 +1,14 @@ + +@group(0) @binding(0) var output: array; + + + + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + + let index = global_id.x; + output[index] = f32(index) * 2.0 + 1.0; + + +} diff --git a/tools/python/wgsl_template/test/testcases/parser-comment-removal/test-config.json b/tools/python/wgsl_template/test/testcases/parser-comment-removal/test-config.json new file mode 100644 index 0000000000000..870639b40e0e4 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-comment-removal/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test that the parser correctly removes single-line and multi-line comments while preserving line structure" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-empty/test-config.json b/tools/python/wgsl_template/test/testcases/parser-empty/test-config.json new file mode 100644 index 0000000000000..b19f2543838d6 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-empty/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test parser behavior on an empty directory with no template files" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template new file mode 100644 index 0000000000000..f1503f4c37959 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template @@ -0,0 +1,7 @@ +#include "utils.wgsl.template" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let result = multiply_by_two(f32(index)); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..8b2f8b53598b7 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-basic/main.wgsl.template.pass1 @@ -0,0 +1,10 @@ +fn multiply_by_two(value: f32) -> f32 { + return value * 2.0; +} + + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let result = multiply_by_two(f32(index)); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-basic/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-basic/test-config.json new file mode 100644 index 0000000000000..71f1284a7d50d --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-basic/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test basic #include preprocessor directive functionality with a simple include scenario" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template new file mode 100644 index 0000000000000..d5bb7dc2e62ef --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template @@ -0,0 +1,3 @@ +fn multiply_by_two(value: f32) -> f32 { + return value * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template.pass1 new file mode 100644 index 0000000000000..d5bb7dc2e62ef --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-basic/utils.wgsl.template.pass1 @@ -0,0 +1,3 @@ +fn multiply_by_two(value: f32) -> f32 { + return value * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template new file mode 100644 index 0000000000000..38d40c7a3be6d --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template @@ -0,0 +1,32 @@ +// This file contains #include directives inside comments +// These should be ignored by the parser + +// Single-line comment with include: #include "nonexistent.wgsl.template" + +/* + * Multi-line comment with include + * #include "another-nonexistent.wgsl.template" + * This should also be ignored + */ + +/* +This is a multi-line comment where #include appears at line start: +#include "linestart-nonexistent.wgsl.template" +And another one: +#include "second-linestart.wgsl.template" +End of multi-line comment +*/ + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + // Another commented include: #include "utils.wgsl.template" + + /* + * Block comment with multiple includes: + * #include "math.wgsl.template" + * #include "graphics.wgsl.template" + */ + + let result = f32(index) * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..a0a830cdbaf0c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/main.wgsl.template.pass1 @@ -0,0 +1,32 @@ + + + + + + + + + + + + + + + + + + + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + + + + + + + + + let result = f32(index) * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/test-config.json new file mode 100644 index 0000000000000..7bb03a205bc41 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test that #include directives inside comments are ignored and not processed" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template new file mode 100644 index 0000000000000..a32991ef55789 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template @@ -0,0 +1,6 @@ +// This file should NOT be included because all #include directives +// referencing it are inside comments + +fn utility_function() -> f32 { + return 42.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template.pass1 new file mode 100644 index 0000000000000..d3c929ae493d5 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-inside-comment/utils.wgsl.template.pass1 @@ -0,0 +1,6 @@ + + + +fn utility_function() -> f32 { + return 42.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/main.wgsl.template new file mode 100644 index 0000000000000..dbd7b8c25bd58 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/main.wgsl.template @@ -0,0 +1,7 @@ +// Test empty include directive +#include "" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/test-config.json new file mode 100644 index 0000000000000..fa1fed525b29f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-empty/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser behavior with empty include directive", + "expectsError": "File \"\" not found" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/main.wgsl.template new file mode 100644 index 0000000000000..940bb184eef6f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/main.wgsl.template @@ -0,0 +1,7 @@ +// Test include directive without quotes +#include utils.wgsl.template + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/test-config.json new file mode 100644 index 0000000000000..9eec2e9ed6ef7 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-no-quotes/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser behavior with include directive missing quotes", + "expectsError": "file path must be enclosed in double quotes" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/main.wgsl.template new file mode 100644 index 0000000000000..d06102a18165b --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/main.wgsl.template @@ -0,0 +1,12 @@ +// Test file that tries to include non-existent files +#include "this-file-does-not-exist.wgsl.template" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + + // This function should exist if includes worked + // let processed = process_data(index); + + let result = f32(index) * 2.0; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/test-config.json new file mode 100644 index 0000000000000..f4bc36dab486a --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-nonexist/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser behavior when trying to include files that do not exist", + "expectsError": "File \"this-file-does-not-exist.wgsl.template\" not found" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/main.wgsl.template new file mode 100644 index 0000000000000..51b55984e4e03 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/main.wgsl.template @@ -0,0 +1,7 @@ +// Test include directive with invalid path characters +#include "../../../invalid<>path|file*.wgsl.template" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/test-config.json new file mode 100644 index 0000000000000..366a3093d2518 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-path/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser behavior with include directive containing invalid path characters", + "expectsError": "File \"../../../invalid<>path|file*.wgsl.template\" not found" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/main.wgsl.template new file mode 100644 index 0000000000000..56827df5b2144 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/main.wgsl.template @@ -0,0 +1,7 @@ +// Test include directive with unclosed quotes +#include "utils.wgsl.template + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/test-config.json new file mode 100644 index 0000000000000..f843b88dbe19a --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-invalid-unclosed-quotes/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser behavior with include directive having unclosed quotes", + "expectsError": "file path must be enclosed in double quotes" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template new file mode 100644 index 0000000000000..a4e4deab2291f --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template @@ -0,0 +1,3 @@ +// Basic constants used throughout the shader pipeline +const MAX_VALUE: f32 = 100.0; +const SCALE_FACTOR: f32 = 1.5; diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template.pass1 new file mode 100644 index 0000000000000..7aa45bf08f0dd --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/constants.wgsl.template.pass1 @@ -0,0 +1,3 @@ + +const MAX_VALUE: f32 = 100.0; +const SCALE_FACTOR: f32 = 1.5; diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template new file mode 100644 index 0000000000000..09f5cb0d6ac6e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template @@ -0,0 +1,11 @@ +// Graphics utilities that depend on math utilities +#include "math.wgsl.template" + +fn create_color(value: f32) -> vec4 { + let normalized = normalize_value(value); + return vec4(normalized, normalized * 0.5, normalized * 0.8, 1.0); +} + +fn transform_vertex(pos: vec3) -> vec3 { + return multiply_vec3(pos, 2.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template.pass1 new file mode 100644 index 0000000000000..479a544cfa92c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/graphics.wgsl.template.pass1 @@ -0,0 +1,24 @@ + + + +const MAX_VALUE: f32 = 100.0; +const SCALE_FACTOR: f32 = 1.5; + + +fn normalize_value(value: f32) -> f32 { + return clamp(value / MAX_VALUE, 0.0, 1.0); +} + +fn multiply_vec3(v: vec3, factor: f32) -> vec3 { + return v * factor * SCALE_FACTOR; +} + + +fn create_color(value: f32) -> vec4 { + let normalized = normalize_value(value); + return vec4(normalized, normalized * 0.5, normalized * 0.8, 1.0); +} + +fn transform_vertex(pos: vec3) -> vec3 { + return multiply_vec3(pos, 2.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template new file mode 100644 index 0000000000000..9bbe79466d947 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template @@ -0,0 +1,9 @@ +// Main shader file that includes higher-level utilities +#include "graphics.wgsl.template" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let color = create_color(f32(index)); + let transformed = transform_vertex(vec3(color.rgb)); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..d8c84e8ed4fd0 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/main.wgsl.template.pass1 @@ -0,0 +1,33 @@ + + + + +const MAX_VALUE: f32 = 100.0; +const SCALE_FACTOR: f32 = 1.5; + + +fn normalize_value(value: f32) -> f32 { + return clamp(value / MAX_VALUE, 0.0, 1.0); +} + +fn multiply_vec3(v: vec3, factor: f32) -> vec3 { + return v * factor * SCALE_FACTOR; +} + + +fn create_color(value: f32) -> vec4 { + let normalized = normalize_value(value); + return vec4(normalized, normalized * 0.5, normalized * 0.8, 1.0); +} + +fn transform_vertex(pos: vec3) -> vec3 { + return multiply_vec3(pos, 2.0); +} + + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let color = create_color(f32(index)); + let transformed = transform_vertex(vec3(color.rgb)); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template new file mode 100644 index 0000000000000..cd924a8600af6 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template @@ -0,0 +1,10 @@ +// Math utilities that depend on basic constants +#include "constants.wgsl.template" + +fn normalize_value(value: f32) -> f32 { + return clamp(value / MAX_VALUE, 0.0, 1.0); +} + +fn multiply_vec3(v: vec3, factor: f32) -> vec3 { + return v * factor * SCALE_FACTOR; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template.pass1 new file mode 100644 index 0000000000000..1c163aafd3dbb --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/math.wgsl.template.pass1 @@ -0,0 +1,13 @@ + + +const MAX_VALUE: f32 = 100.0; +const SCALE_FACTOR: f32 = 1.5; + + +fn normalize_value(value: f32) -> f32 { + return clamp(value / MAX_VALUE, 0.0, 1.0); +} + +fn multiply_vec3(v: vec3, factor: f32) -> vec3 { + return v * factor * SCALE_FACTOR; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/test-config.json new file mode 100644 index 0000000000000..689cd17950e5c --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-multi-layer/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test multi-layer include preprocessing with nested includes (A includes B, B includes C, C includes D)" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template new file mode 100644 index 0000000000000..291863756bafa --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template @@ -0,0 +1 @@ +const PI: f32 = 3.14159; diff --git a/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template.pass1 new file mode 100644 index 0000000000000..291863756bafa --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/constants.wgsl.template.pass1 @@ -0,0 +1 @@ +const PI: f32 = 3.14159; diff --git a/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template new file mode 100644 index 0000000000000..351a3540bc3ff --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template @@ -0,0 +1,8 @@ +// Test #include with space at beginning of line + #include "constants.wgsl.template" + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let value = PI * f32(index); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..b9fd33670db55 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/main.wgsl.template.pass1 @@ -0,0 +1,9 @@ + +const PI: f32 = 3.14159; + + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let value = PI * f32(index); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/test-config.json b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/test-config.json new file mode 100644 index 0000000000000..b59d5809ac053 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-include-space-prefix/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test #include directive with space at the beginning of line" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template new file mode 100644 index 0000000000000..51f4a019abef0 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template @@ -0,0 +1,15 @@ +#define PI 3.14159 +#define MAX_SIZE 1024 +#define SHADER_VERSION 450 + +const pi_value = PI; +const buffer_size = MAX_SIZE; +var version: u32 = SHADER_VERSION; + +fn calculate_area(radius: f32) -> f32 { + return PI * radius * radius; +} + +fn check_bounds(size: u32) -> bool { + return size <= MAX_SIZE; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template.pass1 new file mode 100644 index 0000000000000..fa05b3b9bfe46 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-basic/basic.wgsl.template.pass1 @@ -0,0 +1,15 @@ + + + + +const pi_value = 3.14159; +const buffer_size = 1024; +var version: u32 = 450; + +fn calculate_area(radius: f32) -> f32 { + return 3.14159 * radius * radius; +} + +fn check_bounds(size: u32) -> bool { + return size <= 1024; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-basic/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-basic/test-config.json new file mode 100644 index 0000000000000..3c4485774d888 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-basic/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test basic #define macro directive functionality with simple macro definitions and substitutions" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template new file mode 100644 index 0000000000000..58ead1fbb83cf --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template @@ -0,0 +1,14 @@ +#define A 10 +#define B (A * 2) +#define C (B + 5) +#define PI 3.14159 +#define DOUBLE_PI (PI * 2) + +const value_a = A; +const value_b = B; +const value_c = C; +const circle_circumference = DOUBLE_PI * radius; + +fn calculate() -> f32 { + return A + B + C; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template.pass1 new file mode 100644 index 0000000000000..ea7418ffd2403 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-chain/chain.wgsl.template.pass1 @@ -0,0 +1,14 @@ + + + + + + +const value_a = 10; +const value_b = (10 * 2); +const value_c = ((10 * 2) + 5); +const circle_circumference = (3.14159 * 2) * radius; + +fn calculate() -> f32 { + return 10 + (10 * 2) + ((10 * 2) + 5); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-chain/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-chain/test-config.json new file mode 100644 index 0000000000000..46ecf62a75d97 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-chain/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test macro chaining where macros can reference other previously defined macros" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template new file mode 100644 index 0000000000000..d13139dbb6352 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template @@ -0,0 +1,11 @@ +#define COMPLEX_EXPR (a + b) * (c - d) / 2.0 +#define STRING_LIKE "hello world" +#define WITH_SYMBOLS @#$%^&*(){}[] +#define MATH_EXPR sin(x * PI / 180.0) +#define ARRAY_ACCESS data[index * 2 + 1] + +const result1 = COMPLEX_EXPR; +const result2 = STRING_LIKE; +const result3 = WITH_SYMBOLS; +const result4 = MATH_EXPR; +const result5 = ARRAY_ACCESS; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template.pass1 new file mode 100644 index 0000000000000..5c7664a8d8c82 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/complex.wgsl.template.pass1 @@ -0,0 +1,11 @@ + + + + + + +const result1 = (a + b) * (c - d) / 2.0; +const result2 = "hello world"; +const result3 = @#$%^&*(){}[]; +const result4 = sin(x * PI / 180.0); +const result5 = data[index * 2 + 1]; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/test-config.json new file mode 100644 index 0000000000000..a92934a31c2d5 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-complex-values/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test macro definitions with complex values including expressions, strings, and special characters" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/circular.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/circular.wgsl.template new file mode 100644 index 0000000000000..921113ad512cb --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/circular.wgsl.template @@ -0,0 +1,8 @@ +#define SIZE 1024 +#define HALF (SIZE / 2) +#define A B +#define B A + +const buffer_size = SIZE; +const half_size = HALF; +const value = A; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/test-config.json new file mode 100644 index 0000000000000..95571e996af47 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-circular/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser error handling for circular macro references", + "expectsError": "Circular macro reference in file circular.wgsl.template at line 4: macro \"B\" creates circular dependency with \"A\"" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/duplicate.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/duplicate.wgsl.template new file mode 100644 index 0000000000000..ab3057eb6b0e6 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/duplicate.wgsl.template @@ -0,0 +1,5 @@ +#define PI 3.14159 +#define MAX_SIZE 1024 +#define PI 3.141592653589793 + +const area = PI * radius * radius; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/test-config.json new file mode 100644 index 0000000000000..9e802dd01fbb9 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-duplicate/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser error handling for duplicate macro definitions", + "expectsError": "Duplicate macro definition in file duplicate.wgsl.template at line 3: \"PI\" is already defined" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/empty.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/empty.wgsl.template new file mode 100644 index 0000000000000..d24b51c3e82fd --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/empty.wgsl.template @@ -0,0 +1,3 @@ +#define EMPTY_MACRO + +const test = EMPTY_MACRO; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/test-config.json new file mode 100644 index 0000000000000..c3dcf6d27190d --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-empty-value/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser error handling for macro definition with empty value", + "expectsError": "Invalid macro definition in file empty.wgsl.template at line 1: macro \"EMPTY_MACRO\" has no value" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/invalid.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/invalid.wgsl.template new file mode 100644 index 0000000000000..3d697cd5982c8 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/invalid.wgsl.template @@ -0,0 +1 @@ +#define 123INVALID value1 diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/test-config.json new file mode 100644 index 0000000000000..097e4d639360a --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-syntax/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser error handling for invalid macro name syntax (starts with number)", + "expectsError": "Invalid macro definition in file invalid.wgsl.template at line 1: invalid macro name \"123INVALID\" (must start with letter or underscore, contain only letters, numbers, and underscores)" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/test-config.json new file mode 100644 index 0000000000000..8da33cfc61ccd --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/test-config.json @@ -0,0 +1,5 @@ +{ + "type": "parser", + "description": "Test parser error handling for macro definition with whitespace-only value", + "expectsError": "Invalid macro definition in file whitespace.wgsl.template at line 1: macro \"SPACES_ONLY\" has no value" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/whitespace.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/whitespace.wgsl.template new file mode 100644 index 0000000000000..4dabfef306692 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-invalid-whitespace-value/whitespace.wgsl.template @@ -0,0 +1,9 @@ +#define SPACES_ONLY +#define TABS_ONLY +#define MIXED_SPACES +#define VALID_MACRO actual_value + +const test1 = SPACES_ONLY; +const test2 = TABS_ONLY; +const test3 = MIXED_SPACES; +const test4 = VALID_MACRO; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template new file mode 100644 index 0000000000000..ce3eedc022e58 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template @@ -0,0 +1,13 @@ +#define _ 1 +#define _underscore_start 2 +#define A 3 +#define A123 4 +#define ABC_123_DEF 5 +#define __double_underscore__ 6 + +const test1 = _; +const test2 = _underscore_start; +const test3 = A; +const test4 = A123; +const test5 = ABC_123_DEF; +const test6 = __double_underscore__; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template.pass1 new file mode 100644 index 0000000000000..8fa47b8ea7d58 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/names.wgsl.template.pass1 @@ -0,0 +1,13 @@ + + + + + + + +const test1 = 1; +const test2 = 2; +const test3 = 3; +const test4 = 4; +const test5 = 5; +const test6 = 6; diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/test-config.json new file mode 100644 index 0000000000000..47b33d389e587 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-name-edge-cases/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test valid macro names including edge cases with underscores and single characters" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template new file mode 100644 index 0000000000000..b81839bd8a576 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template @@ -0,0 +1,15 @@ +// Test #define with space at beginning of line + #define PI 3.14159 + #define MAX_SIZE 1024 + +const pi_value = PI; +const buffer_size = MAX_SIZE; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let area = PI * f32(index) * f32(index); + if (index < MAX_SIZE) { + // Process within bounds + } +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..9ad16c66e1c9b --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/main.wgsl.template.pass1 @@ -0,0 +1,15 @@ + + + + +const pi_value = 3.14159; +const buffer_size = 1024; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let index = global_id.x; + let area = 3.14159 * f32(index) * f32(index); + if (index < 1024) { + + } +} diff --git a/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/test-config.json b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/test-config.json new file mode 100644 index 0000000000000..19a6947b9497b --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-macro-space-prefix/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Test #define directive with space at the beginning of line" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template new file mode 100644 index 0000000000000..bad2ed5844332 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template @@ -0,0 +1,8 @@ +// Main shader template +#include "shaders/vertex.wgsl.template" +#include "utils/common.wgsl.template" + +@fragment +fn fs_main() -> @location(0) vec4f { + return vec4f(1.0, 0.0, 0.0, 1.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template.pass1 new file mode 100644 index 0000000000000..f8a9a09ae38b4 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/main.wgsl.template.pass1 @@ -0,0 +1,27 @@ + + +@vertex +fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f { + let pos = array( + vec2f(-0.5, -0.5), + vec2f( 0.5, -0.5), + vec2f( 0.0, 0.5) + ); + return vec4f(pos[vertex_index], 0.0, 1.0); +} + + +fn get_triangle_color(index: u32) -> vec4f { + let colors = array( + vec4f(1.0, 0.0, 0.0, 1.0), + vec4f(0.0, 1.0, 0.0, 1.0), + vec4f(0.0, 0.0, 1.0, 1.0) + ); + return colors[index % 3u]; +} + + +@fragment +fn fs_main() -> @location(0) vec4f { + return vec4f(1.0, 0.0, 0.0, 1.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template new file mode 100644 index 0000000000000..3b6bff3ea4e1d --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template @@ -0,0 +1,10 @@ +// Vertex shader template +@vertex +fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f { + let pos = array( + vec2f(-0.5, -0.5), + vec2f( 0.5, -0.5), + vec2f( 0.0, 0.5) + ); + return vec4f(pos[vertex_index], 0.0, 1.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template.pass1 new file mode 100644 index 0000000000000..4c6bf4f34107e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/shaders/vertex.wgsl.template.pass1 @@ -0,0 +1,10 @@ + +@vertex +fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f { + let pos = array( + vec2f(-0.5, -0.5), + vec2f( 0.5, -0.5), + vec2f( 0.0, 0.5) + ); + return vec4f(pos[vertex_index], 0.0, 1.0); +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/test-config.json b/tools/python/wgsl_template/test/testcases/parser-subfolder/test-config.json new file mode 100644 index 0000000000000..08987971282e2 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/test-config.json @@ -0,0 +1,4 @@ +{ + "type": "parser", + "description": "Parser test with template files in subfolders" +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template b/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template new file mode 100644 index 0000000000000..c1b115f01f221 --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template @@ -0,0 +1,9 @@ +// Common utilities template +fn get_triangle_color(index: u32) -> vec4f { + let colors = array( + vec4f(1.0, 0.0, 0.0, 1.0), // Red + vec4f(0.0, 1.0, 0.0, 1.0), // Green + vec4f(0.0, 0.0, 1.0, 1.0) // Blue + ); + return colors[index % 3u]; +} diff --git a/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template.pass1 b/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template.pass1 new file mode 100644 index 0000000000000..d913f68adbd3e --- /dev/null +++ b/tools/python/wgsl_template/test/testcases/parser-subfolder/utils/common.wgsl.template.pass1 @@ -0,0 +1,9 @@ + +fn get_triangle_color(index: u32) -> vec4f { + let colors = array( + vec4f(1.0, 0.0, 0.0, 1.0), + vec4f(0.0, 1.0, 0.0, 1.0), + vec4f(0.0, 0.0, 1.0, 1.0) + ); + return colors[index % 3u]; +} diff --git a/tools/python/wgsl_template/types.py b/tools/python/wgsl_template/types.py new file mode 100644 index 0000000000000..128a6c338467d --- /dev/null +++ b/tools/python/wgsl_template/types.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Shared data structures for the WGSL template engine.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class CodeReference: + """Where a parsed line came from in the original source.""" + + file_path: str + line_number: int # 1-based + + +@dataclass +class ParsedLine: + """A line after parsing, plus its origin.""" + + line: str + code_reference: CodeReference + + +@dataclass +class TemplatePass0: + """Result of PASS0 (load): raw lines straight from disk.""" + + file_path: str # absolute resolved path on disk + raw: list[str] # split on /\r?\n/ + + +@dataclass +class TemplatePass1: + """Result of PASS1 (parse): pre-processed but not yet code-generated.""" + + file_path: str + raw: list[str] + pass1: list[ParsedLine] + + +@dataclass +class GenerateResult: + """Result of running PASS2 over a single template.""" + + code: str + params: dict[str, str] = field(default_factory=dict) # name -> param-type + variables: dict[str, str] = field(default_factory=dict) # name -> variable-type + has_main_function: bool = False + + +@dataclass +class TemplatePass2: + """Result of PASS2: code-generated string + metadata.""" + + file_path: str + generate_result: GenerateResult + + +@dataclass +class TemplateRepository: + """A collection of templates loaded from disk. + + ``templates`` may hold ``TemplatePass0``, ``TemplatePass1``, or + ``TemplatePass2`` depending on which pipeline stage produced it. + """ + + base_path: str + templates: dict[str, object] + + +# ---------------------------------------------------------------------- +# Source directory specification (loader input) +# ---------------------------------------------------------------------- + + +@dataclass +class SourceDir: + """A source directory with an optional alias prefix. + + The CLI doesn't expose aliases, but the loader API still supports + them so directory-aliased test fixtures keep working. + """ + + path: str + alias: str | None = None