From a120cbe25524b25b58645da0aef24d2ed18526be Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 27 Oct 2026 09:08:35 -0700 Subject: [PATCH 1/4] Support integer and boolean CPU tensors (#24191) Summary: CPU partitioning and delegate I/O required every tensor to be FP32, leaving indices, masks and integer outputs outside the delegate. Accept static uint8, int8, int16, int32, int64 and bool tensors, and validate runtime I/O against the serialized dtype. Existing allocation and ET kernel routes preserve the values without conversion through FP32. Integer delegation also exposes `index.Tensor`/`index_put`, whose indices are `Tensor?[]`. Native serialization emits a plain tensor list when every index is present, but ET's kernels unbox a list of optional tensors. Mark those Reference routes with `optional_tensor_lists` so the CPU runtime binds optional lists for them; the shared serializer and the other Native engines are unchanged. Add exact runtime regressions for integer precision, indices, mixed outputs, constants and optional index lists. Dynamic shapes/scalars and quantized execution remain separate. The Native schema and storage contract do not change. Authored with AI assistance using Codex and Claude Code. Reviewed By: digantdesai Differential Revision: D123332329 --- backends/cpu/BUCK | 13 ++ backends/cpu/README.md | 9 +- backends/cpu/partitioner.py | 16 ++- backends/cpu/preprocess.py | 11 +- backends/cpu/runtime/CPUBackend.cpp | 5 +- backends/cpu/runtime/KernelProvider.cpp | 21 +++- backends/cpu/runtime/KernelProvider.h | 2 +- .../providers/reference/ReferenceProvider.cpp | 34 ++++- backends/cpu/test/test_aot.py | 6 +- backends/cpu/test/test_runtime.py | 118 ++++++++++++++++++ 10 files changed, 208 insertions(+), 27 deletions(-) create mode 100644 backends/cpu/test/test_runtime.py diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index 93afebb6fa2..0171cb67971 100644 --- a/backends/cpu/BUCK +++ b/backends/cpu/BUCK @@ -57,3 +57,16 @@ fbcode_target( "//executorch/extension/pybindings:portable_lib", # @manual ], ) + +fbcode_target( + _kind = runtime.python_test, + name = "test_runtime", + srcs = ["test/test_runtime.py"], + preload_deps = [":cpu_backend"], + deps = [ + ":lib", + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", + ], +) diff --git a/backends/cpu/README.md b/backends/cpu/README.md index 0c105e8cdf0..51c58edebeb 100644 --- a/backends/cpu/README.md +++ b/backends/cpu/README.md @@ -6,8 +6,11 @@ The initial workload is pretrained ConvNeXt-Tiny in FP32 on a Linux server. The registered `CPURecipeType.FP32` (`cpu_fp32`) Recipe retains a semantic Native graph inside one `CpuBackend` delegate in a regular ExecuTorch `.pte`. -This first delegate version supports pure operations on statically shaped, -contiguous FP32 tensors. +The partitioner supports pure operations on statically shaped, contiguous FP32, +uint8, int8, int16, int32, int64 and boolean tensors. Integer and boolean operations +use the existing ET kernel routes, subject to each kernel's operator and dtype +support. Runtime-computed scalar arguments and symbolic tensor shapes remain +unsupported. Ordinary integer tensors do not imply quantized execution support. The `cpu_delegate_version` compile spec stores `CPU_DELEGATE_VERSION` as a four-byte little-endian unsigned integer. The loader accepts explicitly @@ -44,7 +47,7 @@ allocation. Sequential steps share the maximum declared scratch capacity. XNNPAC shares compatible convolution-weight conversions within a plan and borrows weights whose physical order already matches. -For a loaded method, shapes, FP32 dtype, contiguous layout, threadpool identity, +For a loaded method, shapes, tensor dtypes, contiguous layout, threadpool identity, thread count and ISA capabilities remain fixed. Input/output pointers may change and are validated on every invocation. Changing an execution property requires reloading the method; failed validation executes no provider. The trace includes diff --git a/backends/cpu/partitioner.py b/backends/cpu/partitioner.py index 3aaf77a4324..808150838ca 100644 --- a/backends/cpu/partitioner.py +++ b/backends/cpu/partitioner.py @@ -23,11 +23,19 @@ # Serialized as a little-endian uint32. Bump for incompatible artifact contracts. CPU_DELEGATE_VERSION = 1 +CPU_SUPPORTED_DTYPES = ( + torch.float32, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.bool, +) class CPUSemanticOperators(OperatorSupportBase): - """Pure tensor operations representable by the initial static FP32 - delegate version. + """Pure static tensor operations on FP32, integer and boolean values. Runtime providers determine implementation coverage. An unresolved operation is reported by name when the CPU runtime prepares the graph. @@ -64,7 +72,7 @@ def is_node_supported( if any(related_node.meta.get("val") is None for related_node in related_nodes): return False return all( - value.dtype == torch.float32 + value.dtype in CPU_SUPPORTED_DTYPES and all(isinstance(size, int) for size in value.shape) and value.is_contiguous() for related_node in related_nodes @@ -74,7 +82,7 @@ def is_node_supported( class CPUPartitioner(Partitioner): - """Retain the static FP32 semantic graph for runtime provider selection.""" + """Retain the static semantic graph for runtime provider selection.""" def ops_to_not_decompose(self, ep: torch.export.ExportedProgram): support = CPUSemanticOperators() diff --git a/backends/cpu/preprocess.py b/backends/cpu/preprocess.py index b0a7f9c9141..16bc83e19ac 100644 --- a/backends/cpu/preprocess.py +++ b/backends/cpu/preprocess.py @@ -7,7 +7,10 @@ from copy import deepcopy import torch -from executorch.backends.cpu.partitioner import CPU_DELEGATE_VERSION +from executorch.backends.cpu.partitioner import ( + CPU_DELEGATE_VERSION, + CPU_SUPPORTED_DTYPES, +) from executorch.backends.native.serialization import serialize_graph from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import ( @@ -35,7 +38,7 @@ def _validate_boundary_layouts(graph_module: torch.fx.GraphModule) -> None: class CpuBackend(BackendDetails): - """Serialize provider-neutral Native semantics with contiguous FP32 storage. + """Serialize provider-neutral Native semantics with contiguous tensor storage. The delegate version fixes 64-byte alignment and 64 initialized readable tail bytes. Named constants include the tail in their recorded buffer capacity. @@ -73,9 +76,9 @@ def preprocess( ) data = NamedDataStore() for name, tensor in constants.items(): - if tensor.dtype != torch.float32 or tensor.device.type != "cpu": + if tensor.dtype not in CPU_SUPPORTED_DTYPES or tensor.device.type != "cpu": raise ValueError( - f"CpuBackend requires FP32 CPU constant {name}, " + f"CpuBackend requires supported CPU constant {name}, " f"got dtype {tensor.dtype} on {tensor.device}" ) if not tensor.is_contiguous(): diff --git a/backends/cpu/runtime/CPUBackend.cpp b/backends/cpu/runtime/CPUBackend.cpp index a3c39b5e879..963bf3e7f13 100644 --- a/backends/cpu/runtime/CPUBackend.cpp +++ b/backends/cpu/runtime/CPUBackend.cpp @@ -67,7 +67,7 @@ class CPUExecution { ET_CHECK_OR_RETURN_ERROR( spec.kind == ptn::OutputKind::UserOutput, NotSupported, - "CpuBackend static FP32 delegate does not support mutation"); + "CpuBackend static delegate does not support mutation"); } buffers_.resize(graph.values.size()); bytes_.reserve(graph.values.size()); @@ -171,7 +171,8 @@ class CPUExecution { const auto& tensor = args[index]->toTensor(); const auto& expected = method_->graph.value(id).tensor_meta(); ET_CHECK_OR_RETURN_ERROR( - tensor.scalar_type() == aten::ScalarType::Float && + tensor.scalar_type() == + static_cast(expected.dtype) && tensor.dim() == expected.sizes.size() && std::equal( expected.sizes.begin(), diff --git a/backends/cpu/runtime/KernelProvider.cpp b/backends/cpu/runtime/KernelProvider.cpp index eba874f4f31..5e032c4c47b 100644 --- a/backends/cpu/runtime/KernelProvider.cpp +++ b/backends/cpu/runtime/KernelProvider.cpp @@ -5,6 +5,7 @@ // LICENSE file in the root directory of this source tree. #include +#include #include #include @@ -90,10 +91,16 @@ Result tensor_bytes(const ptn::Value& value) { value.name.c_str()); const auto& meta = value.tensor_meta(); ET_CHECK_OR_RETURN_ERROR( - meta.dtype == ptn::ScalarType::Float && meta.is_contiguous() && - meta.sizes.size() <= 16, + (meta.dtype == ptn::ScalarType::Float || + meta.dtype == ptn::ScalarType::Byte || + meta.dtype == ptn::ScalarType::Char || + meta.dtype == ptn::ScalarType::Short || + meta.dtype == ptn::ScalarType::Int || + meta.dtype == ptn::ScalarType::Long || + meta.dtype == ptn::ScalarType::Bool) && + meta.is_contiguous() && meta.sizes.size() <= 16, NotSupported, - "CPU static delegate requires contiguous FP32 rank <= 16: %s", + "CPU static delegate requires supported dtype and contiguous rank <= 16: %s", value.name.c_str()); size_t elements = 1; for (auto dimension = meta.sizes.rbegin(); dimension != meta.sizes.rend(); @@ -107,7 +114,13 @@ Result tensor_bytes(const ptn::Value& value) { value.name.c_str()); elements *= size; } - return elements * sizeof(float); + const size_t element_size = ptn::element_size(meta.dtype); + ET_CHECK_OR_RETURN_ERROR( + elements <= SIZE_MAX / element_size, + NotSupported, + "CPU tensor byte size overflows: %s", + value.name.c_str()); + return elements * element_size; } Result KernelImplementation::requirements( diff --git a/backends/cpu/runtime/KernelProvider.h b/backends/cpu/runtime/KernelProvider.h index a2571d25d52..0db278cd7e2 100644 --- a/backends/cpu/runtime/KernelProvider.h +++ b/backends/cpu/runtime/KernelProvider.h @@ -131,7 +131,7 @@ class KernelImplementation { return false; } /// Pure metadata query, before compile can read constants or allocate - /// storage. The default declares contiguous FP32 storage for every region + /// storage. The default declares contiguous tensor storage for every region /// boundary. virtual runtime::Result requirements( const KernelRegion& region, diff --git a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp index 5110d7eb108..8e5267b5d75 100644 --- a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp +++ b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp @@ -29,6 +29,9 @@ struct KernelSpec { bool tensor_list_output = false; // Schema positions declared float or float?, rather than Scalar. std::string_view float_args{}; + // Tensor-list inputs are declared Tensor?[]. The serializer emits a plain + // tensor list when every entry is present, but the kernel unboxes optionals. + bool optional_tensor_lists = false; }; aten::ScalarType to_aten(ptn::ScalarType dtype) { @@ -296,8 +299,22 @@ constexpr auto kKernels = std::to_array({ {"torch.ops.aten.gt.Scalar", "aten::gt.Scalar_out", 2, 1, nullptr}, {"torch.ops.aten.gt.Tensor", "aten::gt.Tensor_out", 2, 1, nullptr}, {"torch.ops.aten.hardtanh.default", "aten::hardtanh.out", 3, 1, nullptr}, - {"torch.ops.aten.index.Tensor", "aten::index.Tensor_out", 2, 1, nullptr}, - {"torch.ops.aten.index_put.default", "aten::index_put.out", 4, 1, nullptr}, + {"torch.ops.aten.index.Tensor", + "aten::index.Tensor_out", + 2, + 1, + nullptr, + false, + {}, + true}, + {"torch.ops.aten.index_put.default", + "aten::index_put.out", + 4, + 1, + nullptr, + false, + {}, + true}, {"torch.ops.aten.index_select.default", "aten::index_select.out", 3, @@ -1037,11 +1054,16 @@ class ReferenceExecutable final : public Executable { if (input.arg.kind() == ptn::ArgKind::TensorList || input.arg.kind() == ptn::ArgKind::OptionalTensorList) { auto argument = std::make_unique(); - const bool optional = + const bool serialized_optional = input.arg.kind() == ptn::ArgKind::OptionalTensorList; - const auto& ids = optional ? input.arg.as_optional_tensor_list().ids - : input.arg.as_tensor_list().ids; - const auto error = bind_tensor_list(ids, optional, false, *argument); + const auto& ids = serialized_optional + ? input.arg.as_optional_tensor_list().ids + : input.arg.as_tensor_list().ids; + const auto error = bind_tensor_list( + ids, + serialized_optional || spec->optional_tensor_lists, + false, + *argument); if (error != Error::Ok) { return error; } diff --git a/backends/cpu/test/test_aot.py b/backends/cpu/test/test_aot.py index 61442987b50..e3d75d3d7f9 100644 --- a/backends/cpu/test/test_aot.py +++ b/backends/cpu/test/test_aot.py @@ -274,7 +274,7 @@ def test_preprocess_rejects_unsupported_constants(self): ) with self.assertRaisesRegex( ValueError, - f"CpuBackend requires FP32 CPU constant weight, " + f"CpuBackend requires supported CPU constant weight, " f"got dtype {dtype} on {device}", ): CpuBackend.preprocess( @@ -404,7 +404,7 @@ def test_getitem_validates_output_metadata(self): node.meta["val"] = producer.meta["val"][0] support = CPUSemanticOperators() self.assertTrue(support.is_node_supported({}, node)) - for val in (None, 1, (), torch.ones(2, dtype=torch.int32)): + for val in (None, 1, (), torch.ones(2, dtype=torch.float64)): with self.subTest(val=val): node.meta["val"] = val self.assertFalse(support.is_node_supported({}, node)) @@ -414,7 +414,7 @@ class Add(torch.nn.Module): def forward(self, x): return x + x - ep = torch.export.export(Add(), (torch.ones(2, dtype=torch.int32),)) + ep = torch.export.export(Add(), (torch.ones(2, dtype=torch.float64),)) edge = to_edge_transform_and_lower(ep, partitioner=[CPUPartitioner()]) self.assertEqual( len(edge.to_executorch().executorch_program.execution_plan[0].delegates), 0 diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py new file mode 100644 index 00000000000..d4b92a3f1d2 --- /dev/null +++ b/backends/cpu/test/test_runtime.py @@ -0,0 +1,118 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest +from contextlib import contextmanager + +import torch +from executorch.backends.cpu.partitioner import CPUPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) + + +class CPURuntimeTest(unittest.TestCase): + @contextmanager + def _lower(self, model, inputs): + edge = to_edge_transform_and_lower( + torch.export.export(model, inputs), partitioner=[CPUPartitioner()] + ) + program = edge.to_executorch() + plan = program.executorch_program.execution_plan[0] + self.assertEqual([delegate.id for delegate in plan.delegates], ["CpuBackend"]) + self.assertEqual(plan.operators, []) + # The runtime borrows the serialized program buffer. + yield _load_for_executorch_from_buffer(program.buffer) + + def test_integer_inputs_and_outputs_preserve_precision(self): + class Add(torch.nn.Module): + def forward(self, x, y): + return x + y + + for dtype, values in ( + (torch.uint8, [0, 128, 250]), + (torch.int8, [-100, 0, 100]), + (torch.int16, [-30000, 0, 30000]), + (torch.int32, [-(2**24 + 1), 0, 2**24 + 1]), + (torch.int64, [-(2**40 + 3), 0, 2**40 + 3]), + ): + with self.subTest(dtype=dtype): + model = Add() + x = torch.tensor(values, dtype=dtype) + y = torch.ones_like(x) + with self._lower(model, (x, y)) as runtime: + for inputs in ((x, y), (y, x)): + actual = runtime.forward(inputs)[0] + torch.testing.assert_close( + actual, model(*inputs), rtol=0, atol=0 + ) + + def test_integer_indices_and_mixed_outputs(self): + class Select(torch.nn.Module): + def forward(self, x, indices): + selected = torch.index_select(x, 1, indices) + positive = selected > 0 + return ( + torch.where(positive, selected, -selected), + selected.argmax(dim=1), + positive, + ) + + model = Select() + x = torch.tensor([[-3.0, 5.0, 2.0], [7.0, -2.0, -4.0]]) + for dtype in (torch.int32, torch.int64): + with self.subTest(dtype=dtype): + indices = torch.tensor([2, 0], dtype=dtype) + with self._lower(model, (x, indices)) as runtime: + for current in (indices, indices.flip(0)): + actual = runtime.forward((x, current)) + for output, expected in zip( + actual, model(x, current), strict=True + ): + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_boolean_inputs_and_integer_constants(self): + class MaskedAdd(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.tensor([2**40 + 3, -7])) + + def forward(self, x, mask): + return torch.where(mask, x + self.offset, x), torch.logical_not(mask) + + model = MaskedAdd() + x = torch.tensor([2**40 + 1, -(2**40 + 5)]) + mask = torch.tensor([True, False]) + with self._lower(model, (x, mask)) as runtime: + for current in (mask, ~mask): + actual = runtime.forward((x, current)) + for output, expected in zip(actual, model(x, current), strict=True): + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_optional_integer_index_lists(self): + class Index(torch.nn.Module): + def __init__(self, accumulate): + super().__init__() + self.accumulate = accumulate + + def forward(self, x, indices, values): + updated = torch.index_put(x, (indices,), values, self.accumulate) + return updated, updated[indices], updated[:, indices] + + x = torch.arange(9, dtype=torch.float32).reshape(3, 3) + indices = torch.tensor([0, 2]) + values = torch.tensor([[10.0, 20.0, 30.0], [-1.0, -2.0, -3.0]]) + inputs = (x, indices, values) + for accumulate in (False, True): + with self.subTest(accumulate=accumulate): + model = Index(accumulate) + with self._lower(model, inputs) as runtime: + actual = runtime.forward(inputs) + for output, expected in zip(actual, model(*inputs), strict=True): + torch.testing.assert_close(output, expected, rtol=0, atol=0) From 60eb831d340c708e4fec0fa66fcda3f3a0918af0 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 27 Oct 2026 09:08:35 -0700 Subject: [PATCH 2/4] Add shared linear-batch-norm fusion pass (#24192) Summary: Adds `FuseBatchNormWithLinearPass` to `backends/transforms`, the linear counterpart of `FuseBatchNormWithConvPass`. It folds an inference batch norm into the preceding `aten.linear` weight and bias using `torch.nn.utils.fusion.fuse_linear_bn_weights`, so backends can drop the BN without serializing extra weights. Fusion only happens when it is numerically equivalent: the linear output must be rank 2 (the only case where BN's channel dim is linear's output dim), every weight/bias/running-stat input must be a constant parameter, BN must only be consumed through `getitem(0)`, and training-mode `native_batch_norm` is skipped. Authored with Claude Code. Reviewed By: digantdesai Differential Revision: D125532471 --- .../transforms/fuse_batch_norm_with_linear.py | 169 ++++++++++++++++++ backends/transforms/targets.bzl | 27 +++ .../test/test_fuse_batch_norm_with_linear.py | 136 ++++++++++++++ 3 files changed, 332 insertions(+) create mode 100644 backends/transforms/fuse_batch_norm_with_linear.py create mode 100644 backends/transforms/test/test_fuse_batch_norm_with_linear.py diff --git a/backends/transforms/fuse_batch_norm_with_linear.py b/backends/transforms/fuse_batch_norm_with_linear.py new file mode 100644 index 00000000000..66c61d78977 --- /dev/null +++ b/backends/transforms/fuse_batch_norm_with_linear.py @@ -0,0 +1,169 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator + +import torch + +from executorch.backends.transforms.utils import get_param_tensor, is_param_node +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +from torch.nn.utils.fusion import fuse_linear_bn_weights + + +class FuseBatchNormWithLinearPass(ExportPass): + """ + Fold an inference batch norm into the preceding linear's weight and bias. + """ + + def __init__(self, exported_program: ExportedProgram): + super().__init__() + self.exported_program = exported_program + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + counter = 0 + for linear in graph.nodes: + # We want to discover a chain of linear -> batch_norm. + # Only proceed if the current node is a linear node, and has a + # single user/successor. + if ( + linear.target != exir_ops.edge.aten.linear.default + or len(linear.users) != 1 + ): + continue + + # The single user of linear op must be batch_norm. If not, bail. + bn = list(linear.users.keys())[0] + if ( + bn.target != exir_ops.edge.aten.native_batch_norm.default + and bn.target + != exir_ops.edge.aten._native_batch_norm_legit_no_training.default + ): + continue + + if not self.can_fuse(linear, bn, self.exported_program): + continue + + # Linear has 2 or 3 args: (input, weight, bias). + assert len(linear.args) in (2, 3) + + linear_weight = get_param_tensor(self.exported_program, linear.args[1]) + assert linear_weight is not None + + linear_bias = ( + get_param_tensor(self.exported_program, linear.args[2]) + if len(linear.args) == 3 + else None + ) + + # Get the parameters from the batchnorm op + assert ( + bn.target == exir_ops.edge.aten.native_batch_norm.default + and len(bn.args) == 8 + ) or ( + bn.target + == exir_ops.edge.aten._native_batch_norm_legit_no_training.default + and len(bn.args) == 7 + ) + bn_weight = get_param_tensor(self.exported_program, bn.args[1]) + bn_bias = get_param_tensor(self.exported_program, bn.args[2]) + + running_mean = get_param_tensor(self.exported_program, bn.args[3]) + assert running_mean is not None + + running_var = get_param_tensor(self.exported_program, bn.args[4]) + assert running_var is not None + + if bn_bias is None: + bn_bias = torch.zeros_like(running_mean) + + # args[7] for native_batch_norm, but args[6] for + # _native_batch_norm_legit_no_training (which doesn't have training + # as an arg) + eps = bn.args[-1] + + # Compute the updated weight and bias after fusing linear op + # with batchnorm op. + fused_weight, fused_bias = fuse_linear_bn_weights( + linear_weight, + linear_bias, + running_mean, + running_var, + eps, + bn_weight, + bn_bias, + ) + + # Modify the graph by updating the weight and bias of linear op + # with the fused weight and bias params, and replacing all the users + # of getitem(batchnorm) with the linear op. + with graph.inserting_before(linear): + fused_weight_name = f"_fused_with_bn_weight_{counter}" + graph_module.register_parameter(fused_weight_name, fused_weight) + fused_weight_node = graph.get_attr(fused_weight_name) + fused_bias_name = f"_fused_with_bn_bias_{counter}" + graph_module.register_parameter(fused_bias_name, fused_bias) + fused_bias_node = graph.get_attr(fused_bias_name) + + # Update the weight and bias of linear op + linear_args = list(linear.args) + ([None] if len(linear.args) == 2 else []) + linear_args[1] = fused_weight_node + linear_args[2] = fused_bias_node + linear.args = tuple(linear_args) + # Remove any use of batchnorm from the graph + for user in bn.users.copy(): + assert user.target == operator.getitem + user.replace_all_uses_with(linear) + graph.erase_node(user) + + graph.erase_node(bn) + + counter += 1 + + if counter == 0: + return PassResult(graph_module, False) + + graph_module.recompile() + # To Regenerate meta data and shape information, retrace module + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, True) + + @staticmethod + def can_fuse( + linear: torch.fx.Node, bn: torch.fx.Node, program: ExportedProgram + ) -> bool: + """ + Determine whether a batch norm node can be fused with a preceding linear node. + """ + + # All the users of batchnorm node must be getitem ops. batchnorm + # returns a 3-element tuple. Each user must only access the first + # element of the tuple. + if [ + (user.target == operator.getitem and user.args[1] == 0) for user in bn.users + ].count(False): + return False + + if bn.target == exir_ops.edge.aten.native_batch_norm.default and bn.args[5]: + return False + + # Linear and batch norm share their channel axis only for rank-2 outputs. + linear_output = linear.meta.get("val") + if not isinstance(linear_output, torch.Tensor) or linear_output.dim() != 2: + return False + + params = (linear.args[1], bn.args[1], bn.args[3], bn.args[4]) + params += tuple( + param for param in linear.args[2:3] + bn.args[2:3] if param is not None + ) + return all( + isinstance(param, torch.fx.Node) and is_param_node(program, param) + for param in params + ) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 6b153c0fc6c..920f937b37a 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -134,6 +134,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"], @@ -511,6 +525,19 @@ def define_common_targets(): ], ) + runtime.python_test( + name = "test_fuse_batch_norm_with_linear", + srcs = [ + "test/test_fuse_batch_norm_with_linear.py", + ], + deps = [ + "//caffe2:torch", + ":fuse_batch_norm_with_linear", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_quantize_fused_convbn_bias_pass", srcs = [ diff --git a/backends/transforms/test/test_fuse_batch_norm_with_linear.py b/backends/transforms/test/test_fuse_batch_norm_with_linear.py new file mode 100644 index 00000000000..99ed4fbfe1f --- /dev/null +++ b/backends/transforms/test/test_fuse_batch_norm_with_linear.py @@ -0,0 +1,136 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import unittest + +import torch +from executorch.backends.transforms.fuse_batch_norm_with_linear import ( + FuseBatchNormWithLinearPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export + +_BN_OPS = { + exir_ops.edge.aten.native_batch_norm.default, + exir_ops.edge.aten._native_batch_norm_legit_no_training.default, +} + + +class _LinearBn(torch.nn.Module): + def __init__(self, linear: torch.nn.Module, bn: torch.nn.Module): + super().__init__() + self.linear = linear + self.bn = bn + + def forward(self, x): + return self.bn(self.linear(x)) + + +class _LinearBnWithRuntimeParameter(_LinearBn): + def __init__(self, parameter: str): + super().__init__(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8)) + self.parameter = parameter + + def forward(self, x, parameter): + x = torch.nn.functional.linear( + x, + self.linear.weight, + parameter if self.parameter == "linear_bias" else self.linear.bias, + ) + return torch.nn.functional.batch_norm( + x, + parameter if self.parameter == "running_mean" else self.bn.running_mean, + parameter if self.parameter == "running_var" else self.bn.running_var, + self.bn.weight, + parameter if self.parameter == "bn_bias" else self.bn.bias, + training=False, + eps=self.bn.eps, + ) + + +def _seeded(module: torch.nn.Module) -> torch.nn.Module: + torch.manual_seed(0) + bn = module.bn + bn.running_mean.uniform_(-1, 1) + bn.running_var.uniform_(0.5, 2) + bn.weight.data.uniform_(0.5, 2) + if bn.bias is not None: + bn.bias.data.uniform_(-1, 1) + return module.eval() + + +class TestFuseBatchNormWithLinear(unittest.TestCase): + def _check(self, module, inputs, *, fuse=True): + module = _seeded(module) + expected = module(*inputs) + + edge = to_edge( + export(module, inputs, strict=True), + compile_config=EdgeCompileConfig( + preserve_ops=[torch.ops.aten.linear.default], + ), + ) + self.assertTrue( + any(n.target in _BN_OPS for n in edge.exported_program().graph.nodes) + ) + result = FuseBatchNormWithLinearPass(edge.exported_program())( + copy.deepcopy(edge.exported_program().graph_module) + ) + self.assertEqual(result.modified, fuse) + edge = edge.transform([FuseBatchNormWithLinearPass(edge.exported_program())]) + ep = edge.exported_program() + + self.assertEqual(any(n.target in _BN_OPS for n in ep.graph.nodes), not fuse) + torch.testing.assert_close(ep.module()(*inputs), expected, rtol=1e-5, atol=1e-5) + + def test_linear(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8)), + (torch.randn(4, 16),), + ) + + def test_linear_no_bias(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8, bias=False), torch.nn.BatchNorm1d(8)), + (torch.randn(4, 16),), + ) + + def test_linear_wide(self): + self._check( + _LinearBn(torch.nn.Linear(64, 32), torch.nn.BatchNorm1d(32)), + (torch.randn(2, 64),), + ) + + def test_batch_norm_no_bias(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8, bias=False)), + (torch.randn(4, 16),), + ) + + def test_non_rank_two_is_not_fused(self): + for bn_type, spatial_shape in ( + (torch.nn.BatchNorm1d, ()), + (torch.nn.BatchNorm2d, (3,)), + (torch.nn.BatchNorm3d, (3, 2)), + ): + for channels in (4, 8): + with self.subTest(batch_norm=bn_type.__name__, channels=channels): + self._check( + _LinearBn(torch.nn.Linear(16, 8), bn_type(channels)), + (torch.randn(2, channels, *spatial_shape, 16),), + fuse=False, + ) + + def test_runtime_parameters_are_not_fused(self): + for parameter in ("linear_bias", "bn_bias", "running_mean", "running_var"): + with self.subTest(parameter=parameter): + self._check( + _LinearBnWithRuntimeParameter(parameter), + (torch.randn(4, 16), torch.rand(8) + 0.5), + fuse=False, + ) From 082d9ed2c47ac62e3feebb68ede4fa0b3a5d522d Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 27 Oct 2026 09:08:35 -0700 Subject: [PATCH 3/4] Wire batch-norm fusion into CpuBackend.preprocess (#24193) Summary: Run the shared `FuseBatchNormWithConvPass` and `FuseBatchNormWithLinearPass` in `CpuBackend.preprocess`. Unfused `_native_batch_norm_legit_no_training` already executes through the reference provider, but as a separate kernel with its own serialized weight and stats; folding it into the preceding conv/linear removes that kernel from the delegate. The passes are listed in `CPU_PREPROCESS_PASSES` (mirroring `XNNPACKPassManager.passes`) so every lowering flow gets them from one extension point. Fusion introduces `get_attr` constants, so preprocess follows with `lift_constant_tensor_pass` (the native serializer rejects unlifted constants), then drops parameter/buffer/constant inputs left without users. Without that step the delegate serialized the original weights and batch norm stats next to the fused ones. Authored with Claude Code. Reviewed By: digantdesai Differential Revision: D126784606 --- backends/cpu/BUCK | 2 + backends/cpu/preprocess.py | 63 ++++++++++++++++++++++++++++- backends/cpu/test/test_aot.py | 66 ++++++++++++++++++++++++++++++- backends/cpu/test/test_runtime.py | 26 ++++++++++++ 4 files changed, 154 insertions(+), 3 deletions(-) diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index 0171cb67971..24daafd26fe 100644 --- a/backends/cpu/BUCK +++ b/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", diff --git a/backends/cpu/preprocess.py b/backends/cpu/preprocess.py index 16bc83e19ac..42eb88e298b 100644 --- a/backends/cpu/preprocess.py +++ b/backends/cpu/preprocess.py @@ -12,6 +12,12 @@ 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, @@ -19,8 +25,20 @@ 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( diff --git a/backends/cpu/test/test_aot.py b/backends/cpu/test/test_aot.py index e3d75d3d7f9..42fd78b9546 100644 --- a/backends/cpu/test/test_aot.py +++ b/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 ( diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py index d4b92a3f1d2..62ceaed5c91 100644 --- a/backends/cpu/test/test_runtime.py +++ b/backends/cpu/test/test_runtime.py @@ -116,3 +116,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 + ) From 65b4187fa63a3ea3ed3962e915650cfc334efdcf Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 27 Oct 2026 09:08:35 -0700 Subject: [PATCH 4/4] Route _neg_view_copy to ET neg kernel (#24194) Summary: `aten::_neg_view_copy` (negation with view-copy autograd metadata, identical numerics to neg for inference) had no provider route, crashing at backend init (0x10). Add it to the ET spec table mapped to `aten::neg.out`. Reviewed By: digantdesai Differential Revision: D126929421 --- .../providers/reference/ReferenceProvider.cpp | 15 ++++++++++++ backends/cpu/test/test_runtime.py | 23 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp index 8e5267b5d75..8dc54961c0a 100644 --- a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp +++ b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp @@ -106,6 +106,7 @@ constexpr auto kKernels = std::to_array({ 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->node = &node; instruction->kernel = route->kernel; diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py index 62ceaed5c91..b3d1df05982 100644 --- a/backends/cpu/test/test_runtime.py +++ b/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):