diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp index 372176f8d2d..ee6718ebacc 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp @@ -649,8 +649,8 @@ void GraphBuilder::lower(const Node& node) { 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 @@ void GraphBuilder::lower(const Node& node) { } 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,31 +755,21 @@ void GraphBuilder::lower(const Node& node) { } // 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 kSameValues = { - "aten._to_copy.default", - "aten.alias.default", - "aten.clone.default", - "aten.contiguous.default", - "aten.expand.default", - "aten.expand_copy.default", - "aten.permute.default", - "aten.permute_copy.default", - "aten.reshape.default", - "aten._unsafe_view.default", - "aten.select.int", - "aten.select_copy.int", - "aten.slice.Tensor", - "aten.slice_copy.Tensor", - "aten.squeeze.dims", - "aten.squeeze_copy.dims", - "aten.transpose_copy.int", - "aten.unsqueeze.default", - "aten.unsqueeze_copy.default", - "aten.view.default", - "aten.view_copy.default", + "aten._to_copy.default", "aten.alias.default", + "aten.clone.default", "aten.contiguous.default", + "aten.expand.default", "aten.expand_copy.default", + "aten.permute.default", "aten.permute_copy.default", + "aten.repeat.default", "aten.reshape.default", + "aten._unsafe_view.default", "aten.select.int", + "aten.select_copy.int", "aten.slice.Tensor", + "aten.slice_copy.Tensor", "aten.squeeze.dims", + "aten.squeeze_copy.dims", "aten.transpose_copy.int", + "aten.unsqueeze.default", "aten.unsqueeze_copy.default", + "aten.view.default", "aten.view_copy.default", }; const auto is_index = [this](ValueId value) { const ScalarType dtype = method_.graph.value(value).tensor_meta().dtype; @@ -796,6 +788,19 @@ void GraphBuilder::check_row_indices(ValueId id, int64_t rows) { 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 +810,7 @@ void GraphBuilder::check_row_indices(ValueId id, int64_t rows) { } 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) { diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h index ee8db65dfc7..2f69e368fe0 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h +++ b/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); diff --git a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp b/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp index 93281445b80..01625b42a09 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp +++ b/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 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); diff --git a/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp b/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp index 2927f5b5dba..96382fd6e94 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -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 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::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) { diff --git a/backends/native/test/runtime/ggml/test_ggml_ops.py b/backends/native/test/runtime/ggml/test_ggml_ops.py index 54d31653a77..2a7169ffecd 100644 --- a/backends/native/test/runtime/ggml/test_ggml_ops.py +++ b/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),))