Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/CMakeLists.txt

@@ -161,12 +161,6 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
   add_webgpu_native_test(
     webgpu_update_cache_test test/native/test_update_cache.cpp
   )
-  add_webgpu_native_test(
-    webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
-  )
-  add_webgpu_native_test(
-    webgpu_dispatch_2d_test test/native/test_dispatch_2d.cpp
-  )
 
   # Manifest-driven op-test framework: a generic gtest driver (webgpu_op_test) +
   # its device-free util unit test. GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
@@ -200,6 +194,21 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
     )
     target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions)
     set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17)
+
+    # Dynamic-shape integration test: a gtest binary with its own main() that
+    # brings up the device once (like webgpu_op_test).
+    add_webgpu_native_test(
+      webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
+    )
+    target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
+
+    # Device-free fold unit test (gtest_main provides main; no device needed).
+    add_webgpu_native_test(
+      webgpu_dispatch_2d_test test/native/test_dispatch_2d.cpp
+    )
+    target_link_libraries(
+      webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main
+    )
   endif()
   add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 endif()

backends/webgpu/runtime/WebGPUUtils.h

@@ -11,6 +11,7 @@
 #include <webgpu/webgpu.h>
 
 #include <algorithm>
+#include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <stdexcept>
@@ -63,23 +64,31 @@ inline uint32_t queried_max_workgroups(WGPUDevice device) {
 }
 
 // Pure 2D fold of a 1D workgroup count (device-free, unit-testable): {count,1}
-// when count <= max, else {max, div_up(count, max)} so a >max workload fits the
-// per-dimension cap; throws if a 3rd dimension would be needed (out of scope).
-// The shader reconstructs the linear index from @builtin(num_workgroups).
+// when count <= max, else a near-square {x, y} with x ~ ceil(sqrt(count)) so
+// the launched grid stays close to count. A flat {max, div_up(count, max)}
+// split would leave up to ~half the workgroups inactive when count just exceeds
+// max, and inactive workgroups still cost launch/scheduling; the near-square
+// split keeps the waste to O(sqrt(count)). Throws if even a max*max grid is too
+// small (a 3rd dispatch dimension, out of scope). The shader reconstructs the
+// linear index from @builtin(num_workgroups), so any x/y factoring works.
 inline WgCount fold_workgroup_count_2d(
     uint32_t count,
     uint32_t max_count,
     const char* op_name) {
   if (count <= max_count) {
     return {count, 1u};
   }
-  uint32_t y = (count + max_count - 1) / max_count;
+  uint32_t x =
+      static_cast<uint32_t>(std::ceil(std::sqrt(static_cast<double>(count))));
+  x = std::min(x, max_count);
+  // ceil-div written overflow-safe (count >= 1 here) as count nears UINT32_MAX.
+  uint32_t y = 1u + (count - 1u) / x;
   if (y > max_count) {
     throw std::runtime_error(
         std::string("WebGPU ") + op_name +
         ": workgroup count needs a 3rd dispatch dimension (unsupported)");
   }
-  return {max_count, y};
+  return {x, y};
 }
 
 // 1D dispatch count (mirrors Vulkan div_up); throws if > device limit.
@@ -104,7 +113,7 @@ inline WgCount compute_2d_workgroup_count(
     uint32_t num_threads,
     uint32_t workgroup_size,
     const char* op_name) {
-  uint32_t count = (num_threads + workgroup_size - 1) / workgroup_size;
+  uint32_t count = div_up(num_threads, workgroup_size);
   return fold_workgroup_count_2d(
       count, queried_max_workgroups(device), op_name);
 }

backends/webgpu/test/native/test_dispatch_2d.cpp

@@ -6,83 +6,55 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// Device-free unit test for the pure 2D workgroup-count fold that lifts the
+// 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU.
+
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 
+#include <gtest/gtest.h>
+
+#include <cmath>
 #include <cstdint>
-#include <cstdio>
-#include <stdexcept>
 
 using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
 using executorch::backends::webgpu::utils::WgCount;
 
 namespace {
 
-// Device-free unit test for the pure 2D workgroup-count fold that lifts the
-// 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU.
-int g_failures = 0;
+constexpr uint32_t kMax = 65535u;
 
-void expect_fold(
-    uint32_t count,
-    uint32_t max_count,
-    uint32_t want_x,
-    uint32_t want_y) {
-  WgCount got = fold_workgroup_count_2d(count, max_count, "test");
-  bool ok = got.x == want_x && got.y == want_y &&
-      static_cast<uint64_t>(got.x) * got.y >= count;
-  printf(
-      "%s fold(%u, max=%u) = {%u, %u} (want {%u, %u})\n",
-      ok ? "PASS:" : "FAIL:",
-      count,
-      max_count,
-      got.x,
-      got.y,
-      want_x,
-      want_y);
-  if (!ok) {
-    g_failures++;
+// count <= max -> {count, 1}: the 1D fast path, byte-identical to the old path.
+TEST(DispatchFold, FastPath1D) {
+  for (uint32_t count : {1u, kMax - 1u, kMax}) {
+    const WgCount got = fold_workgroup_count_2d(count, kMax, "test");
+    EXPECT_EQ(got.x, count);
+    EXPECT_EQ(got.y, 1u);
   }
 }
 
-void expect_throw(uint32_t count, uint32_t max_count) {
-  bool threw = false;
-  try {
-    fold_workgroup_count_2d(count, max_count, "test");
-  } catch (const std::exception&) {
-    threw = true;
-  }
-  printf(
-      "%s fold(%u, max=%u) throws (needs a 3rd dispatch dimension)\n",
-      threw ? "PASS:" : "FAIL:",
-      count,
-      max_count);
-  if (!threw) {
-    g_failures++;
+// count > max -> near-square {x, y}: fits the per-dim cap, covers every
+// workgroup, and stays near-square so few invocations are inactive (launched -
+// count is O(sqrt(count)); a flat {max, div_up} split would idle up to ~half).
+TEST(DispatchFold, NearSquareFold) {
+  // Includes prefill-scale QK counts (Hq*ceil(S/4)*ceil(ctx/4)/wg) that fold:
+  // 131072 = S=2048 (32*512*512/64); 2097152 = large-S stress.
+  for (uint32_t count :
+       {kMax + 1u, 2u * kMax, 2u * kMax + 1u, 131072u, 2097152u}) {
+    const WgCount got = fold_workgroup_count_2d(count, kMax, "test");
+    const uint64_t launched = static_cast<uint64_t>(got.x) * got.y;
+    const uint32_t root =
+        static_cast<uint32_t>(std::ceil(std::sqrt(static_cast<double>(count))));
+    EXPECT_LE(got.x, kMax) << "count=" << count;
+    EXPECT_LE(got.y, kMax) << "count=" << count;
+    EXPECT_GE(launched, count) << "count=" << count;
+    EXPECT_LT(launched - count, 2ull * root)
+        << "count=" << count << " launched=" << launched;
   }
 }
 
-} // namespace
-
-int main() {
-  const uint32_t kMax = 65535u;
-  // 1D fast path: count <= max -> {count, 1}, byte-identical to the old path.
-  expect_fold(1u, kMax, 1u, 1u);
-  expect_fold(kMax - 1u, kMax, kMax - 1u, 1u);
-  expect_fold(kMax, kMax, kMax, 1u);
-  // Fold to 2D: count > max -> {max, div_up(count, max)}.
-  expect_fold(kMax + 1u, kMax, kMax, 2u);
-  expect_fold(2u * kMax, kMax, kMax, 2u);
-  expect_fold(2u * kMax + 1u, kMax, kMax, 3u);
-  // Prefill-scale QK counts (tiled grid = Hq*ceil(S/4)*ceil(ctx/4)/wg) that
-  // exceed kMax and must fold.
-  expect_fold(131072u, kMax, kMax, 3u); // S=2048: 32*512*512/64
-  expect_fold(2097152u, kMax, kMax, 33u); // deep fold (large-S stress)
-  // count > max^2 needs a 3rd dispatch dimension -> throws (out of scope).
-  expect_throw(kMax * kMax + 1u, kMax);
-
-  if (g_failures != 0) {
-    printf("\nFAIL: %d dispatch_2d fold case(s) failed\n", g_failures);
-    return 1;
-  }
-  printf("\nAll dispatch_2d fold tests passed\n");
-  return 0;
+// count > max^2 needs a 3rd dispatch dimension -> throws (out of scope).
+TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) {
+  EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test"));
 }
+
+} // namespace