From a1092a2310ab49bd6e0c8344033e58cc09b71324 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 15:34:01 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/test/op_tests/cases.py | 35 ++++++++++++++ backends/webgpu/test/op_tests/driver_util.cpp | 14 ++++++ backends/webgpu/test/op_tests/driver_util.h | 5 +- .../webgpu/test/op_tests/generate_op_tests.py | 11 +++++ .../webgpu/test/op_tests/op_test_driver.cpp | 17 +++++++ backends/webgpu/test/ops/test_argmax.py | 46 +++++++++++++++++++ 6 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 backends/webgpu/test/ops/test_argmax.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 95343de3f74..777111a808f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -35,6 +35,12 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule +from executorch.backends.webgpu.test.ops.test_argmax import ( + argmax_tie_gen, + argmin_tie_gen, + ArgmaxModule, + ArgminModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule @@ -1023,3 +1029,32 @@ def _dequant_const_suite() -> WebGPUTestSuite: ], golden_dtype="float32", ) + + +@register_op_test("argmax") +def _argmax_suite() -> WebGPUTestSuite: + # Last-dim argmax -> int64 index. randn puts the max at an INTERIOR index + # (exercises the walk); the tie case pins the strict-`>` first-occurrence. + return WebGPUTestSuite( + module_factory=lambda: ArgmaxModule(), + cases=[ + Case(name="2d", inputs=((M1, M2),)), + Case(name="3d", inputs=((S, S1, S2),)), + Case(name="tie", inputs=(InputSpec(shape=(3, 6), gen=argmax_tie_gen),)), + ], + golden_dtype="float32", + ) + + +@register_op_test("argmin") +def _argmin_suite() -> WebGPUTestSuite: + # Last-dim argmin -> int64 index; randn interior min + strict-`<` tie case. + return WebGPUTestSuite( + module_factory=lambda: ArgminModule(), + cases=[ + Case(name="2d", inputs=((M1, M2),)), + Case(name="3d", inputs=((S, S1, S2),)), + Case(name="tie", inputs=(InputSpec(shape=(3, 6), gen=argmin_tie_gen),)), + ], + golden_dtype="float32", + ) diff --git a/backends/webgpu/test/op_tests/driver_util.cpp b/backends/webgpu/test/op_tests/driver_util.cpp index 3f9a1d21a32..d45a183d654 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -93,6 +93,20 @@ std::vector load_int8_bin(const std::string& path, size_t numel) { return g; } +std::vector load_int64_bin(const std::string& path, size_t numel) { + FILE* f = std::fopen(path.c_str(), "rb"); + if (!f) { + return {}; + } + std::vector g(numel); + const size_t n = std::fread(g.data(), sizeof(int64_t), numel, f); + std::fclose(f); + if (n != numel) { + return {}; + } + return g; +} + bool within_tol( const float* out, const float* golden, diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index 22936629774..c12850fc2d1 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -23,7 +23,7 @@ struct GoldenRef { std::string path; std::vector shape; int output_index = 0; - std::string dtype = "float32"; // "float32" or "int8" (int8-output ops) + std::string dtype = "float32"; // "float32" | "int8" | "int64" (argmax index) }; struct ManifestEntry { @@ -47,6 +47,9 @@ std::vector load_fp32_bin(const std::string& path, size_t numel); /// Load raw int8 (one byte per element); empty on size/IO mismatch. std::vector load_int8_bin(const std::string& path, size_t numel); +/// Load raw little-endian int64 (8B/elem); empty on size/IO mismatch. +std::vector load_int64_bin(const std::string& path, size_t numel); + /// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at /// |golden|=1e-6). Sets the reported maxima; true iff all elements pass. bool within_tol( diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index a18fd824c74..d194d746f89 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -79,6 +79,10 @@ def _write_int8(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" None: + t.detach().contiguous().cpu().numpy().astype(" list[dict]: """Export one case; write its .pte + input/golden .bin(s). Returns one manifest entry per output tensor (multi-output ops emit N; single-output ops emit one, @@ -125,10 +129,15 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d for out_index in range(n_out): raw = golden_outs[out_index] is_int8 = raw.dtype == torch.int8 + is_int64 = raw.dtype in (torch.int64, torch.int32) if is_int8: # int8-output ops (quantize): byte-exact golden, no fp32 cast/oracle. out_t = raw out_dtype = "int8" + elif is_int64: + # int64-index ops (argmax/argmin): exact golden, no fp32 cast/oracle. + out_t = raw.to(torch.int64) + out_dtype = "int64" else: out_t = raw.to(torch.float32) out_dtype = "float32" @@ -146,6 +155,8 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d golden_rel = f"{case_id}{suffix}.golden.bin" if is_int8: _write_int8(out_t, os.path.join(out_dir, golden_rel)) + elif is_int64: + _write_int64(out_t, os.path.join(out_dir, golden_rel)) else: _write_fp32(out_t, os.path.join(out_dir, golden_rel)) entries.append( diff --git a/backends/webgpu/test/op_tests/op_test_driver.cpp b/backends/webgpu/test/op_tests/op_test_driver.cpp index 5ad7afb507e..f50987664e6 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -102,6 +102,23 @@ class OpCase : public ::testing::Test { EXPECT_EQ(mism, -1) << "int8 mismatch at index " << mism << ": out=" << (mism >= 0 ? int(out_p[mism]) : 0) << " golden=" << (mism >= 0 ? int(golden[mism]) : 0); + } else if (e_.golden.dtype == "int64") { + // int64-index ops (argmax/argmin) compare exact: indices are discrete. + auto golden = load_int64_bin(e_.golden.path, gn); + ASSERT_FALSE(golden.empty()) + << "missing/short golden: " << e_.golden.path; + const int64_t* out_p = out_tensor.const_data_ptr(); + int mism = -1; + for (size_t i = 0; i < gn; i++) { + if (out_p[i] != golden[i]) { + mism = static_cast(i); + break; + } + } + EXPECT_EQ(mism, -1) << "int64 mismatch at index " << mism << ": out=" + << (mism >= 0 ? (long long)out_p[mism] : 0) + << " golden=" + << (mism >= 0 ? (long long)golden[mism] : 0); } else { auto golden = load_fp32_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) diff --git a/backends/webgpu/test/ops/test_argmax.py b/backends/webgpu/test/ops/test_argmax.py new file mode 100644 index 00000000000..73aed92ca23 --- /dev/null +++ b/backends/webgpu/test/ops/test_argmax.py @@ -0,0 +1,46 @@ +# 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. + +"""`aten.argmax.default` / `aten.argmin.default` modules for the op-test framework. + +Last-dim arg-reduction -> int64 index. The kernel writes an int32 index (the AOT +downcasts the int64 output), which `copy_outputs` widens to the int64 program +output. `golden_dtype="float32"` so the golden's fp32 argmax matches the kernel's +fp32 argmax exactly (avoids an fp32-vs-fp64 tie-break discriminator flip). +""" + +import torch + + +class ArgmaxModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.argmax(x, dim=-1) + + +class ArgminModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.argmin(x, dim=-1) + + +def argmax_tie_gen(shape): + # Repeated max in the last dim (idx 1 and 3); the FIRST wins (index 1) — + # discriminates the strict-`>` tie-break from a `>=` (last-wins) bug. + assert shape[-1] >= 4, "tie generator writes idx 1 and 3; last dim must be >= 4" + base = torch.arange(shape[-1], dtype=torch.float32) * 0.01 + t = base.expand(shape).contiguous() + t[..., 1] = 10.0 + t[..., 3] = 10.0 + return t + + +def argmin_tie_gen(shape): + # Repeated min in the last dim (idx 1 and 3); the FIRST wins (index 1). + assert shape[-1] >= 4, "tie generator writes idx 1 and 3; last dim must be >= 4" + base = 5.0 + torch.arange(shape[-1], dtype=torch.float32) * 0.01 + t = base.expand(shape).contiguous() + t[..., 1] = -10.0 + t[..., 3] = -10.0 + return t