[Native] Gather the rows of x[:, rows] in the ggml engine

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

@@ -48,19 +48,31 @@ void GraphBuilder::lower_linear(const Node& node) {
 }
 
 void GraphBuilder::lower_row_gather(const Node& node) {
-  const TensorRef weight = tensor_arg(node, "weight", 0);
+  TensorRef weight = tensor_arg(node, "weight", 0);
   const Argument& arg = require_arg(node, "indices", 1).arg;
   ValueId indices_id = kInvalid;
+  // x[:, rows] of a [1, T, D] x, as Hugging Face selects logits_to_keep rows,
+  // gathers rows of the [T, D] matrix and keeps the leading 1.
+  bool leading_one = false;
   if (op_key(node.target) == "aten.embedding.default") {
     indices_id = arg.as_tensor().id;
   } else {
     const std::vector<ValueId> ids = arg.kind() == ArgKind::TensorList
         ? arg.as_tensor_list().ids
         : arg.as_optional_tensor_list().ids;
-    if (ids.size() != 1 || !valid(ids.front())) {
-      throw std::runtime_error("ggml: index supports one row-index tensor");
+    leading_one = ids.size() == 2 && !valid(ids.front()) &&
+        weight.sizes.size() == 3 && weight.sizes.front() == 1;
+    if (leading_one) {
+      weight.sizes.erase(weight.sizes.begin());
+    } else if (ids.size() != 1) {
+      throw std::runtime_error(
+          "ggml: index supports one row-index tensor, after at most one "
+          "leading size-1 dimension");
+    }
+    indices_id = ids.back();
+    if (!valid(indices_id)) {
+      throw std::runtime_error("ggml: index needs a row-index tensor");
     }
-    indices_id = ids.front();
   }
   const ScalarType dtype = method_.graph.value(indices_id).tensor_meta().dtype;
   if (weight.sizes.size() != 2 ||
@@ -79,6 +91,9 @@ void GraphBuilder::lower_row_gather(const Node& node) {
       reshape(ctx, indices.tensor, {ggml_nelements(indices.tensor)});
   std::vector<int64_t> sizes = indices.sizes;
   sizes.push_back(weight.sizes.at(1));
+  if (leading_one) {
+    sizes.insert(sizes.begin(), 1);
+  }
   ggml_tensor* result =
       reshape(ctx, ggml_get_rows(ctx, weight.tensor, flat), sizes);
   set_single_output(node, TensorRef{result, std::move(sizes)});

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

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -282,9 +282,18 @@
       host, program_bytes, kquant_block(kQ4KBlockBytes), options);
 }
 
-// A tensor (by name), an int, a float, an int list, or None.
-using TestArg = std::
-    variant<std::string, int64_t, double, std::vector<int64_t>, std::monostate>;
+// An optional tensor list: a tensor name per entry, or nullopt for None.
+using TestOptionalTensors = std::vector<std::optional<std::string>>;
+
+// A tensor (by name), an int, a float, an int list, None, or an optional
+// tensor list.
+using TestArg = std::variant<
+    std::string,
+    int64_t,
+    double,
+    std::vector<int64_t>,
+    std::monostate,
+    TestOptionalTensors>;
 
 // A one-method program assembled value by value; every node's output value
 // carries the node's name.
@@ -451,6 +460,22 @@
           fbs::ArgumentValue::IntListArg,
           fbs::CreateIntListArgDirect(builder_, list).Union());
     }
+    if (const auto* tensors = std::get_if<TestOptionalTensors>(&arg)) {
+      std::vector<flatbuffers::Offset<flatbuffers::String>> names;
+      std::vector<uint8_t> has_value;
+      for (const auto& name : *tensors) {
+        names.push_back(builder_.CreateString(name.value_or("")));
+        has_value.push_back(name.has_value());
+      }
+      return fbs::CreateArgument(
+          builder_,
+          fbs::ArgumentValue::OptionalTensorListArg,
+          fbs::CreateOptionalTensorListArg(
+              builder_,
+              builder_.CreateVector(names),
+              builder_.CreateVector(has_value))
+              .Union());
+    }
     return fbs::CreateArgument(
         builder_,
         fbs::ArgumentValue::FloatArg,
@@ -2291,6 +2316,39 @@
       std::runtime_error);
 }
 
+// x[:, rows] of a [1, T, D] x, as Hugging Face selects logits_to_keep rows.
+TEST(GgmlEngineTest, GathersRowsAfterALeadingOne) {
+  TestProgram program;
+  program.input("x", fbs::ScalarType::FLOAT, {1, 6, 8});
+  program.input("rows", fbs::ScalarType::LONG, {2});
+  program.call(
+      "selected",
+      "torch.ops.aten.index.Tensor",
+      {{"self", std::string("x")},
+       {"indices", TestOptionalTensors{std::nullopt, "rows"}}},
+      fbs::ScalarType::FLOAT,
+      {1, 2, 8});
+  auto host = GgmlEngineHost::create();
+  auto context = make_context(*host, program.finish("selected"));
+  auto executable = context->compile("forward");
+  std::vector<float> x(48);
+  std::iota(x.begin(), x.end(), 0.0f);
+  const std::vector<int64_t> rows = {4, 1};
+  executable->set_input(0, x.data(), x.size(), ScalarType::Float);
+  executable->set_input(1, rows.data(), rows.size(), ScalarType::Long);
+  executable->execute();
+  EXPECT_EQ(executable->output_sizes(0), (std::vector<int64_t>{1, 2, 8}));
+  std::vector<float> out(16);
+  executable->get_output(0, out.data(), out.size(), ScalarType::Float);
+  for (int64_t i = 0; i < 8; ++i) {
+    EXPECT_EQ(out[i], x[4 * 8 + i]);
+    EXPECT_EQ(out[8 + i], x[1 * 8 + i]);
+  }
+  const std::vector<int64_t> outside = {6, 0};
+  executable->set_input(1, outside.data(), outside.size(), ScalarType::Long);
+  EXPECT_ANY_THROW(executable->execute());
+}
+
 TEST(GgmlEngineTest, RejectsQuantizedGraphInputsAndOutputs) {
   TestProgram as_input;
   as_input.input("x", fbs::ScalarType::FLOAT, {1, 256});