Skip to content

Commit 6f4c21e

Browse files
JCNTHSS-JIA
JCNTH
authored and
SS-JIA
committed
[ExecuTorch][WebGPU] Op-test codegen framework (cases.py -> generated .pte+golden -> gtest driver)
Pull Request resolved: pytorch#20339 A manifest-driven op-test framework for the WebGPU backend, mirroring Vulkan's `op_tests/cases.py` (declarative per-op suites) but with a torch-computed golden loaded in C++, since the native test binary has no ATen. An op fits when it is stateless and expressible as one `module(inputs) -> golden` forward; stateful KV-cache ops (`sdpa`, `update_cache`) stay hand-written. Lands first as the shared foundation for the following op test-diffs — adding an op's test becomes one `cases.py` entry. Composition: - `test_suite.py` — schema (`WebGPUTestSuite`/`Case`/`InputSpec`) + a `register_op_test` decorator; per-case `required`/`heavy`/`golden_fn`, per-suite `golden_dtype`. - `cases.py` — the declarative suites; registers the landed `add` + `rms_norm`; later ops append one entry each. - `generate_op_tests.py` — per case: export via `VulkanPartitioner` to `.pte`, compute the fp64 torch golden (dual-oracle gate), serialize inputs+golden as fp32, emit `manifest.json`. - `op_test_driver.cpp` + `driver_util.{h,cpp}` — generic gtest driver: one test per manifest entry, runs forward on-device, abs/rel + shape + reconciliation checks. - `CMakeLists`/`ci.sh` — `webgpu_op_test` + device-free `webgpu_op_test_util_test`, wired into the Dawn(Tint)+SwiftShader CI. ghstack-source-id: 394712836 @exported-using-ghexport Differential Revision: [D108816389](https://our.internmc.facebook.com/intern/diff/D108816389/)
1 parent 5241b4e commit 6f4c21e

13 files changed

Lines changed: 1149 additions & 0 deletions

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,4 +148,38 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
148148
add_webgpu_native_test(
149149
webgpu_update_cache_test test/native/test_update_cache.cpp
150150
)
151+
152+
# Manifest-driven op-test framework: a generic gtest driver (webgpu_op_test) +
153+
# its device-free util unit test. GTest needs -DEXECUTORCH_BUILD_TESTS=ON.
154+
if(NOT TARGET GTest::gtest)
155+
find_package(GTest QUIET)
156+
endif()
157+
if(TARGET GTest::gtest)
158+
# Reuse add_webgpu_native_test for the backend link + frameworks + flags;
159+
# add only driver_util, GTest, and the header-only nlohmann/json include.
160+
add_webgpu_native_test(webgpu_op_test test/op_tests/op_test_driver.cpp)
161+
target_sources(webgpu_op_test PRIVATE test/op_tests/driver_util.cpp)
162+
target_link_libraries(webgpu_op_test PRIVATE GTest::gtest)
163+
target_include_directories(
164+
webgpu_op_test
165+
PRIVATE "${EXECUTORCH_ROOT}/third-party/json/single_include"
166+
)
167+
168+
# Device-free util unit test: no backend/Dawn link (pure manifest/tolerance
169+
# helpers), so it does NOT use the native-test helper.
170+
add_executable(
171+
webgpu_op_test_util_test test/op_tests/test_driver_util.cpp
172+
test/op_tests/driver_util.cpp
173+
)
174+
target_include_directories(
175+
webgpu_op_test_util_test
176+
PRIVATE $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
177+
"${EXECUTORCH_ROOT}/third-party/json/single_include"
178+
)
179+
target_link_libraries(
180+
webgpu_op_test_util_test PRIVATE GTest::gtest GTest::gtest_main
181+
)
182+
target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions)
183+
set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17)
184+
endif()
151185
endif()

‎backends/webgpu/scripts/test_webgpu_native_ci.sh‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,3 +165,25 @@ fi
165165
[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test"
166166

167167
echo "=== WebGPU native tests on Dawn: all run targets passed ==="
168+
169+
# ── Op-test codegen framework: generate manifest → build → run (Dawn+SwiftShader) ──
170+
# Reconfigure the SAME build dir adding GTest (EXECUTORCH_BUILD_TESTS=ON), then run
171+
# every op in cases.py against its torch golden. Self-skips if the generator can't run.
172+
OP_TEST_DIR="/tmp/webgpu_op_tests"
173+
if $PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests \
174+
--output "${OP_TEST_DIR}"; then
175+
echo "=== Reconfigure with GTest + build/run op-test framework ==="
176+
cmake -DEXECUTORCH_BUILD_TESTS=ON -B "${BUILD_DIR}" "${EXECUTORCH_ROOT}"
177+
OP_DEFINED="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)"
178+
if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test_util_test; then
179+
cmake --build "${BUILD_DIR}" --target webgpu_op_test_util_test -j"${NPROC}"
180+
"${BIN_DIR}/webgpu_op_test_util_test"
181+
fi
182+
if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test; then
183+
cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}"
184+
"${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json"
185+
fi
186+
echo "=== WebGPU op-test framework on Dawn: passed ==="
187+
else
188+
echo "WARN: op-test manifest generation failed (needs the executorch wheel); skipping"
189+
fi

‎backends/webgpu/test/op_tests/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.

‎backends/webgpu/test/op_tests/cases.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
"""WebGPU op-test cases.
8+
9+
Declarative per-op suites for the manifest-driven op-test framework, mirroring the
10+
Vulkan op-test authoring ergonomics. Each op reuses its `nn.Module` + generators from
11+
the per-op `test_*.py`; new ops append a `@register_op_test` entry in their own tests-diff.
12+
"""
13+
14+
import torch
15+
16+
from executorch.backends.webgpu.test.op_tests.test_suite import (
17+
Case,
18+
InputSpec,
19+
M1,
20+
M2,
21+
register_op_test,
22+
S,
23+
S1,
24+
S2,
25+
WebGPUTestSuite,
26+
XS,
27+
)
28+
from executorch.backends.webgpu.test.ops.add.test_add import (
29+
AddChainedModule,
30+
AddModule,
31+
AddSelfModule,
32+
)
33+
from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import (
34+
_CASES,
35+
_linspace_weight,
36+
_ramp,
37+
RmsNormModule,
38+
)
39+
40+
# rms_norm coverage is exactly the 14 cases the native test covered.
41+
RMS_NORM_CASES = _CASES
42+
43+
44+
def _add_factory(variant: str = "regular") -> torch.nn.Module:
45+
return {
46+
"regular": AddModule,
47+
"self": AddSelfModule,
48+
"chained": AddChainedModule,
49+
}[variant]()
50+
51+
52+
@register_op_test("add")
53+
def _add_suite() -> WebGPUTestSuite:
54+
# Same-shape numeric coverage only: broadcast adds stay export-smoke in
55+
# ops/add/test_add.py because the kernel can't broadcast.
56+
return WebGPUTestSuite(
57+
module_factory=_add_factory,
58+
cases=[
59+
Case(
60+
name="regular_2d",
61+
construct={"variant": "regular"},
62+
inputs=((M1, M2), (M1, M2)),
63+
),
64+
Case(
65+
name="regular_3d",
66+
construct={"variant": "regular"},
67+
inputs=((S, S1, S2), (S, S1, S2)),
68+
),
69+
Case(
70+
name="regular_4d",
71+
construct={"variant": "regular"},
72+
inputs=((XS, S, S1, S2), (XS, S, S1, S2)),
73+
),
74+
Case(name="self", construct={"variant": "self"}, inputs=((M1, M2),)),
75+
# "scalar" (x+3.0) is intentionally OMITTED — the WebGPU add kernel can't
76+
# do scalar/broadcast adds (0x30 at runtime); it stays export-smoke.
77+
Case(
78+
name="chained",
79+
construct={"variant": "chained"},
80+
inputs=((M1, M2), (M1, M2)),
81+
),
82+
],
83+
)
84+
85+
86+
def _rms_norm_factory(hidden: int, eps: float, weight_fn) -> torch.nn.Module:
87+
model = RmsNormModule(hidden, eps=eps)
88+
with torch.no_grad():
89+
model.weight.copy_(weight_fn(hidden))
90+
return model
91+
92+
93+
@register_op_test("rms_norm")
94+
def _rms_norm_suite() -> WebGPUTestSuite:
95+
cases = []
96+
for c in RMS_NORM_CASES:
97+
shape = c["shape"]
98+
hidden = shape[-1]
99+
weight_fn = c.get("weight_fn", _linspace_weight)
100+
input_fn = c.get("input_fn", _ramp)
101+
cases.append(
102+
Case(
103+
name=c["name"],
104+
construct={"hidden": hidden, "eps": 1e-6, "weight_fn": weight_fn},
105+
inputs=(InputSpec(shape=shape, gen=input_fn),),
106+
)
107+
)
108+
return WebGPUTestSuite(module_factory=_rms_norm_factory, cases=cases)

‎backends/webgpu/test/op_tests/driver_util.cpp‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/webgpu/test/op_tests/driver_util.h>
10+
11+
#include <algorithm>
12+
#include <cmath>
13+
#include <cstdint>
14+
#include <cstdio>
15+
#include <fstream>
16+
17+
#include <nlohmann/json.hpp>
18+
19+
namespace executorch::backends::webgpu {
20+
21+
namespace {
22+
std::string dir_of(const std::string& p) {
23+
const auto pos = p.find_last_of('/');
24+
return pos == std::string::npos ? std::string(".") : p.substr(0, pos);
25+
}
26+
std::string join(const std::string& a, const std::string& b) {
27+
return a + "/" + b;
28+
}
29+
} // namespace
30+
31+
std::vector<ManifestEntry> parse_manifest(const std::string& manifest_path) {
32+
std::ifstream in(manifest_path);
33+
if (!in.good()) {
34+
std::fprintf(
35+
stderr, "ERROR: cannot open manifest: %s\n", manifest_path.c_str());
36+
return {};
37+
}
38+
nlohmann::json j;
39+
in >> j;
40+
const std::string base = dir_of(manifest_path);
41+
42+
std::vector<ManifestEntry> out;
43+
for (const auto& e : j) {
44+
ManifestEntry m;
45+
m.op = e.at("op").get<std::string>();
46+
m.name = e.at("case").get<std::string>();
47+
m.pte = join(base, e.at("pte").get<std::string>());
48+
for (const auto& ie : e.at("inputs")) {
49+
InputRef ir;
50+
ir.path = join(base, ie.at("path").get<std::string>());
51+
ir.shape = ie.at("shape").get<std::vector<int>>();
52+
m.inputs.push_back(std::move(ir));
53+
}
54+
const auto& g = e.at("golden");
55+
m.golden.path = join(base, g.at("path").get<std::string>());
56+
m.golden.shape = g.at("shape").get<std::vector<int>>();
57+
m.golden.output_index = g.value("output_index", 0);
58+
m.atol = e.value("atol", 1e-3f);
59+
m.rtol = e.value("rtol", 1e-3f);
60+
m.required = e.value("required", true);
61+
m.heavy = e.value("heavy", false);
62+
out.push_back(std::move(m));
63+
}
64+
return out;
65+
}
66+
67+
std::vector<float> load_fp32_bin(const std::string& path, size_t numel) {
68+
FILE* f = std::fopen(path.c_str(), "rb");
69+
if (!f) {
70+
return {};
71+
}
72+
std::vector<float> g(numel);
73+
const size_t n = std::fread(g.data(), sizeof(float), numel, f);
74+
std::fclose(f);
75+
if (n != numel) {
76+
return {};
77+
}
78+
return g;
79+
}
80+
81+
bool within_tol(
82+
const float* out,
83+
const float* golden,
84+
int n,
85+
float atol,
86+
float rtol,
87+
float* max_abs,
88+
float* max_rel) {
89+
float ma = 0.0f, mr = 0.0f;
90+
bool ok = true;
91+
for (int i = 0; i < n; i++) {
92+
if (std::isnan(out[i]) || std::isnan(golden[i])) {
93+
ok = false; // NaN never passes a tolerance check
94+
continue;
95+
}
96+
const float ae = std::abs(out[i] - golden[i]);
97+
const float re = ae / std::max(std::abs(golden[i]), 1e-6f);
98+
ma = std::max(ma, ae);
99+
mr = std::max(mr, re);
100+
if (ae > atol && re > rtol) {
101+
ok = false;
102+
}
103+
}
104+
*max_abs = ma;
105+
*max_rel = mr;
106+
return ok;
107+
}
108+
109+
size_t numel(const std::vector<int>& shape) {
110+
size_t n = 1;
111+
for (int d : shape) {
112+
if (d <= 0) {
113+
return 0; // empty or malformed dim
114+
}
115+
const auto dd = static_cast<size_t>(d);
116+
if (n > SIZE_MAX / dd) {
117+
return 0; // overflow: malformed shape
118+
}
119+
n *= dd;
120+
}
121+
return n;
122+
}
123+
124+
} // namespace executorch::backends::webgpu

‎backends/webgpu/test/op_tests/driver_util.h‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstddef>
12+
#include <string>
13+
#include <vector>
14+
15+
namespace executorch::backends::webgpu {
16+
17+
struct InputRef {
18+
std::string path;
19+
std::vector<int> shape;
20+
};
21+
22+
struct GoldenRef {
23+
std::string path;
24+
std::vector<int> shape;
25+
int output_index = 0;
26+
};
27+
28+
struct ManifestEntry {
29+
std::string op;
30+
std::string name;
31+
std::string pte;
32+
std::vector<InputRef> inputs;
33+
GoldenRef golden;
34+
float atol = 1e-3f;
35+
float rtol = 1e-3f;
36+
bool required = true; // a missing .pte is a FAIL, not a silent skip
37+
bool heavy = false; // export-gated behind WEBGPU_TEST_HEAVY (never required)
38+
};
39+
40+
/// Parse manifest.json; resolve every relative path against the manifest's dir.
41+
std::vector<ManifestEntry> parse_manifest(const std::string& manifest_path);
42+
43+
/// Load raw little-endian fp32; empty on size/IO mismatch.
44+
std::vector<float> load_fp32_bin(const std::string& path, size_t numel);
45+
46+
/// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at
47+
/// |golden|=1e-6). Sets the reported maxima; true iff all elements pass.
48+
bool within_tol(
49+
const float* out,
50+
const float* golden,
51+
int n,
52+
float atol,
53+
float rtol,
54+
float* max_abs,
55+
float* max_rel);
56+
57+
size_t numel(const std::vector<int>& shape);
58+
59+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)