Add shader key validation step in WebGPU CI pipeline (#24243)

fs-eire · web-flow · commit 4eeefd7260b7 · 2025-03-31T10:36:12.000-07:00
### Description This PR adds a shader key validation step to the WebGPU CI pipeline. The shader key validation works in this way: - first, run onnxruntime_test_all with verbose logging, dumping the logs into a file - then, parse the file and found WebGPU EP program logs. The log contains the following information: - the shader cache key - the corresponding shader code The script will aggregate those information and make sure for each cache key, the corresponding shader code must be consistent. To make the validation work, this PR also modified a few things: - set the locale of `std::wclog` to ".UTF-8" to support Unicode characters. Otherwise the logger will fail and no longer output future logs. A fix is submitted in PR #24237 but there is a concern if this may potentially break some users. Setting inside onnxruntime_test_all is pretty safe. - re-enable the WebGPU device auto collect which was introduced in #24115. Now we have a better way to detect cache key inconsistency. ### Next Step The newly added test is marked as `continue-on-error: true`, which means even if it failed it does not block the CI pipeline. We should fix those failures one-by-one and eventually the test should pass. then we can remove the `continue-on-error: true` flag.
diff --git a/.github/actions/webgpu-validate-shader-key/action.yml b/.github/actions/webgpu-validate-shader-key/action.yml
@@ -14,14 +14,15 @@ runs:
   using: "composite"
   steps:
     - name: Validate shader keys (chromium log)
-      if: ${{ inputs.is_chromium_log }}
+      # GitHub Actions treats all inputs as strings even if it's specified as a boolean.
+      if: ${{ inputs.is_chromium_log == 'true' }}
       shell: cmd
       run: |
         node parse-chromium-debug-log.js < "${{ inputs.log_file_path }}" | node validate-shader-key.js
       working-directory: ${{ github.action_path }}
 
     - name: Validate shader keys (native log)
-      if: ${{ !inputs.is_chromium_log }}
+      if: ${{ !inputs.is_chromium_log != 'true' }}
       shell: cmd
       run: |
         node validate-shader-key.js < "${{ inputs.log_file_path }}"
diff --git a/.github/actions/webgpu-validate-shader-key/validate-shader-key.js b/.github/actions/webgpu-validate-shader-key/validate-shader-key.js
@@ -11,21 +11,30 @@ const readline = require("readline");
 
 const shaderMap = new Map();
 
+const regexStartingProgram =
+  /onnxruntime::webgpu::WebGpuContext::Run.+Starting program \"(?<key>.+)\"/;
 const regexShaderStart =
-  /^===\ WebGPU\ Shader\ code\ \[.+?Key=\"(?<key>.+)\"]\ Start\ ===$/;
+  /^===\ WebGPU\ Shader\ code\ \[.+?(Key=\"(?<key>.+)\")?]\ Start\ ===$/;
 const regexShaderEnd =
-  /^===\ WebGPU\ Shader\ code\ \[.+?Key=\"(?<key>.+)\"]\ End\ ===$/;
+  /^===\ WebGPU\ Shader\ code\ \[.+?(Key=\"(?<key>.+)\")?]\ End\ ===$/;
 
 async function processVerboseLog() {
   const rl = readline.createInterface({
     input: process.stdin,
     crlfDelay: Infinity,
   });
 
+  let lastProgramKey = null;
   let currentShaderKey = null;
   let currentShaderCode = null;
 
   for await (const line of rl) {
+    const startingProgram = regexStartingProgram.exec(line);
+    if (startingProgram) {
+      lastProgramKey = startingProgram.groups.key;
+      continue;
+    }
+
     const resultStart = regexShaderStart.exec(line);
     if (resultStart) {
       if (currentShaderKey) {
@@ -34,7 +43,18 @@ async function processVerboseLog() {
         );
       }
 
-      currentShaderKey = resultStart.groups.key;
+      const key = resultStart.groups.key ?? lastProgramKey;
+      if (!key) {
+        throw new Error(
+          'No shader key is found in the log. Please use debug build or enable verbose logging in session options in release build.'
+        );
+      }
+      if (lastProgramKey && key !== lastProgramKey) {
+        throw new Error(
+          `Found incorrect shader key from log. Expected "${lastProgramKey}", but got "${key}".`
+        );
+      }
+      currentShaderKey = key;
       currentShaderCode = "";
       continue;
     }
@@ -45,9 +65,17 @@ async function processVerboseLog() {
         throw new Error(
           `Found unexpected shader end for key "${resultEnd.groups.key}".`
         );
-      } else if (currentShaderKey !== resultEnd.groups.key) {
+      }
+
+      const key = resultEnd.groups.key ?? lastProgramKey;
+      if (!key) {
+        throw new Error(
+          'No shader key is found in the log. Please use debug build or enable verbose logging in session options in release build.'
+        );
+      }
+      if (lastProgramKey && key !== lastProgramKey) {
         throw new Error(
-          `Found inconsistent shader key. Expected "${currentShaderKey}", but got "${resultEnd.groups.key}".`
+          `Found incorrect shader key from log. Expected "${lastProgramKey}", but got "${key}".`
         );
       }
 
@@ -87,7 +115,7 @@ ${currentShaderCode}
   }
 
   console.log(
-    `All shader code is consistent. Total ${shaderMap.size} shader code found.`
+    `All shader code is consistent. Total ${shaderMap.size} shader keys found.`
   );
 }
 
diff --git a/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml
@@ -24,10 +24,10 @@ jobs:
       OnnxRuntimeBuildDirectory: ${{ github.workspace }}
       DOTNET_SKIP_FIRST_TIME_EXPERIENCE: true
       setVcvars: true
-      ALLOW_RELEASED_ONNX_OPSET_ONLY: '0'
+      ALLOW_RELEASED_ONNX_OPSET_ONLY: "0"
       DocUpdateNeeded: false
-      NVIDIA_TF32_OVERRIDE: '0'
-      ONNXRUNTIME_TEST_GPU_DEVICE_ID: '0'
+      NVIDIA_TF32_OVERRIDE: "0"
+      ONNXRUNTIME_TEST_GPU_DEVICE_ID: "0"
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -38,7 +38,7 @@ jobs:
       - name: Setup Python 3.12
         uses: actions/setup-python@v5
         with:
-          python-version: '3.12'
+          python-version: "3.12"
           architecture: x64
 
       - name: Locate vcvarsall and Setup Env
@@ -54,13 +54,13 @@ jobs:
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
-          node-version: '20.x'
+          node-version: "20.x"
 
       - name: Setup Java
         uses: actions/setup-java@v4
         with:
-          distribution: 'temurin'
-          java-version: '17'
+          distribution: "temurin"
+          java-version: "17"
           architecture: x64
 
       - name: API Documentation Check and generate
@@ -78,12 +78,12 @@ jobs:
         env:
           PROCESSOR_ARCHITECTURE: x64
         with:
-          dotnet-version: '8.x'
+          dotnet-version: "8.x"
 
       - name: Use Nuget 6.x
         uses: nuget/setup-nuget@v2
         with:
-          nuget-version: '6.x'
+          nuget-version: "6.x"
 
       - name: NuGet restore
         run: |
@@ -113,13 +113,30 @@ jobs:
           }
           Remove-Item "${{ github.workspace }}\RelWithDebInfo" -Include "*.obj" -Recurse
 
+      - name: Run tests (onnxruntime_test_all) with verbose logging
+        shell: pwsh
+        run: |
+          $env:ORT_UNIT_TEST_MAIN_LOG_LEVEL = "0"
+          .\onnxruntime_test_all.exe 2>.\onnxruntime_test_all_stderr.log
+        working-directory: ${{ github.workspace }}\RelWithDebInfo\RelWithDebInfo
+
+      - name: Check log file
+        shell: cmd
+        run: |
+          dir ${{ github.workspace }}\RelWithDebInfo\RelWithDebInfo\onnxruntime_test_all_stderr.log
+
+      - name: Validate shader keys
+        continue-on-error: true
+        uses: ./.github/actions/webgpu-validate-shader-key
+        with:
+          log_file_path: ${{ github.workspace }}\RelWithDebInfo\RelWithDebInfo\onnxruntime_test_all_stderr.log
+
       - name: Validate C# native delegates
         run: python tools\ValidateNativeDelegateAttributes.py
         shell: cmd
         working-directory: ${{ github.workspace }}\csharp
         continue-on-error: true
 
-
   webgpu_external_dawn_build_x64_RelWithDebInfo:
     runs-on: ["self-hosted", "1ES.Pool=onnxruntime-github-Win2022-GPU-A10"]
     timeout-minutes: 300
@@ -133,7 +150,7 @@ jobs:
       - name: Setup Python 3.12
         uses: actions/setup-python@v5
         with:
-          python-version: '3.12'
+          python-version: "3.12"
           architecture: x64
 
       - name: Locate vcvarsall and Setup Env
@@ -177,12 +194,12 @@ jobs:
     runs-on: ["self-hosted", "1ES.Pool=onnxruntime-github-Win2022-GPU-A10"]
     timeout-minutes: 300
     env:
-        OrtPackageId: Microsoft.ML.OnnxRuntime
-        OnnxRuntimeBuildDirectory: ${{ github.workspace }}
-        DOTNET_SKIP_FIRST_TIME_EXPERIENCE: true
-        ALLOW_RELEASED_ONNX_OPSET_ONLY: '0'
-        DocUpdateNeeded: false
-        ONNXRUNTIME_TEST_GPU_DEVICE_ID: '0'
+      OrtPackageId: Microsoft.ML.OnnxRuntime
+      OnnxRuntimeBuildDirectory: ${{ github.workspace }}
+      DOTNET_SKIP_FIRST_TIME_EXPERIENCE: true
+      ALLOW_RELEASED_ONNX_OPSET_ONLY: "0"
+      DocUpdateNeeded: false
+      ONNXRUNTIME_TEST_GPU_DEVICE_ID: "0"
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -193,7 +210,7 @@ jobs:
       - name: Setup Python 3.12
         uses: actions/setup-python@v5
         with:
-          python-version: '3.12'
+          python-version: "3.12"
           architecture: x64
 
       - name: Locate vcvarsall and Setup Env
@@ -209,13 +226,13 @@ jobs:
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
-          node-version: '20.x'
+          node-version: "20.x"
 
       - name: Setup Java
         uses: actions/setup-java@v4
         with:
-          distribution: 'temurin'
-          java-version: '17'
+          distribution: "temurin"
+          java-version: "17"
           architecture: x64
 
       - name: API Documentation Check and generate
@@ -233,12 +250,12 @@ jobs:
         env:
           PROCESSOR_ARCHITECTURE: x64
         with:
-          dotnet-version: '8.x'
+          dotnet-version: "8.x"
 
       - name: Use Nuget 6.x
         uses: nuget/setup-nuget@v2
         with:
-          nuget-version: '6.x'
+          nuget-version: "6.x"
 
       - name: NuGet restore
         run: |
diff --git a/onnxruntime/test/unittest_main/test_main.cc b/onnxruntime/test/unittest_main/test_main.cc
@@ -5,6 +5,10 @@
 #include <cstdlib>
 #include <optional>
 #include <string>
+#ifdef _WIN32
+#include <iostream>
+#include <locale>
+#endif
 
 #ifndef USE_ONNXRUNTIME_DLL
 #ifdef __GNUC__
@@ -29,6 +33,11 @@ std::unique_ptr<Ort::Env> ort_env;
 
 // ortenv_setup() and ortenv_teardown() are used by onnxruntime/test/xctest/xcgtest.mm so can't be file local
 extern "C" void ortenv_setup() {
+#ifdef _WIN32
+  // Set the locale to UTF-8 to ensure proper handling of wide characters on Windows
+  std::wclog.imbue(std::locale(".UTF-8", std::locale::ctype));
+#endif
+
   OrtThreadingOptions tpo;
 
   // allow verbose logging to be enabled by setting this environment variable to a numeric log level
diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc
@@ -303,10 +303,6 @@ std::unique_ptr<IExecutionProvider> DefaultWebGpuExecutionProvider() {
   ORT_ENFORCE(config_options.AddConfigEntry(webgpu::options::kStorageBufferCacheMode,
                                             webgpu::options::kBufferCacheMode_Disabled)
                   .IsOK());
-  // Disable device auto collect
-  ORT_ENFORCE(config_options.AddConfigEntry(webgpu::options::kPreserveDevice,
-                                            webgpu::options::kPreserveDevice_ON)
-                  .IsOK());
   return WebGpuProviderFactoryCreator::Create(config_options)->CreateProvider();
 #else
   return nullptr;
