diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index f0e1f3a1487..18e9006b51a 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 + ) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 8314bbb5695..22fced32b8a 100644 --- a/backends/transforms/targets.bzl +++ b/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"],