Skip to content

Commit 2bbe265

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] et_vk.prepack test suite (export + native golden)
Pull Request resolved: #20292 Test suite for the `et_vk.prepack` constant-materialization op, split into its own diff (op below, tests above) per the per-op test-split convention. The prepack op is how a serialized constant becomes a GPU tensor: the constant arrives as a CPU-side reference (sizes + a pointer into the .pte bytes), and the prepack node is the sole materialization — one CPU->GPU transfer straight into the consumer's buffer. The model `M(x) = x + w` (w a constant) routes `w` through a prepack node, so the delegate must run the materialization for the output to equal `x + w` rather than `x + 0`. ghstack-source-id: 395555139 @exported-using-ghexport Differential Revision: [D108678631](https://our.internmc.facebook.com/intern/diff/D108678631/)
1 parent 229da78 commit 2bbe265

4 files changed

Lines changed: 276 additions & 0 deletions

File tree

‎backends/webgpu/scripts/test_webgpu_native_ci.sh‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ ROPE_XK_GOLDEN="/tmp/webgpu_rope_xk_golden.bin"
5757
ROPE_DECODE_MODEL="/tmp/webgpu_rope_decode.pte"
5858
ROPE_DECODE_XQ_GOLDEN="/tmp/webgpu_rope_decode_xq_golden.bin"
5959
ROPE_DECODE_XK_GOLDEN="/tmp/webgpu_rope_decode_xk_golden.bin"
60+
PREPACK_MODEL="/tmp/webgpu_prepack.pte"
61+
PREPACK_GOLDEN="/tmp/webgpu_prepack_golden.bin"
62+
PREPACK2_MODEL="/tmp/webgpu_prepack_two_const.pte"
63+
PREPACK2_GOLDEN="/tmp/webgpu_prepack_two_const_golden.bin"
64+
PREPACK_TIED_MODEL="/tmp/webgpu_prepack_tied_const.pte"
65+
PREPACK_TIED_GOLDEN="/tmp/webgpu_prepack_tied_const_golden.bin"
6066

6167
$PYTHON_EXECUTABLE -c "
6268
from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import export_all_quantized_linear_models
@@ -75,6 +81,13 @@ export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}')
7581
export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode')
7682
" || echo "WARN: rope export failed; apply_rotary_emb configs will FAIL in webgpu_native_test"
7783

84+
$PYTHON_EXECUTABLE -c "
85+
from executorch.backends.webgpu.test.ops.prepack.test_prepack import export_prepack_model, export_prepack_two_const_model, export_prepack_tied_const_model
86+
export_prepack_model('${PREPACK_MODEL}', '${PREPACK_GOLDEN}')
87+
export_prepack_two_const_model('${PREPACK2_MODEL}', '${PREPACK2_GOLDEN}')
88+
export_prepack_tied_const_model('${PREPACK_TIED_MODEL}', '${PREPACK_TIED_GOLDEN}')
89+
" || echo "WARN: prepack export failed; prepack configs will FAIL in webgpu_native_test"
90+
7891
$PYTHON_EXECUTABLE -c "
7992
from executorch.backends.webgpu.test.ops.dispatch_order.test_dispatch_order import export_dispatch_order_cases
8093
export_dispatch_order_cases('${DISPATCH_ORDER_DIR}')
@@ -172,6 +185,12 @@ if [[ -x "${BIN_DIR}/webgpu_native_test" ]] &&
172185
WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \
173186
WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \
174187
WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \
188+
WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \
189+
WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \
190+
WEBGPU_TEST_PREPACK2_MODEL="${PREPACK2_MODEL}" \
191+
WEBGPU_TEST_PREPACK2_GOLDEN="${PREPACK2_GOLDEN}" \
192+
WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \
193+
WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \
175194
"${BIN_DIR}/webgpu_native_test"
176195
else
177196
echo "(skipping webgpu_native_test: executorch wheel absent — exports did not run)"

‎backends/webgpu/test/ops/prepack/__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/ops/prepack/test_prepack.py‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
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+
"""Constant-tensor prepack (`et_vk.prepack`) export + golden for the WebGPU
8+
backend.
9+
10+
The VulkanPartitioner wraps every constant feeding a delegated op in an
11+
`et_vk.prepack.default` node that materializes the constant into a GPU buffer at
12+
init. Model `M(x) = x + w` (w a constant) routes `w` through prepack, so the
13+
delegate must run the prepack copy for the output to equal `x + w` rather than
14+
`x + 0 = x`. The input is a deterministic /16 ramp so the native binary
15+
reconstructs it bit-for-bit; the torch-computed golden is written for the native
16+
binary to compare (it has no ATen).
17+
"""
18+
19+
import unittest
20+
21+
import executorch.backends.vulkan.custom_ops_lib # noqa: F401
22+
23+
import torch
24+
from executorch.backends.vulkan import VulkanPartitioner
25+
from executorch.exir import to_edge_transform_and_lower
26+
27+
# 4x4 constant weight, small enough to dump and reason about by hand.
28+
N = 4
29+
30+
31+
class _AddConst(torch.nn.Module):
32+
def __init__(self) -> None:
33+
super().__init__()
34+
# arange weight: non-zero everywhere so an unrun prepack (out = x + 0 = x)
35+
# is unambiguously distinguishable from a correct one (out = x + w).
36+
self.w = torch.nn.Parameter(
37+
torch.arange(N * N, dtype=torch.float32).reshape(N, N)
38+
)
39+
40+
def forward(self, x: torch.Tensor) -> torch.Tensor:
41+
return x + self.w
42+
43+
44+
class _AddTwoConst(torch.nn.Module):
45+
# Two constants => two prepack nodes (the multi-copy path E2E Llama needs);
46+
# add-only so it stays delegated with just this stack's registered ops.
47+
def __init__(self) -> None:
48+
super().__init__()
49+
self.w1 = torch.nn.Parameter(
50+
torch.arange(N * N, dtype=torch.float32).reshape(N, N)
51+
)
52+
self.w2 = torch.nn.Parameter(
53+
torch.arange(N * N, dtype=torch.float32).reshape(N, N) * 0.5 - 3.0
54+
)
55+
56+
def forward(self, x: torch.Tensor) -> torch.Tensor:
57+
return x + self.w1 + self.w2
58+
59+
60+
class _AddTiedConst(torch.nn.Module):
61+
# Two BYTE-IDENTICAL constants => two prepack nodes sharing ONE SHA256
62+
# named-data key (tied/duplicate weights). Exercises the prepack handler
63+
# materializing the same key twice (independent get_data + Free per call).
64+
def __init__(self) -> None:
65+
super().__init__()
66+
self.w1 = torch.nn.Parameter(
67+
torch.arange(N * N, dtype=torch.float32).reshape(N, N)
68+
)
69+
self.w2 = torch.nn.Parameter(
70+
torch.arange(N * N, dtype=torch.float32).reshape(N, N)
71+
)
72+
# Pin the tied premise; the dedup to one key is assumed, not asserted.
73+
assert torch.equal(self.w1, self.w2)
74+
75+
def forward(self, x: torch.Tensor) -> torch.Tensor:
76+
return x + self.w1 + self.w2
77+
78+
79+
def _inputs() -> tuple[torch.Tensor]:
80+
# ((i % 13) - 6) / 16: exact in fp32, matches test_webgpu_native.cpp.
81+
idx = torch.arange(N * N, dtype=torch.int64)
82+
x = (((idx % 13) - 6).to(torch.float32) / 16.0).reshape(N, N)
83+
return (x,)
84+
85+
86+
def _export(model, inputs):
87+
ep = torch.export.export(model.eval(), inputs)
88+
return to_edge_transform_and_lower(
89+
ep, partitioner=[VulkanPartitioner()]
90+
).to_executorch()
91+
92+
93+
class TestPrepack(unittest.TestCase):
94+
def test_export_delegates(self) -> None:
95+
# Each model must fully delegate -- every constant wrapped in a prepack
96+
# node inside a VulkanBackend delegate (single, multi-const, tied).
97+
for name, model in (
98+
("x + w", _AddConst()),
99+
("x + w1 + w2", _AddTwoConst()),
100+
("x + w + w (tied)", _AddTiedConst()),
101+
):
102+
with self.subTest(model=name):
103+
et = _export(model, _inputs())
104+
found = any(
105+
d.id == "VulkanBackend"
106+
for plan in et.executorch_program.execution_plan
107+
for d in plan.delegates
108+
)
109+
self.assertTrue(found, f"Expected a VulkanBackend delegate: {name}")
110+
111+
112+
def _write(model, pte_path: str, golden_path: str) -> None:
113+
(x,) = _inputs()
114+
golden = model.eval()(x)
115+
et = _export(model, (x,))
116+
with open(pte_path, "wb") as f:
117+
f.write(et.buffer)
118+
golden.detach().numpy().astype("<f4").tofile(golden_path)
119+
print(f"Exported {pte_path}; golden {golden_path} ({golden.numel()} floats)")
120+
121+
122+
def export_prepack_model(pte_path: str, golden_path: str) -> None:
123+
"""Write the x + w .pte + torch golden (raw LE fp32). One prepacked constant.
124+
The input is a /16 ramp reconstructed in the native test."""
125+
_write(_AddConst(), pte_path, golden_path)
126+
127+
128+
def export_prepack_two_const_model(pte_path: str, golden_path: str) -> None:
129+
"""Write the x + w1 + w2 .pte + golden. Two prepacked constants, exercising
130+
the multi-copy path."""
131+
_write(_AddTwoConst(), pte_path, golden_path)
132+
133+
134+
def export_prepack_tied_const_model(pte_path: str, golden_path: str) -> None:
135+
"""Write the x + w1 + w2 .pte + golden where w1 and w2 are BYTE-IDENTICAL,
136+
so they share one named-data key -> two prepack nodes materialize the same
137+
key (verifies per-call buffer ownership / no double-free on tied weights)."""
138+
_write(_AddTiedConst(), pte_path, golden_path)
139+
140+
141+
if __name__ == "__main__":
142+
unittest.main()

‎backends/webgpu/test/test_webgpu_native.cpp‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,74 @@ static bool test_rope(
536536
return true;
537537
}
538538

539+
static bool test_prepack(
540+
const std::string& model_path,
541+
const std::string& golden_path,
542+
const std::string& label = "x + const w") {
543+
// et_vk.prepack copy vs golden; unrun copy leaves zeros. See test_prepack.py.
544+
constexpr int n = 4;
545+
constexpr int numel = n * n;
546+
printf("\n--- Test: prepack (%s, %dx%d) ---\n", label.c_str(), n, n);
547+
548+
Module module(model_path);
549+
auto err = module.load_forward();
550+
if (err != Error::Ok) {
551+
printf("FAIL: could not load forward method (error %d)\n", (int)err);
552+
return false;
553+
}
554+
printf("Model loaded: %s\n", model_path.c_str());
555+
556+
std::vector<float> golden = load_golden(golden_path, numel);
557+
if (golden.empty()) {
558+
printf("FAIL: could not load golden %s\n", golden_path.c_str());
559+
return false;
560+
}
561+
562+
// ((i % 13) - 6) / 16: exact in fp32, matches test_prepack.py::_inputs.
563+
std::vector<float> x_data(numel);
564+
for (int i = 0; i < numel; i++) {
565+
x_data[i] = static_cast<float>((i % 13) - 6) / 16.0f;
566+
}
567+
auto x = make_tensor_ptr({n, n}, std::vector<float>(x_data));
568+
569+
auto result = module.forward({EValue(x)});
570+
if (!result.ok()) {
571+
printf("FAIL: forward failed (error %d)\n", (int)result.error());
572+
return false;
573+
}
574+
const auto& outputs = result.get();
575+
if (outputs.empty() || !outputs[0].isTensor()) {
576+
printf("FAIL: no tensor output\n");
577+
return false;
578+
}
579+
const auto& out_tensor = outputs[0].toTensor();
580+
if (out_tensor.numel() != numel) {
581+
printf(
582+
"FAIL: output numel %zu != expected %d\n",
583+
(size_t)out_tensor.numel(),
584+
numel);
585+
return false;
586+
}
587+
const float* out_data = out_tensor.const_data_ptr<float>();
588+
589+
float max_abs_err = 0.0f, max_rel_err = 0.0f;
590+
// Per-element abs-OR-rel (quant_within_tol): a global rel gate spuriously
591+
// fails near-zero outputs where rel error explodes.
592+
const bool within = quant_within_tol(
593+
out_data, golden.data(), numel, 1e-3f, 1e-3f, &max_abs_err, &max_rel_err);
594+
printf(
595+
"Max abs error: %e Max rel error: %e (checked %d elements)\n",
596+
max_abs_err,
597+
max_rel_err,
598+
numel);
599+
if (!within) {
600+
printf("FAIL: prepack exceeds tolerance 1e-3\n");
601+
return false;
602+
}
603+
printf("PASS: prepack test\n");
604+
return true;
605+
}
606+
539607
// Reconstruct _ramp_input bit-for-bit, run the op, compare to the fp64 golden.
540608
static bool test_q4gsw_config(
541609
const Q4gswConfig& cfg,
@@ -1614,6 +1682,30 @@ int main(int argc, char** argv) {
16141682
64},
16151683
};
16161684

1685+
std::string prepack_model_path, prepack_golden_path;
1686+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_MODEL")) {
1687+
prepack_model_path = env;
1688+
}
1689+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_GOLDEN")) {
1690+
prepack_golden_path = env;
1691+
}
1692+
1693+
std::string prepack2_model_path, prepack2_golden_path;
1694+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_MODEL")) {
1695+
prepack2_model_path = env;
1696+
}
1697+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK2_GOLDEN")) {
1698+
prepack2_golden_path = env;
1699+
}
1700+
1701+
std::string prepack_tied_model_path, prepack_tied_golden_path;
1702+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_MODEL")) {
1703+
prepack_tied_model_path = env;
1704+
}
1705+
if (const char* env = std::getenv("WEBGPU_TEST_PREPACK_TIED_GOLDEN")) {
1706+
prepack_tied_golden_path = env;
1707+
}
1708+
16171709
// SDPA sweep: configs self-discover their sdpa_<name>.pte/.golden.bin under
16181710
// this directory (default "" = the embedded-file root / cwd). Set
16191711
// WEBGPU_TEST_SDPA_DIR to point at the exported .pte directory (e.g. /tmp/).
@@ -1679,6 +1771,24 @@ int main(int argc, char** argv) {
16791771
}
16801772
}
16811773

1774+
if (!prepack_model_path.empty() && !prepack_golden_path.empty()) {
1775+
ok = test_prepack(prepack_model_path, prepack_golden_path) && ok;
1776+
}
1777+
1778+
if (!prepack2_model_path.empty() && !prepack2_golden_path.empty()) {
1779+
ok = test_prepack(
1780+
prepack2_model_path, prepack2_golden_path, "x + w1 + w2") &&
1781+
ok;
1782+
}
1783+
1784+
if (!prepack_tied_model_path.empty() && !prepack_tied_golden_path.empty()) {
1785+
ok = test_prepack(
1786+
prepack_tied_model_path,
1787+
prepack_tied_golden_path,
1788+
"x + w + w (tied weights, shared key)") &&
1789+
ok;
1790+
}
1791+
16821792
bool sdpa_ran = false;
16831793
bool sdpa_ok = test_sdpa_sweep(sdpa_dir, &sdpa_ran);
16841794
if (sdpa_ran) {

0 commit comments

Comments
 (0)