From 79be7131926d727e2b328b0211be9ef7250813c0 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 28 Oct 2026 13:21:10 -0700 Subject: [PATCH 1/2] Support 1D linear in the CPU Reference provider Summary: `aten::linear` with a 1D input or a 1D weight failed at execution with `0x12` (InvalidArgument). The Reference provider routes every `linear` to `opt_linear_out`, which requires an input of rank >= 2 and a 2D weight. Portable has no `linear` kernel, and XNNPACK declines non-constant or non-2D weights. The Reference provider now binds 1D operands as single-row views: a `[K]` input or weight becomes `[1, K]`, and the output gains the matching unit dimension. Scalar bias is also viewed as `[1]` to satisfy the optimized kernel's bias contract. Views share the value's buffer, so nothing is copied. The partitioner keeps preserving these `linear` nodes instead of decomposing them, so a provider can later claim the 1D shapes with a dedicated kernel. Authored with assistance from Claude Code and Codex. Differential Revision: D128153701 --- .../providers/reference/ReferenceProvider.cpp | 72 ++++++++++++++++--- backends/cpu/test/test_aot.py | 17 +++++ backends/cpu/test/test_runtime.py | 19 +++++ 3 files changed, 97 insertions(+), 11 deletions(-) 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 ( From 1ef4175fd71e40b512e2e8b3454609bbf061a6e9 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 28 Oct 2026 13:27:46 -0700 Subject: [PATCH 2/2] =?UTF-8?q?Fold=20constant=20int64=E2=86=92FP32=20cast?= =?UTF-8?q?s=20into=20CPU=20delegate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: The experimental FP32 CPU backend split the graph around hard activations: `hardsigmoid`/`hardswish` decompose into a lifted scalar `int64` constant cast to FP32 via `_to_dim_order_copy`, and `CPUSemanticOperators` rejected the `int64` input, leaving two copies per activation outside the partition. The `to_backend` contract forbids a partitioner from mutating the graph (`is_identical_graph`), so (mirroring the Samsung backend) the fold is split across the two owning files: - `partitioner.py`: a read-only `get_propagated_const_tensor_dict` analysis marks call-function nodes that resolve to an FP32 constant as supported, pulling the casts into the partition instead of forming a boundary. Restricting to FP32 keeps the FP32-only ABI safe and leaves int64 position-index subgraphs out of scope. - `preprocess.py`: reuse `constant_prop_pass` to actually fold those now-in-partition casts to FP32 constants before serialization, so a non-FP32 constant is never serialized. `constant_prop_pass` preserves dtype/shape/value, skips inputs/mutable buffers/impure ops, and already guards `_prop_tensor_constant*` name collisions. Differential Revision: D123138812 --- backends/cpu/BUCK | 1 + backends/cpu/preprocess.py | 23 +++++++++++++++++++ backends/cpu/test/test_aot.py | 38 +++++++++++++++++++++++++++++++ backends/cpu/test/test_runtime.py | 30 ++++++++++++++++++++++++ 4 files changed, 92 insertions(+) diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index 18e9006b51a..294008607a2 100644 --- a/backends/cpu/BUCK +++ b/backends/cpu/BUCK @@ -30,6 +30,7 @@ fbcode_target( "//executorch/exir/backend:compile_spec_schema", "//executorch/exir/backend:partitioner", "//executorch/exir/backend:utils", + "//executorch/exir/passes:constant_prop_pass", "//executorch/export:lib", ], ) diff --git a/backends/cpu/preprocess.py b/backends/cpu/preprocess.py index 42eb88e298b..003ea858e2b 100644 --- a/backends/cpu/preprocess.py +++ b/backends/cpu/preprocess.py @@ -25,6 +25,11 @@ ExportedProgram, PreprocessResult, ) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.constant_prop_pass import ( + _DEFAULT_SKIP_TARGETS, + constant_prop_pass, +) from executorch.exir.program._program import ( _transform, _update_exported_program_graph_module, @@ -87,6 +92,15 @@ def _remove_unused_constants(edge_program: ExportedProgram) -> None: FuseBatchNormWithLinearPass, ] +# Constant folding must not materialize uninitialized allocations: an +# `empty` node's bytes are undefined, so it stays a runtime allocation +# instead of becoming a serialized constant. +_CONSTANT_PROP_SKIP_TARGETS = _DEFAULT_SKIP_TARGETS | { + exir_ops.edge.aten.empty.memory_format, + exir_ops.edge.aten.empty_strided.default, + exir_ops.edge.dim_order_ops._empty_dim_order.default, +} + class CpuBackend(BackendDetails): """Serialize provider-neutral Native semantics with contiguous tensor storage. @@ -118,6 +132,15 @@ def preprocess( edge_program = _transform(edge_program, pass_type(edge_program)) # Passes can introduce tensor attributes that must become program inputs. edge_program = lift_constant_tensor_pass(edge_program) + # Fold pure constant casts the partitioner pulled in (e.g. the int64->FP32 + # scalars from hard-activation decomposition) to FP32 constants before + # serialization. This runs on the program ahead of the graph-module + # rewrites below, which derive the serialization graph from it. + constant_prop_pass( + edge_program, + custom_skip_targets=_CONSTANT_PROP_SKIP_TARGETS, + _skip_dim_order=False, + ) # Fusion leaves the original weights and batch norm stats unused. _remove_unused_constants(edge_program) graph_module = edge_program.graph_module diff --git a/backends/cpu/test/test_aot.py b/backends/cpu/test/test_aot.py index f20a0093753..9434b4e7686 100644 --- a/backends/cpu/test/test_aot.py +++ b/backends/cpu/test/test_aot.py @@ -214,6 +214,44 @@ def forward(self, x, w): ] self.assertEqual(delegate_targets.count("torch.ops.aten.linear.default"), 1) + def test_hard_activation_constant_casts_are_folded_into_delegate(self): + for activation in ( + torch.nn.functional.hardsigmoid, + torch.nn.functional.hardswish, + ): + + class Model(torch.nn.Module): + def forward(self, x): + return activation(x) + + inputs = (torch.randn(2, 3, 4, 5),) + edge = to_edge_transform_and_lower( + torch.export.export(Model(), inputs), + partitioner=[CPUPartitioner()], + ) + exported = edge.exported_program() + # The lifted int64->FP32 casts must be pulled into the delegate, + # leaving no copy boundaries in the top-level graph. + self.assertEqual( + [ + node + for node in exported.graph.nodes + if node.op == "call_function" + and "to_dim_order_copy" in str(node.target) + ], + [], + ) + program = edge.to_executorch().executorch_program + self.assertEqual( + [d.id for d in program.execution_plan[0].delegates], ["CpuBackend"] + ) + graph = deserialize_graph(bytes(program.backend_delegate_data[0].data)) + self.assertNotIn( + "torch.ops.dim_order_ops._to_dim_order_copy.default", + [node.target for node in graph.nodes], + ) + torch.testing.assert_close(exported.module()(*inputs), Model()(*inputs)) + 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 7359fc54b3f..c7d71316857 100644 --- a/backends/cpu/test/test_runtime.py +++ b/backends/cpu/test/test_runtime.py @@ -121,6 +121,36 @@ def forward(self, x, mask): for output, expected in zip(actual, model(x, current), strict=True): torch.testing.assert_close(output, expected, rtol=0, atol=0) + def test_fp64_constant_chain_stays_outside_delegate(self): + class AddConstant(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "constant", torch.arange(6, dtype=torch.float64).reshape(2, 3) + ) + + def forward(self, x): + return x + self.constant.sin().float() + + model = AddConstant() + x = torch.arange(6, dtype=torch.float32).reshape(2, 3) + edge = to_edge_transform_and_lower( + torch.export.export(model, (x,)), partitioner=[CPUPartitioner()] + ) + program = edge.to_executorch() + plan = program.executorch_program.execution_plan[0] + self.assertEqual([delegate.id for delegate in plan.delegates], ["CpuBackend"]) + runtime = _load_for_executorch_from_buffer(program.buffer) + for current in (x, -x): + torch.testing.assert_close(runtime.forward((current,))[0], model(current)) + + def test_hard_activations_execute_with_folded_constants(self): + x = torch.tensor([-4.0, -3.0, -1.0, 0.0, 1.0, 3.0, 4.0]) + for model in (torch.nn.Hardsigmoid(), torch.nn.Hardswish()): + with self.subTest(activation=type(model).__name__): + with self._lower(model, (x,)) as runtime: + torch.testing.assert_close(runtime.forward((x,))[0], model(x)) + def test_empty_tensors_execute(self): class Add(torch.nn.Module): def forward(self, x):