diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 94792fd0f5c..8b0446955b2 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 @@ -1683,3 +1689,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 75d398fdacf..455f61c9532 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -108,6 +108,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 c171c114309..96998cec49e 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -24,7 +24,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 { @@ -49,6 +49,9 @@ std::vector load_int32_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 862d5eaa98d..1907f20bbe1 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -82,6 +82,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, @@ -135,10 +139,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" @@ -156,6 +165,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 ba5995ceb81..c29daf6e101 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -108,6 +108,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