[ExecuTorch][WebGPU] Add plain Gemma 4 source-closure tests

backends/webgpu/CMakeLists.txt

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -261,6 +261,10 @@
       webgpu_update_cache_test test/native/test_update_cache.cpp
     )
     target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
+    target_include_directories(
+      webgpu_update_cache_test
+      PRIVATE "${EXECUTORCH_ROOT}/third-party/json/single_include"
+    )
     add_webgpu_native_test(
       webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
     )
@@ -275,6 +279,19 @@
     target_link_libraries(
       webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main
     )
+    add_webgpu_native_test(
+      webgpu_update_cache_state_test
+      test/native/test_update_cache_state.cpp
+    )
+    target_link_libraries(
+      webgpu_update_cache_state_test PRIVATE GTest::gtest GTest::gtest_main
+    )
+    add_webgpu_native_test(
+      webgpu_slice_chain_test test/native/test_slice_chain.cpp
+    )
+    target_link_libraries(
+      webgpu_slice_chain_test PRIVATE GTest::gtest GTest::gtest_main
+    )
     add_executable(
       webgpu_execution_options_test test/native/test_execution_options.cpp
                                     runtime/WebGPUExecutionOptions.cpp

backends/webgpu/scripts/test_webgpu_native_ci.sh

@@ -61,6 +61,35 @@ run_with_required_device() {
   fi
 }
 
+run_required_gtests() {
+  local output
+  if ! output="$(run_with_required_device "$@" 2>&1)"; then
+    printf '%s\n' "${output}"
+    return 1
+  fi
+  printf '%s\n' "${output}"
+
+  local test_name
+  for test_name in \
+    DynamicShape.SliceCrosses2dDispatchBoundary \
+    DynamicShape.CatCrosses2dDispatchBoundary \
+    DynamicShape.SliceDualStoreWritesBothDestinations; do
+    if ! grep -Eq "^\\[       OK \\] ${test_name}( \\([0-9]+ ms\\))?$" \
+        <<<"${output}"; then
+      echo "ERROR: required WebGPU test did not pass: ${test_name}" >&2
+      return 1
+    fi
+  done
+  if ! grep -Fxq '[  PASSED  ] 3 tests.' <<<"${output}"; then
+    echo "ERROR: required WebGPU run did not pass exactly three tests" >&2
+    return 1
+  fi
+  if grep -Eq '^\\[  SKIPPED \\]' <<<"${output}"; then
+    echo "ERROR: required WebGPU run skipped a test" >&2
+    return 1
+  fi
+}
+
 DISPATCH_ORDER_DIR="/tmp/dispatch_order"
 UPDATE_CACHE_DIR="/tmp/update_cache"
 INDEX_DIR="/tmp/index"
@@ -138,10 +167,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
 export_all_index_models('${INDEX_DIR}')
 "
 
-$PYTHON_EXECUTABLE -c "
+WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.dynamic_shape.test_dynamic_shape_export import export_dynamic_shape_cases
 export_dynamic_shape_cases('${DYNAMIC_SHAPE_DIR}')
 "
+require_file "${DYNAMIC_SHAPE_DIR}/dyn_cat_2d.pte"
 
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.test_sdpa import (
@@ -221,6 +251,9 @@ run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \
 "${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}"
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
+run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 WEBGPU_TEST_HEAVY=1 \
+    "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" \
+    --gtest_filter=DynamicShape.SliceCrosses2dDispatchBoundary:DynamicShape.CatCrosses2dDispatchBoundary:DynamicShape.SliceDualStoreWritesBothDestinations
 "${BIN_DIR}/webgpu_scratch_buffer_test"
 "${BIN_DIR}/webgpu_dispatch_2d_test"
 "${BIN_DIR}/webgpu_compute_dispatch_test"
@@ -238,4 +271,11 @@ $PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests
   --output "${OP_TEST_DIR}"
 cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}"
 "${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json"
+CAT_2D_TEST_DIR="/tmp/webgpu_cat_2d_test"
+WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE \
+  -m executorch.backends.webgpu.test.op_tests.generate_op_tests \
+  --output "${CAT_2D_TEST_DIR}" --ops cat
+run_with_required_device env WEBGPU_REQUIRE_DEVICE=1 \
+  "${BIN_DIR}/webgpu_op_test" \
+  --manifest "${CAT_2D_TEST_DIR}/manifest.json"
 echo "=== WebGPU op-test framework on Dawn: passed ==="

backends/webgpu/test/BUCK

@@ -19,6 +19,21 @@ fbcode_target(
     ],
 )
 
+fbcode_target(
+    _kind = python_unittest,
+    name = "test_et_vk_sdpa",
+    srcs = [
+        "ops/test_et_vk_sdpa.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/vulkan:custom_ops_lib",
+        "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
+        "//executorch/backends/vulkan:vulkan_preprocess",
+        "//executorch/exir:lib",
+    ],
+)
+
 non_fbcode_target(
     _kind = runtime.python_test,
     name = "test_webgpu_artifact_manifest",
@@ -42,6 +57,18 @@ fbcode_target(
     ],
 )
 
+fbcode_target(
+    _kind = python_unittest,
+    name = "test_rope_hf_single",
+    srcs = [
+        "ops/test_rope_hf_single.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/vulkan:vulkan_preprocess",
+    ],
+)
+
 fbcode_target(
     _kind = runtime.python_library,
     name = "tester",
@@ -60,3 +87,30 @@ fbcode_target(
         "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest",
     ],
 )
+
+fbcode_target(
+    _kind = python_unittest,
+    name = "test_wgsl_codegen",
+    srcs = ["test_wgsl_codegen.py"],
+    deps = [
+        "fbsource//third-party/pypi/pyyaml:pyyaml",
+    ],
+)
+
+fbcode_target(
+    _kind = runtime.python_test,
+    name = "test_update_cache",
+    srcs = ["ops/test_update_cache.py"],
+    preload_deps = [
+        "//executorch/extension/llm/custom_ops:custom_ops_aot_lib",
+        "//executorch/extension/llm/custom_ops:custom_ops_aot_py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
+        "//executorch/backends/vulkan/serialization:lib",
+        "//executorch/backends/vulkan:vulkan_preprocess",
+        "//executorch/examples/models/gemma4:webgpu_support",
+        "//executorch/exir:lib",
+    ],
+)

backends/webgpu/test/native/RequiredDevicePolicy.h

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+namespace executorch::backends::webgpu {
+
+inline int required_device_failure_exit_code(bool required) {
+  return required ? 1 : 0;
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/test/native/test_dispatch_2d.cpp

@@ -12,6 +12,9 @@
 #include <executorch/backends/webgpu/runtime/WebGPUBackend.h>
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/cat/CatDispatch.h>
+#include <executorch/backends/webgpu/runtime/ops/slice/SliceDispatch.h>
+#include <executorch/backends/webgpu/test/native/RequiredDevicePolicy.h>
 
 #include <gtest/gtest.h>
 
@@ -20,6 +23,9 @@
 #include <utility>
 #include <vector>
 
+using executorch::backends::webgpu::required_device_failure_exit_code;
+using executorch::backends::webgpu::set_cat_dispatch_grid;
+using executorch::backends::webgpu::set_slice_dispatch_grid;
 using executorch::backends::webgpu::WebGPUDispatch;
 using executorch::backends::webgpu::WebGPUGraph;
 using executorch::backends::webgpu::utils::DispatchRange;
@@ -65,6 +71,38 @@ TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) {
   EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test"));
 }
 
+TEST(SliceDispatchGrid, RestoresBothDimensionsAcrossResize) {
+  WebGPUGraph graph;
+  const size_t dispatch_index = graph.add_dispatch(WebGPUDispatch{});
+
+  for (const WgCount grid :
+       {WgCount{256u, 256u}, WgCount{65535u, 1u}, WgCount{256u, 256u}}) {
+    set_slice_dispatch_grid(graph, dispatch_index, grid);
+    EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_x, grid.x);
+    EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_y, grid.y);
+  }
+}
+
+TEST(CatDispatchGrid, RestoresBothDimensionsAcrossResize) {
+  WebGPUGraph graph;
+  const size_t dispatch_index = graph.add_dispatch(WebGPUDispatch{});
+
+  for (const WgCount grid :
+       {WgCount{257u, 256u}, WgCount{65535u, 1u}, WgCount{257u, 256u}}) {
+    set_cat_dispatch_grid(graph, dispatch_index, grid);
+    EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_x, grid.x);
+    EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_y, grid.y);
+  }
+}
+
+TEST(RequiredDevicePolicy, DefaultDeviceFailureRemainsASkip) {
+  EXPECT_EQ(required_device_failure_exit_code(false), 0);
+}
+
+TEST(RequiredDevicePolicy, RequiredDeviceFailureIsAnError) {
+  EXPECT_NE(required_device_failure_exit_code(true), 0);
+}
+
 void expect_grid(const WgCount& grid, uint32_t x, uint32_t y) {
   EXPECT_EQ(grid.x, x);
   EXPECT_EQ(grid.y, y);