Lower amax over trailing dims (#24735) (#24735)

Summary:

## Ulterior Motive

Run Whisper's log-mel front end on ggml, where `ReplaceGlobalMaxPass`
rewrites the global `max()` as `amax`.

## Rationale

**What**: Lower `aten.amax.default` over trailing dims on ggml.

**Why**: ggml has no max reduction, and the graph builder rejected
`amax`.

## Details

- ggml's `amax` reuses the `trailing_rows` helper shared with `mean` and
`sum`, and gathers each row's `ggml_argmax` element with
`ggml_get_rows`.
- Unlike `ggml_pool_1d` max, which starts from `-FLT_MAX`, this keeps an
all `-inf` row at `-inf`. NaN does not propagate, as in ggml's argmax.
- A non-trailing dim is rejected by name.

Authored with AI using Claude Code.

Reviewed By: SS-JIA

Differential Revision: D129976771

Author: JakeStevens

backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp

@@ -709,6 +709,8 @@ void GraphBuilder::lower(const Node& node) {
     lower_mean(node);
   } else if (key == "aten.sum.dim_IntList") {
     lower_sum(node);
+  } else if (key == "aten.amax.default") {
+    lower_amax(node);
   } else if (key == "aten.slice_copy.Tensor" || key == "aten.slice.Tensor") {
     lower_slice(node, key == "aten.slice_copy.Tensor");
   } else if (key == "aten.select_copy.int" || key == "aten.select.int") {

backends/native/runtime/engine/ggml/GgmlGraphBuilder.h

@@ -241,8 +241,9 @@ class GraphBuilder {
   void lower_max_pool2d(const Node& node);
   void lower_mean(const Node& node);
   void lower_sum(const Node& node);
-  // The `self` of a mean or sum over trailing dims as ggml rows of the reduced
-  // elements, and the result sizes.
+  void lower_amax(const Node& node);
+  // The `self` of a mean, sum or amax over trailing dims as ggml rows of the
+  // reduced elements, and the result sizes.
   std::pair<ggml_tensor*, std::vector<int64_t>> trailing_rows(
       const Node& node,
       const std::string& op);

backends/native/runtime/engine/ggml/GgmlLowerElementwise.cpp

@@ -294,6 +294,20 @@ void GraphBuilder::lower_sum(const Node& node) {
   set_single_output(node, TensorRef{result, std::move(sizes)});
 }
 
+// Gathers each row's argmax element. Unlike max pooling, which starts from
+// -FLT_MAX, this keeps an all -inf row at -inf. NaN does not propagate.
+void GraphBuilder::lower_amax(const Node& node) {
+  auto [rows, sizes] = trailing_rows(node, "amax");
+  ggml_context* ctx = instance_.context.get();
+  const int64_t width = rows->ne[0];
+  const int64_t count = rows->ne[1];
+  ggml_tensor* index = ggml_reshape_2d(ctx, ggml_argmax(ctx, rows), 1, count);
+  ggml_tensor* values =
+      ggml_get_rows(ctx, ggml_reshape_3d(ctx, rows, 1, width, count), index);
+  ggml_tensor* result = reshape(ctx, values, sizes);
+  set_single_output(node, TensorRef{result, std::move(sizes)});
+}
+
 // Comparisons produce 0/1 floats: a <= b is 1 - step(a - b) and a < b is
 // step(b - a).
 void GraphBuilder::lower_compare(const Node& node, bool strict) {

backends/native/test/runtime/ggml/test_ggml_engine.cpp

@@ -2032,6 +2032,72 @@ TEST(GgmlEngineTest, RejectsSumOverANonTrailingDim) {
       program.finish("y"), "sum supports trailing dimensions only");
 }
 
+TEST(GgmlEngineTest, AmaxOverTrailingDims) {
+  const std::vector<int64_t> shape = {2, 3, 4, 5};
+  struct Case {
+    std::vector<int64_t> dims;
+    bool keepdim;
+    size_t width;
+    std::vector<int64_t> out_shape;
+  };
+  constexpr float kInf = std::numeric_limits<float>::infinity();
+  for (const Case& c :
+       {Case{{-1}, true, 5, {2, 3, 4, 1}},
+        Case{{2, 3}, false, 20, {2, 3}},
+        Case{{0, 1, 2, 3}, false, 120, {}}}) {
+    for (const size_t cached : {size_t{0}, size_t{1}}) {
+      SCOPED_TRACE(std::to_string(c.width) + std::to_string(cached));
+      TestProgram program;
+      program.input("x", fbs::ScalarType::FLOAT, shape);
+      program.call(
+          "y",
+          "torch.ops.aten.amax.default",
+          {{"self", std::string("x")},
+           {"dim", c.dims},
+           {"keepdim", BoolValue{c.keepdim}}},
+          fbs::ScalarType::FLOAT,
+          c.out_shape);
+      auto host = GgmlEngineHost::create(
+          GgmlEngineOptions{.max_cached_graphs = cached});
+      auto context = make_context(*host, program.finish("y"));
+      auto executable = context->compile("forward");
+      for (const float scale : {1.0f, -3.0f}) {
+        std::vector<float> x(120);
+        for (size_t i = 0; i < x.size(); ++i) {
+          x[i] = scale * std::cos(static_cast<float>(i) * 1.3f);
+        }
+        // The first row is all -inf, and the last holds +inf.
+        std::fill_n(x.begin(), c.width, -kInf);
+        x.back() = kInf;
+        executable->set_input(0, x.data(), x.size(), ScalarType::Float);
+        executable->execute();
+        std::vector<float> y(120 / c.width);
+        executable->get_output(0, y.data(), y.size(), ScalarType::Float);
+        std::vector<float> expected(y.size(), -kInf);
+        for (size_t i = 0; i < x.size(); ++i) {
+          expected[i / c.width] = std::max(expected[i / c.width], x[i]);
+        }
+        EXPECT_EQ(y, expected);
+      }
+    }
+  }
+}
+
+TEST(GgmlEngineTest, RejectsAmaxOverANonTrailingDim) {
+  TestProgram program;
+  program.input("x", fbs::ScalarType::FLOAT, {2, 3, 4});
+  program.call(
+      "y",
+      "torch.ops.aten.amax.default",
+      {{"self", std::string("x")},
+       {"dim", std::vector<int64_t>{1}},
+       {"keepdim", BoolValue{false}}},
+      fbs::ScalarType::FLOAT,
+      {2, 4});
+  expect_compile_error(
+      program.finish("y"), "amax supports trailing dimensions only");
+}
+
 TEST(GgmlEngineTest, GpuUnavailableOnCpuVariant) {
   GgmlEngineOptions options;
   options.device = GgmlDevice::Gpu;

backends/native/test/runtime/ggml/test_ggml_ops.py

@@ -879,6 +879,12 @@ def test_sum_trailing_dims(self):
         )
         self._check(Lambda(lambda x: x.sum(dim=[2, 3])), (torch.randn(2, 3, 4, 5),))
 
+    def test_amax_trailing_dims(self):
+        self._check(
+            Lambda(lambda x: x.amax(dim=-1, keepdim=True)), (torch.randn(2, 3, 16, 32),)
+        )
+        self._check(Lambda(lambda x: x - x.amax(dim=[0, 1])), (torch.randn(80, 300),))
+
     def test_linear(self):
         self._check(Linear(bias=False), (torch.randn(4, 32),), **_MATMUL_TOL)