[ExecuTorch][WebGPU] Op-tests for argmax/argmin + int64-golden harness

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -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 @@
         ],
         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",
+    )

backends/webgpu/test/op_tests/driver_util.cpp

@@ -93,6 +93,20 @@ std::vector<int8_t> load_int8_bin(const std::string& path, size_t numel) {
   return g;
 }
 
+std::vector<int64_t> load_int64_bin(const std::string& path, size_t numel) {
+  FILE* f = std::fopen(path.c_str(), "rb");
+  if (!f) {
+    return {};
+  }
+  std::vector<int64_t> 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,

backends/webgpu/test/op_tests/driver_util.h

@@ -23,7 +23,7 @@ struct GoldenRef {
   std::string path;
   std::vector<int> 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<float> load_fp32_bin(const std::string& path, size_t numel);
 /// Load raw int8 (one byte per element); empty on size/IO mismatch.
 std::vector<int8_t> load_int8_bin(const std::string& path, size_t numel);
 
+/// Load raw little-endian int64 (8B/elem); empty on size/IO mismatch.
+std::vector<int64_t> 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(

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("<i1").tofile(path)
 
 
+def _write_int64(t: torch.Tensor, path: str) -> None:
+    t.detach().contiguous().cpu().numpy().astype("<i8").tofile(path)
+
+
 def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> 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(

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<int64_t>();
+      int mism = -1;
+      for (size_t i = 0; i < gn; i++) {
+        if (out_p[i] != golden[i]) {
+          mism = static_cast<int>(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())

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