[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
4 commits merged into from
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test #20723

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
db65e54
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…
JCNTH Jul 3, 2026
b2dc380
[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)
JCNTH Jul 3, 2026
f60c388
[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…
JCNTH Jul 3, 2026
5fedf64
[ExecuTorch][WebGPU] Convert remaining native tests to GTest
JCNTH Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 29 additions & 16 deletions backends/webgpu/CMakeLists.txt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
add_webgpu_native_test(
webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
)
add_webgpu_native_test(
webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
)
add_webgpu_native_test(
webgpu_update_cache_test test/native/test_update_cache.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.
# All WebGPU native tests use GTest (device-dependent ones bring up the device
# in their own main(); the fold unit test is device-free via gtest_main).
# GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
if(NOT TARGET GTest::gtest)
find_package(GTest QUIET)
endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
# Device-dependent native tests: each has its own main() that brings up the
# device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
add_webgpu_native_test(
webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
)
target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
add_webgpu_native_test(
webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
)
target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
add_webgpu_native_test(
webgpu_update_cache_test test/native/test_update_cache.cpp
)
target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
add_webgpu_native_test(
webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
)
target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
target_link_libraries(webgpu_index_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()
4 changes: 3 additions & 1 deletion backends/webgpu/runtime/WebGPUDevice.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

// TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
// Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
// Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
instance_desc.capabilities.timedWaitAnyEnable = true;
instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down
17 changes: 9 additions & 8 deletions backends/webgpu/runtime/ops/mul/BinaryOp.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
const auto& in2_tensor = graph.get_tensor(in2_id);
const auto& out_tensor = graph.get_tensor(out_id);

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
// Rank guard (NCHW backend is <= 4 dims).
if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
in1_tensor.dims.size() > kTensorMetaMaxNdim ||
in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

uint32_t wg_size =
utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

WGPUConstantEntry wg_size_constant = {};
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
bg_desc.entries = bg_entries;
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
const size_t dispatch_idx = graph.add_dispatch(
{pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

// Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
g.dispatch_at(dispatch_idx).workgroup_count_x =
utils::compute_1d_workgroup_count(
g.device(), om.numel, wg_size, "mul(resize)");
const utils::WgCount wgc = utils::compute_2d_workgroup_count(
g.device(), om.numel, wg_size, "mul(resize)");
g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
};
graph.add_tensor_resize_hook(in1_id, mul_resize);
graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down
7 changes: 5 additions & 2 deletions backends/webgpu/runtime/ops/mul/binary_mul.wgsl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let idx = gid.x;
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
if (idx >= out_meta.numel) {
return;
}
Expand Down
9 changes: 6 additions & 3 deletions backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let idx = gid.x;
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
if (idx >= out_meta.numel) {
return;
}
Expand Down
5 changes: 3 additions & 2 deletions backends/webgpu/runtime/ops/permute/Permute.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

uint32_t wg_size =
utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
device, out_meta.numel, wg_size, "permute");

WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
bg_desc.entries = bg_entries;
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

graph.add_dispatch({pipeline, bind_group, workgroup_count});
graph.add_dispatch(
{pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

wgpuShaderModuleRelease(shader);
wgpuBindGroupLayoutRelease(bgl);
Expand Down
7 changes: 5 additions & 2 deletions backends/webgpu/runtime/ops/permute/permute.wgsl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let out_bufi = gid.x;
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
if (out_bufi >= out_meta.numel) {
return;
}
Expand Down
9 changes: 6 additions & 3 deletions backends/webgpu/runtime/ops/permute/permute_wgsl.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let out_bufi = gid.x;
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>) {
// 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
if (out_bufi >= out_meta.numel) {
return;
}
Expand Down
15 changes: 14 additions & 1 deletion backends/webgpu/scripts/test_webgpu_native_ci.sh
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
-DEXECUTORCH_BUILD_WEBGPU=ON \
-DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
-DEXECUTORCH_BUILD_TESTS=ON \
-DDawn_DIR="${Dawn_DIR}" \
-DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
-DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
"${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
"${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down
60 changes: 60 additions & 0 deletions backends/webgpu/test/native/test_dispatch_2d.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
* 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.
*/

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
}
}

// 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;
}
}

// 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
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test#20723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged


4 commits merged into

from

Jul 4, 2026
 Merged


[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fold unit test
#20723
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
db65e54

[ExecuTorch][WebGPU] 2D compute dispatch tests — prefill golden + fol…


JCNTH Jul 3, 2026
b2dc380

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)


JCNTH Jul 3, 2026
f60c388

[ExecuTorch][WebGPU] Use requiredFeatures instance API on native + em…


JCNTH Jul 3, 2026
5fedf64

[ExecuTorch][WebGPU] Convert remaining native tests to GTest


JCNTH Jul 3, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


45 changes: 29 additions & 16 deletions
backends/webgpu/CMakeLists.txt

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -151,19 +151,9 @@ function(add_webgpu_native_test test_name test_src)
endfunction()

if(EXECUTORCH_BUILD_WEBGPU_TEST)
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.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.
 # All WebGPU native tests use GTest (device-dependent ones bring up the device
 # in their own main(); the fold unit test is device-free via gtest_main).
 # GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
 if(NOT TARGET GTest::gtest)
 find_package(GTest QUIET)
 endif()
Expand DownExpand Up@@ -195,12 +185,35 @@ 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).
 # Device-dependent native tests: each has its own main() that brings up the
 # device once, then RUN_ALL_TESTS(); link GTest::gtest (not gtest_main).
 add_webgpu_native_test(webgpu_native_test test/test_webgpu_native.cpp)
 target_link_libraries(webgpu_native_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dispatch_order_test test/native/test_dispatch_order.cpp
 )
 target_link_libraries(webgpu_dispatch_order_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_scratch_buffer_test test/native/test_scratch_buffer.cpp
 )
 target_link_libraries(webgpu_scratch_buffer_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_update_cache_test test/native/test_update_cache.cpp
 )
 target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest)
 add_webgpu_native_test(
 webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp
 )
 target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest)
 add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp)
 target_link_libraries(webgpu_index_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()

4 changes: 3 additions & 1 deletion
backends/webgpu/runtime/WebGPUDevice.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -90,7 +90,9 @@ WebGPUContext create_webgpu_context() {

 // TimedWaitAny lets webgpu_wait() block on futures via wgpuInstanceWaitAny.
 WGPUInstanceDescriptor instance_desc = {};
#if defined(__EMSCRIPTEN__)
 // Vendored (buck) Dawn uses the older capabilities.* API; the rig's native
 // Dawn and emscripten's emdawnwebgpu (emcc 4.0.19+) use requiredFeatures.
#if defined(WEBGPU_DAWN_INSTANCE_CAPABILITIES)
 instance_desc.capabilities.timedWaitAnyEnable = true;
 instance_desc.capabilities.timedWaitAnyMaxCount = 1;
#else
Expand Down

17 changes: 9 additions & 8 deletions
backends/webgpu/runtime/ops/mul/BinaryOp.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 const auto& in2_tensor = graph.get_tensor(in2_id);
 const auto& out_tensor = graph.get_tensor(out_id);

 // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).
 // Rank guard (NCHW backend is <= 4 dims).
 if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
 in1_tensor.dims.size() > kTensorMetaMaxNdim ||
 in2_tensor.dims.size() > kTensorMetaMaxNdim) {
Expand DownExpand Up@@ -63,8 +63,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kBinaryMulWorkgroupSizeX);
uint32_t workgroup_count =
 utils::compute_1d_workgroup_count(device, out_meta.numel, wg_size, "mul");
utils::WgCount workgroup_count =
 utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "mul");

 WGPUConstantEntry wg_size_constant = {};
 wg_size_constant.key = {"wg_size", WGPU_STRLEN};
Expand DownExpand Up@@ -165,8 +165,8 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, workgroup_count});
 const size_t dispatch_idx = graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y});

 // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch.
 WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf;
Expand DownExpand Up@@ -199,9 +199,10 @@ void mul_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
 wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
 wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm));
 g.dispatch_at(dispatch_idx).workgroup_count_x =
 utils::compute_1d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 const utils::WgCount wgc = utils::compute_2d_workgroup_count(
 g.device(), om.numel, wg_size, "mul(resize)");
 g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
 g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
 };
 graph.add_tensor_resize_hook(in1_id, mul_resize);
 graph.add_tensor_resize_hook(in2_id, mul_resize);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/mul/binary_mul.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from binary_mul.wgsl - DO NOT EDIT.
// wgsl-sha256: e7f77426cbaf48e6085e0d882522c027302ec97ef017b86a2275eed9820f7891
// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e
inline constexpr const char* kBinaryMulWGSL = R"(
@group(0) @binding(0) var<storage, read> input1: array<f32>;
@group(0) @binding(1) var<storage, read> input2: array<f32>;
Expand All@@ -32,8 +32,11 @@ struct TensorMeta {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let idx = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (idx >= out_meta.numel) {
 return;
 }
Expand Down

5 changes: 3 additions & 2 deletions
backends/webgpu/runtime/ops/permute/Permute.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -92,7 +92,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {

 uint32_t wg_size =
 utils::clamp_workgroup_size(device, kPermuteWorkgroupSizeX);
uint32_t workgroup_count = utils::compute_1d_workgroup_count(
utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
 device, out_meta.numel, wg_size, "permute");

 WGPUConstantEntry wg_size_constant = {};
Expand DownExpand Up@@ -176,7 +176,8 @@ void permute_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 bg_desc.entries = bg_entries;
 WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

 graph.add_dispatch({pipeline, bind_group, workgroup_count});
 graph.add_dispatch(
 {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y});

 wgpuShaderModuleRelease(shader);
 wgpuBindGroupLayoutRelease(bgl);
Expand Down

7 changes: 5 additions & 2 deletions
backends/webgpu/runtime/ops/permute/permute.wgsl

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,8 +18,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

9 changes: 6 additions & 3 deletions
backends/webgpu/runtime/ops/permute/permute_wgsl.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from permute.wgsl - DO NOT EDIT.
// wgsl-sha256: d34f59730cda7317589b6ed5691a1ccab8666b9c94e17ac2cb3658b036300197
// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e
inline constexpr const char* kPermuteWGSL = R"(
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
Expand All@@ -35,8 +35,11 @@ struct Params {
override wg_size: u32 = 64u;

@compute @workgroup_size(wg_size, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 let out_bufi = gid.x;
fn main(
 @builtin(global_invocation_id) gid: vec3<u32>,
 @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
 let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
 if (out_bufi >= out_meta.numel) {
 return;
 }
Expand Down

15 changes: 14 additions & 1 deletion
backends/webgpu/scripts/test_webgpu_native_ci.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,8 @@ UPDATE_CACHE_DIR="/tmp/update_cache"
UPDATE_CACHE_OK=1
INDEX_DIR="/tmp/index"
INDEX_OK=1
DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape"
DYNAMIC_SHAPE_OK=1
EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
Expand DownExpand Up@@ -111,6 +113,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde
export_all_index_models('${INDEX_DIR}')
" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; }

$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}')
" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; }

# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
$PYTHON_EXECUTABLE -c "
Expand All@@ -132,6 +139,7 @@ rm -rf "${BUILD_DIR}"
cmake \
 -DEXECUTORCH_BUILD_WEBGPU=ON \
 -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \
 -DEXECUTORCH_BUILD_TESTS=ON \
 -DDawn_DIR="${Dawn_DIR}" \
 -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
 -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
Expand All@@ -143,7 +151,7 @@ cmake \
 "${EXECUTORCH_ROOT}"

# ── Build + run every native test target that exists in this tree ────────────
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test)
TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test)
BIN_DIR="${BUILD_DIR}/backends/webgpu"

# Which targets are defined depends on which diffs are landed (native_test +
Expand DownExpand Up@@ -211,7 +219,12 @@ fi
if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then
 "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}"
fi
if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then
 "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"
fi
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU).
[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test"

echo "=== WebGPU native tests on Dawn: all run targets passed ==="

Expand Down

60 changes: 60 additions & 0 deletions
backends/webgpu/test/native/test_dispatch_2d.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
/*
 * 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.
 */

// 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>

using executorch::backends::webgpu::utils::fold_workgroup_count_2d;
using executorch::backends::webgpu::utils::WgCount;

namespace {

constexpr uint32_t kMax = 65535u;

// 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);
 }
}

// 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;
 }
}

// 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
Loading
Loading