Support Whisper positional indexing (#23903)

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

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -649,8 +649,8 @@
   const std::string key = op_key(node.target);
   if (key == "aten.linear.default") {
     lower_linear(node);
-  } else if (key == "aten.embedding.default") {
-    lower_embedding(node);
+  } else if (key == "aten.embedding.default" || key == "aten.index.Tensor") {
+    lower_row_gather(node);
   } else if (key == "aten.rms_norm.default") {
     lower_rms_norm(node);
   } else if (
@@ -713,6 +713,8 @@
   } else if (
       key == "aten.expand_copy.default" || key == "aten.expand.default") {
     lower_expand(node);
+  } else if (key == "aten.repeat.default") {
+    lower_repeat(node);
   } else if (key == "aten.cat.default") {
     lower_cat(node);
   } else if (
@@ -753,8 +755,8 @@
 }
 
 // Row indices are checked on the host, which needs them to come from a
-// graph input, possibly through integer casts and views; any other source
-// is rejected.
+// graph input through integer casts, views or repeats, or an integer arange
+// whose bounds are resolved when building the graph.
 void GraphBuilder::check_row_indices(ValueId id, int64_t rows) {
   static const std::unordered_set<std::string> kSameValues = {
       "aten._to_copy.default",
@@ -765,6 +767,7 @@
       "aten.expand_copy.default",
       "aten.permute.default",
       "aten.permute_copy.default",
+      "aten.repeat.default",
       "aten.reshape.default",
       "aten._unsafe_view.default",
       "aten.select.int",
@@ -796,6 +799,19 @@
       break;
     }
     const Node& node = method_.graph.node(producer);
+    if (node.is_call() && op_key(node.target) == "aten.arange.start_step") {
+      const int64_t start = int_arg(node, "start", 0, 0);
+      const int64_t step = int_arg(node, "step", 2, 1, true);
+      const int64_t count = ggml_nelements(tensor(id).tensor);
+      int64_t last = 0;
+      if (count > 0 &&
+          (__builtin_mul_overflow(count - 1, step, &last) ||
+           __builtin_add_overflow(start, last, &last) ||
+           std::min(start, last) < 0 || std::max(start, last) >= rows)) {
+        throw std::out_of_range("ggml: arange row index is outside the table");
+      }
+      return;
+    }
     const NamedArgument* source = find_arg(node, "self", 0);
     if (!node.is_call() || !kSameValues.contains(op_key(node.target)) ||
         source == nullptr || source->arg.kind() != ArgKind::Tensor) {
@@ -805,7 +821,7 @@
   }
   throw std::runtime_error(
       "ggml: row indices must come from a graph input, through integer "
-      "casts and views only");
+      "casts, views or repeats, or an integer arange");
 }
 
 ggml_tensor* GraphBuilder::as_f32(ggml_tensor* tensor) {

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

@@ -186,7 +186,7 @@ class GraphBuilder {
   void set_single_output(const Node& node, TensorRef result);
   void lower(const Node& node);
   void lower_linear(const Node& node);
-  void lower_embedding(const Node& node);
+  void lower_row_gather(const Node& node);
   void check_row_indices(ValueId id, int64_t rows);
   void lower_rms_norm(const Node& node);
   void lower_layer_norm(const Node& node);
@@ -244,6 +244,7 @@ class GraphBuilder {
   void lower_unsqueeze(const Node& node);
   void lower_squeeze(const Node& node);
   void lower_expand(const Node& node);
+  void lower_repeat(const Node& node);
   void lower_cat(const Node& node);
   void lower_matmul(const Node& node);
   void lower_update_cache(const Node& node);

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

@@ -47,11 +47,29 @@ void GraphBuilder::lower_linear(const Node& node) {
   set_single_output(node, TensorRef{result, std::move(sizes)});
 }
 
-void GraphBuilder::lower_embedding(const Node& node) {
+void GraphBuilder::lower_row_gather(const Node& node) {
   const TensorRef weight = tensor_arg(node, "weight", 0);
-  TensorRef indices = tensor_arg(node, "indices", 1);
-  check_row_indices(
-      require_arg(node, "indices", 1).arg.as_tensor().id, weight.sizes.at(0));
+  const Argument& arg = require_arg(node, "indices", 1).arg;
+  ValueId indices_id = kInvalid;
+  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");
+    }
+    indices_id = ids.front();
+  }
+  const ScalarType dtype = method_.graph.value(indices_id).tensor_meta().dtype;
+  if (weight.sizes.size() != 2 ||
+      (dtype != ScalarType::Int && dtype != ScalarType::Long)) {
+    throw std::runtime_error(
+        "ggml: row gather needs a matrix and integer indices");
+  }
+  TensorRef indices = tensor(indices_id);
+  check_row_indices(indices_id, weight.sizes[0]);
   if (indices.tensor->type != GGML_TYPE_I32) {
     indices.tensor =
         ggml_cast(instance_.context.get(), indices.tensor, GGML_TYPE_I32);

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

@@ -8,6 +8,7 @@
 
 #include <algorithm>
 #include <array>
+#include <limits>
 #include <numeric>
 #include <optional>
 #include <stdexcept>
@@ -294,6 +295,32 @@ void GraphBuilder::lower_expand(const Node& node) {
   set_single_output(node, TensorRef{result, std::move(sizes)});
 }
 
+void GraphBuilder::lower_repeat(const Node& node) {
+  const TensorRef input = tensor_arg(node, "self", 0);
+  if (input.tensor == nullptr) {
+    throw std::runtime_error("ggml: repeat input has no tensor");
+  }
+  std::vector<int64_t> sizes = int_list(require_arg(node, "repeats", 1).arg);
+  if (sizes.size() < input.sizes.size() || sizes.size() > 4) {
+    throw std::runtime_error(
+        "ggml: repeat needs input rank <= repeats rank <= 4");
+  }
+  const size_t offset = sizes.size() - input.sizes.size();
+  for (size_t i = 0; i < sizes.size(); ++i) {
+    const int64_t extent = i < offset ? 1 : input.sizes[i - offset];
+    if (extent <= 0 || sizes[i] <= 0 ||
+        sizes[i] > std::numeric_limits<int64_t>::max() / extent) {
+      throw std::runtime_error("ggml: repeat requires positive, bounded sizes");
+    }
+    sizes[i] *= extent;
+  }
+  ggml_context* ctx = instance_.context.get();
+  ggml_tensor* shape = new_tensor(ctx, input.tensor->type, sizes);
+  ggml_tensor* result = ggml_repeat(ctx, input.tensor, shape);
+  // A no-op repeat must still snapshot storage that can be written later.
+  set_single_output(node, TensorRef{snapshot(result), std::move(sizes)});
+}
+
 void GraphBuilder::lower_cat(const Node& node) {
   const Argument& tensors_arg = require_arg(node, "tensors", 0).arg;
   if (tensors_arg.kind() != ArgKind::TensorList) {

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

@@ -884,6 +884,56 @@ def forward(self, indices):
 
         self._check(SqueezedEmbedding(), (torch.tensor([[1, 3, 9]]),))
 
+    def test_embedding_arange_indices(self):
+        class PositionEmbedding(Embedding):
+            def forward(self, x):
+                return self.embedding(torch.arange(1, 10, 2)) + x
+
+        self._check(PositionEmbedding(), (torch.randn(5, 8),))
+
+    def test_embedding_arange_indices_rejects_out_of_bounds(self):
+        class InvalidPositionEmbedding(Embedding):
+            def forward(self, x):
+                return self.embedding(torch.arange(-1, 4)) + x
+
+        with self.assertRaisesRegex(AssertionError, "arange row index"):
+            # Export does not execute the invalid gather; the runner must reject it.
+            with tempfile.TemporaryDirectory() as tmp:
+                outdir = Path(tmp)
+                x = torch.randn(5, 8)
+                package = outdir / "model.ptn"
+                ep = torch.export.export(InvalidPositionEmbedding().eval(), (x,))
+                to_native(ep).save(str(package))
+                _run_ggml(package, [(x,)], outdir, 1)
+
+    def test_repeated_row_gather_rechecks_indices(self):
+        class RepeatedRowGather(Embedding):
+            def forward(self, indices):
+                return self.embedding.weight[indices.unsqueeze(0).repeat(2, 1)]
+
+        model = RepeatedRowGather().eval()
+        inputs = (torch.tensor([1, 3, 9]),)
+        self._check(model, inputs)
+        with tempfile.TemporaryDirectory() as tmp:
+            outdir = Path(tmp)
+            package = outdir / "model.ptn"
+            to_native(torch.export.export(model, inputs)).save(str(package))
+            for invalid in (-1, 10):
+                with self.subTest(invalid=invalid):
+                    with self.assertRaisesRegex(AssertionError, "row index is outside"):
+                        _run_ggml(
+                            package,
+                            [inputs, (torch.tensor([1, invalid, 3]),)],
+                            outdir,
+                            1,
+                        )
+
+    def test_repeat_tiles_non_singleton_dimensions(self):
+        self._check(
+            Lambda(lambda x: x.repeat(2, 3, 2)),
+            (torch.arange(6, dtype=torch.float32).reshape(2, 3),),
+        )
+
     def test_rms_norm(self):
         self._check(RMSNorm(), (torch.randn(2, 3, 16),))