Route _neg_view_copy to ET neg kernel (#24194)

backends/cpu/BUCK

@@ -22,6 +22,8 @@ fbcode_target(
     deps = [
         "//caffe2:torch",
         "//executorch/backends/native:lib",
+        "//executorch/backends/transforms:fuse_batch_norm_with_conv",
+        "//executorch/backends/transforms:fuse_batch_norm_with_linear",
         "//executorch/exir:lib",
         "//executorch/exir/_serialize:lib",
         "//executorch/exir/backend:backend_details",

backends/cpu/preprocess.py

@@ -12,15 +12,33 @@
     CPU_SUPPORTED_DTYPES,
 )
 from executorch.backends.native.serialization import serialize_graph
+from executorch.backends.transforms.fuse_batch_norm_with_conv import (
+    FuseBatchNormWithConvPass,
+)
+from executorch.backends.transforms.fuse_batch_norm_with_linear import (
+    FuseBatchNormWithLinearPass,
+)
 from executorch.exir._serialize._named_data_store import NamedDataStore
 from executorch.exir.backend.backend_details import (
     BackendDetails,
     CompileSpec,
     ExportedProgram,
     PreprocessResult,
 )
+from executorch.exir.program._program import (
+    _transform,
+    _update_exported_program_graph_module,
+    lift_constant_tensor_pass,
+)
+from torch.export.graph_signature import InputKind
 from torch.utils._pytree import tree_leaves, tree_map
 
+_CONSTANT_INPUT_KINDS = (
+    InputKind.PARAMETER,
+    InputKind.BUFFER,
+    InputKind.CONSTANT_TENSOR,
+)
+
 
 def _validate_boundary_layouts(graph_module: torch.fx.GraphModule) -> None:
     for node in graph_module.graph.nodes:
@@ -37,6 +55,39 @@ def _validate_boundary_layouts(graph_module: torch.fx.GraphModule) -> None:
             )
 
 
+def _remove_unused_constants(edge_program: ExportedProgram) -> None:
+    signature = edge_program.graph_signature
+    mutated = set(signature.buffers_to_mutate.values())
+    placeholders = {
+        node.name: node for node in edge_program.graph.nodes if node.op == "placeholder"
+    }
+    kept_specs = []
+    for spec in signature.input_specs:
+        node = placeholders.get(spec.arg.name)
+        if (
+            node is not None
+            and spec.kind in _CONSTANT_INPUT_KINDS
+            and spec.target not in mutated
+            and not node.users
+        ):
+            edge_program.graph.erase_node(node)
+            edge_program.state_dict.pop(spec.target, None)
+            edge_program.constants.pop(spec.target, None)
+        else:
+            kept_specs.append(spec)
+    signature.input_specs = kept_specs
+    edge_program.graph_module.recompile()
+
+
+# Ordered passes applied to every program in preprocess, mirroring
+# XNNPACKPassManager.passes. Each entry takes the current program, so later
+# passes observe earlier passes' rewrites; extend by appending new pass types.
+CPU_PREPROCESS_PASSES = [
+    FuseBatchNormWithConvPass,
+    FuseBatchNormWithLinearPass,
+]
+
+
 class CpuBackend(BackendDetails):
     """Serialize provider-neutral Native semantics with contiguous tensor storage.
 
@@ -59,7 +110,17 @@ def preprocess(
                 f"expected exactly {expected_specs!r}, got {specs!r}"
             )
         _validate_boundary_layouts(edge_program.graph_module)
-        graph_module = deepcopy(edge_program.graph_module)
+        edge_program = deepcopy(edge_program)
+        edge_program = _update_exported_program_graph_module(
+            edge_program, edge_program.graph_module
+        )
+        for pass_type in CPU_PREPROCESS_PASSES:
+            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)
+        # Fusion leaves the original weights and batch norm stats unused.
+        _remove_unused_constants(edge_program)
+        graph_module = edge_program.graph_module
         for node in graph_module.graph.nodes:
             if "val" in node.meta:
                 node.meta["val"] = tree_map(

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

@@ -106,6 +106,7 @@ constexpr auto kKernels = std::to_array<KernelSpec>({
      nullptr,
      false,
      "56"},
+    {"torch.ops.aten._neg_view_copy.default", "aten::neg.out", 1, 1, nullptr},
     {"torch.ops.aten._pdist_forward.default",
      "aten::_pdist_forward.out",
      2,
@@ -1021,6 +1022,20 @@ class ReferenceExecutable final : public Executable {
             reason.data());
         return Error::NotSupported;
       }
+      if (node.target == "torch.ops.aten._neg_view_copy.default" &&
+          tensor_bytes(context_.graph.value(node.inputs[0].arg.as_tensor().id))
+                  .get() == 0) {
+        const auto& input =
+            context_.graph.value(node.inputs[0].arg.as_tensor().id)
+                .tensor_meta();
+        const auto& output =
+            context_.graph.value(node.outputs[0].value_id).tensor_meta();
+        ET_CHECK_OR_RETURN_ERROR(
+            input.dtype == output.dtype && input.sizes == output.sizes,
+            InvalidProgram,
+            "Empty negative view copy requires matching input/output metadata");
+        continue;
+      }
       auto instruction = std::make_unique<Instruction>();
       instruction->node = &node;
       instruction->kernel = route->kernel;

backends/cpu/test/test_aot.py

@@ -14,12 +14,16 @@
     CPUPartitioner,
     CPUSemanticOperators,
 )
-from executorch.backends.cpu.preprocess import CpuBackend
+from executorch.backends.cpu.preprocess import _remove_unused_constants, CpuBackend
 from executorch.backends.cpu.recipes import CPURecipeType
 from executorch.backends.native.serialization import deserialize_graph, serialize_graph
-from executorch.exir import to_edge_transform_and_lower
+from executorch.backends.transforms.fuse_batch_norm_with_linear import (
+    FuseBatchNormWithLinearPass,
+)
+from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.program._program import _transform, lift_constant_tensor_pass
 from executorch.exir.schema import DelegateCall, Tensor
 from executorch.export import export, ExportRecipe
 
@@ -213,6 +217,64 @@ def test_constants_have_recorded_readable_tail(self):
             self.assertEqual(len(buffer), tensor.numel() * tensor.element_size() + 64)
             self.assertEqual(buffer[-64:], bytes(64))
 
+    def test_preprocess_fuses_batch_norm_into_conv_and_linear(self):
+        class ConvBnLinearBn(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(4, 4, 3, padding=1)
+                self.bn1 = torch.nn.BatchNorm2d(4)
+                self.linear = torch.nn.Linear(4, 4)
+                self.bn2 = torch.nn.BatchNorm1d(4)
+
+            def forward(self, x):
+                y = self.bn1(self.conv(x))
+                return self.bn2(self.linear(y.mean(dim=(2, 3))))
+
+        model = ConvBnLinearBn().eval()
+        edge = to_edge(
+            torch.export.export(model, (torch.randn(1, 4, 8, 8),)),
+            compile_config=EdgeCompileConfig(
+                preserve_ops=[torch.ops.aten.linear.default]
+            ),
+        ).exported_program()
+        self.assertTrue(
+            any("batch_norm" in str(node.target) for node in edge.graph.nodes)
+        )
+        result = CpuBackend.preprocess(
+            edge,
+            [
+                CompileSpec(
+                    "cpu_delegate_version",
+                    CPU_DELEGATE_VERSION.to_bytes(4, "little"),
+                )
+            ],
+        )
+        graph = deserialize_graph(bytes(result.processed_bytes))
+        targets = [node.target for node in graph.nodes]
+        self.assertNotIn(
+            "torch.ops.aten._native_batch_norm_legit_no_training.default", targets
+        )
+        self.assertIn("torch.ops.aten.convolution.default", targets)
+        self.assertIn("torch.ops.aten.linear.default", targets)
+        # Only the fused conv and linear weights and biases remain.
+        self.assertEqual(len(result.data_store_output.pte_data), 4)
+
+    def test_remove_unused_constants_prunes_backing_store(self):
+        model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.BatchNorm1d(4))
+        ep = to_edge(
+            torch.export.export(model.eval(), (torch.randn(2, 4),)),
+            compile_config=EdgeCompileConfig(
+                preserve_ops=[torch.ops.aten.linear.default]
+            ),
+        ).exported_program()
+        ep = _transform(ep, FuseBatchNormWithLinearPass(ep))
+        ep = lift_constant_tensor_pass(ep)
+        self.assertIn("0.weight", ep.state_dict)
+        _remove_unused_constants(ep)
+        targets = {spec.target for spec in ep.graph_signature.input_specs}
+        self.assertNotIn("0.weight", ep.state_dict)
+        self.assertLessEqual(set(ep.state_dict) | set(ep.constants), targets)
+
     def test_preprocess_rejects_compile_spec_mismatch(self):
         ep = torch.export.export(torch.nn.Linear(2, 3), (torch.ones(1, 2),))
         for specs in (

backends/cpu/test/test_runtime.py

@@ -95,6 +95,29 @@ 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_neg_view_copy_routes_to_neg(self):
+        class NegView(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten._neg_view_copy.default(x)
+
+        model = NegView()
+        x = torch.randn(2, 3)
+        with self._lower(model, (x,)) as runtime:
+            actual = runtime.forward((x,))[0]
+            torch.testing.assert_close(actual, model(x))
+
+    def test_neg_view_copy_empty_bool(self):
+        class NegView(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten._neg_view_copy.default(x)
+
+        model = NegView()
+        x = torch.empty(0, dtype=torch.bool)
+        expected = model(x)
+        with self._lower(model, (x,)) as runtime:
+            actual = runtime.forward((x,))[0]
+            torch.testing.assert_close(actual, expected)
+
     def test_optional_integer_index_lists(self):
         class Index(torch.nn.Module):
             def __init__(self, accumulate):
@@ -116,3 +139,29 @@ def forward(self, x, indices, values):
                     actual = runtime.forward(inputs)
                     for output, expected in zip(actual, model(*inputs), strict=True):
                         torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+    def test_batch_norm_fused_into_conv_and_linear(self):
+        class ConvBnLinearBn(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(4, 4, 3, padding=1)
+                self.bn1 = torch.nn.BatchNorm2d(4)
+                self.linear = torch.nn.Linear(4, 8)
+                self.bn2 = torch.nn.BatchNorm1d(8)
+
+            def forward(self, x):
+                y = self.bn1(self.conv(x))
+                return self.bn2(self.linear(y.mean(dim=(2, 3))))
+
+        torch.manual_seed(0)
+        model = ConvBnLinearBn().eval()
+        for bn in (model.bn1, model.bn2):
+            bn.running_mean.uniform_(-1, 1)
+            bn.running_var.uniform_(0.5, 2)
+            bn.weight.data.uniform_(0.5, 2)
+            bn.bias.data.uniform_(-1, 1)
+        inputs = (torch.randn(2, 4, 8, 8),)
+        with self._lower(model, inputs) as runtime:
+            torch.testing.assert_close(
+                runtime.forward(inputs)[0], model(*inputs), rtol=1e-5, atol=1e-5
+            )

backends/transforms/targets.bzl

@@ -148,6 +148,20 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_library(
+        name = "fuse_batch_norm_with_linear",
+        srcs = ["fuse_batch_norm_with_linear.py"],
+        visibility = [
+            "//executorch/backends/...",
+        ],
+        deps = [
+            ":utils",
+            "//caffe2:torch",
+            "//executorch/exir:pass_base",
+            "//executorch/exir/dialects:lib",
+        ],
+    )
+
     runtime.python_library(
         name = "fuse_conv_with_clamp",
         srcs = ["fuse_conv_with_clamp.py"],