diff --git a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp index 8dc54961c0a..73405391f46 100644 --- a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp +++ b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp @@ -884,6 +884,46 @@ std::string_view check_node( return {}; } +using Sizes = std::vector; + +struct LinearViews { + std::array inputs; + Sizes output; +}; + +// aten::linear.out requires an input of rank >= 2, a 2D weight and a 1D bias. +std::optional linear_views( + const Kernel& node, + const ptn::Graph& graph) { + if (node.target != "torch.ops.aten.linear.default" || + node.inputs.size() < 2 || + node.inputs[0].arg.kind() != ptn::ArgKind::Tensor || + node.inputs[1].arg.kind() != ptn::ArgKind::Tensor) { + return std::nullopt; + } + const auto& input = + graph.value(node.inputs[0].arg.as_tensor().id).tensor_meta().sizes; + const auto& weight = + graph.value(node.inputs[1].arg.as_tensor().id).tensor_meta().sizes; + if (input.empty() || (input.size() >= 2 && weight.size() == 2)) { + return std::nullopt; + } + const auto row = [](const auto& sizes) { + return sizes.size() == 1 ? Sizes{1, static_cast(sizes[0])} + : Sizes(sizes.begin(), sizes.end()); + }; + LinearViews views{{row(input), row(weight), {}}, {}}; + if (node.inputs.size() > 2 && + node.inputs[2].arg.kind() == ptn::ArgKind::Tensor) { + const auto& bias = + graph.value(node.inputs[2].arg.as_tensor().id).tensor_meta().sizes; + views.inputs[2] = bias.empty() ? Sizes{1} : Sizes(bias.begin(), bias.end()); + } + views.output.assign(views.inputs[0].begin(), views.inputs[0].end() - 1); + views.output.push_back(views.inputs[1][0]); + return views; +} + struct Instruction { const ptn::Node* node = nullptr; OpFunction kernel = nullptr; @@ -908,9 +948,16 @@ class ReferenceExecutable final : public Executable { }; Result tensor_value(ptn::ValueId id, bool writable) { + const auto& shape = context_.graph.value(id).tensor_meta().sizes; + return tensor_value(id, writable, Sizes(shape.begin(), shape.end())); + } + + // Views of one value share its buffer but keep separate tensor metadata. + Result + tensor_value(ptn::ValueId id, bool writable, const Sizes& sizes) { const auto existing = std::find_if( - tensors_.begin(), tensors_.end(), [id](const auto& storage) { - return storage->id == id; + tensors_.begin(), tensors_.end(), [id, &sizes](const auto& storage) { + return storage->id == id && storage->sizes == sizes; }); if (existing != tensors_.end()) { (*existing)->writable |= writable; @@ -921,23 +968,22 @@ class ReferenceExecutable final : public Executable { return bytes.error(); } const auto& meta = context_.graph.value(id).tensor_meta(); - const auto& shape = meta.sizes; auto storage = std::make_unique(); storage->id = id; storage->writable = writable; - storage->sizes.assign(shape.begin(), shape.end()); - storage->order.resize(shape.size()); - storage->strides.resize(shape.size()); + storage->sizes = sizes; + storage->order.resize(sizes.size()); + storage->strides.resize(sizes.size()); int64_t stride = 1; - for (size_t reverse = shape.size(); reverse > 0; --reverse) { + for (size_t reverse = sizes.size(); reverse > 0; --reverse) { const auto dim = reverse - 1; storage->order[dim] = dim; storage->strides[dim] = static_cast(stride); - stride *= shape[dim]; + stride *= sizes[dim]; } storage->tensor = std::make_unique( to_aten(meta.dtype), - shape.size(), + sizes.size(), storage->sizes.data(), nullptr, storage->order.data(), @@ -1036,6 +1082,7 @@ class ReferenceExecutable final : public Executable { "Empty negative view copy requires matching input/output metadata"); continue; } + const auto views = linear_views(node, context_.graph); auto instruction = std::make_unique(); instruction->node = &node; instruction->kernel = route->kernel; @@ -1059,7 +1106,9 @@ class ReferenceExecutable final : public Executable { const auto& input = node.inputs[index]; if (input.arg.kind() == ptn::ArgKind::Tensor) { const auto value_id = input.arg.as_tensor().id; - auto value = tensor_value(value_id, false); + auto value = views && index < views->inputs.size() + ? tensor_value(value_id, false, views->inputs[index]) + : tensor_value(value_id, false); if (!value.ok()) { return value.error(); } @@ -1114,7 +1163,8 @@ class ReferenceExecutable final : public Executable { InvalidProgram, "Reference provider requires tensor outputs: %s", node.target.c_str()); - auto value = tensor_value(output.value_id, true); + auto value = views ? tensor_value(output.value_id, true, views->output) + : tensor_value(output.value_id, true); if (!value.ok()) { return value.error(); } diff --git a/backends/cpu/test/test_aot.py b/backends/cpu/test/test_aot.py index cc8e8c70ec0..f20a0093753 100644 --- a/backends/cpu/test/test_aot.py +++ b/backends/cpu/test/test_aot.py @@ -197,6 +197,23 @@ def test_non_fp32_linear_decomposes(self): self.assertEqual(program.executorch_program.execution_plan[0].delegates, []) self.assertGreater(len(program.buffer), 0) + def test_1d_linear_is_preserved_and_delegated(self): + class Linear1D(torch.nn.Module): + def forward(self, x, w): + return torch.ops.aten.linear.default(x, w, None) + + edge = to_edge_transform_and_lower( + torch.export.export(Linear1D(), (torch.randn(3), torch.randn(3))), + partitioner=[CPUPartitioner()], + ) + program = edge.to_executorch() + delegate_targets = [ + node.target + for data in program.executorch_program.backend_delegate_data + for node in deserialize_graph(bytes(data.data)).nodes + ] + self.assertEqual(delegate_targets.count("torch.ops.aten.linear.default"), 1) + def test_constants_have_recorded_readable_tail(self): model = torch.nn.Linear(2, 3).eval() result = CpuBackend.preprocess( diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py index 7195c9a4a14..7359fc54b3f 100644 --- a/backends/cpu/test/test_runtime.py +++ b/backends/cpu/test/test_runtime.py @@ -132,6 +132,25 @@ def forward(self, x): actual = runtime.forward((x,))[0] torch.testing.assert_close(actual, model(x)) + def test_1d_linear_reference_fallback(self): + class Linear(torch.nn.Module): + def forward(self, *args): + return torch.nn.functional.linear(*args) + + model = Linear() + for inputs in ( + (torch.randn(4), torch.randn(3, 4), torch.randn(3)), + (torch.randn(4), torch.randn(4)), + (torch.randn(2, 4), torch.randn(4)), + (torch.randn(4), torch.randn(4), torch.randn(())), + (torch.randn(2, 3, 4), torch.randn(4), torch.randn(())), + ): + with self.subTest(shapes=[x.shape for x in inputs]): + with self._lower(model, inputs) as runtime: + torch.testing.assert_close( + runtime.forward(inputs)[0], model(*inputs) + ) + def test_maxpool2d_reference_fallback(self): model = torch.nn.MaxPool2d(2) for x in (