Fold constant int64→FP32 casts into CPU delegate

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",
     ],
 )

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

backends/cpu/runtime/providers/reference/ReferenceProvider.cpp

@@ -884,6 +884,46 @@ std::string_view check_node(
   return {};
 }
 
+using Sizes = std::vector<aten::SizesType>;
+
+struct LinearViews {
+  std::array<Sizes, 3> inputs;
+  Sizes output;
+};
+
+// aten::linear.out requires an input of rank >= 2, a 2D weight and a 1D bias.
+std::optional<LinearViews> 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<aten::SizesType>(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<EValue*> 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<EValue*>
+  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<TensorStorage>();
     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<aten::StridesType>(stride);
-      stride *= shape[dim];
+      stride *= sizes[dim];
     }
     storage->tensor = std::make_unique<aten::TensorImpl>(
         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>();
       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();
         }

backends/cpu/test/test_aot.py

@@ -197,6 +197,61 @@
         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_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(

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):
@@ -132,6 +162,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 (