diff --git a/backends/xnnpack/partition/BUCK b/backends/xnnpack/partition/BUCK index 2581043d5b2..c5688669064 100644 --- a/backends/xnnpack/partition/BUCK +++ b/backends/xnnpack/partition/BUCK @@ -18,6 +18,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/exir/backend:partitioner", "//executorch/exir/backend:utils", "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib", + "//executorch/exir/passes:constant_prop_pass", ], ) diff --git a/backends/xnnpack/partition/xnnpack_partitioner.py b/backends/xnnpack/partition/xnnpack_partitioner.py index 5327121b158..8d3e9ca0e1f 100644 --- a/backends/xnnpack/partition/xnnpack_partitioner.py +++ b/backends/xnnpack/partition/xnnpack_partitioner.py @@ -9,6 +9,8 @@ import logging from typing import List, Optional, Type, Union +import torch + from executorch.backends.xnnpack.partition.config import ALL_PARTITIONER_CONFIGS from executorch.backends.xnnpack.partition.config.xnnpack_config import ( ConfigPrecisionType, @@ -21,6 +23,7 @@ ConfigerationBasedPartitioner, ) from executorch.exir.backend.partitioner import DelegationSpec +from executorch.exir.passes.constant_prop_pass import constant_prop_pass from torch.fx.passes.infra.partitioner import Partition logging.basicConfig(level=logging.WARNING) @@ -28,6 +31,21 @@ class XnnpackPartitioner(ConfigerationBasedPartitioner): + # constant_prop_pass skips aten.full at the edge level so that a scalar + # fill does not become a stored tensor. Before decomposition the same + # tensors come from these factory ops. + _CONSTANT_PROP_SKIP_TARGETS = frozenset( + { + torch.ops.aten.full.default, + torch.ops.aten.full_like.default, + torch.ops.aten.ones.default, + torch.ops.aten.ones_like.default, + torch.ops.aten.zeros.default, + torch.ops.aten.zeros_like.default, + } + ) + _CONSTANT_PROP_SKIP_NAMESPACES = ("quantized_decomposed", "torchao") + def __init__( self, configs: Optional[List[Type[XNNPartitionerConfig]]] = None, @@ -83,6 +101,30 @@ def _check_if_called_from_to_backend(self) -> bool: return True return False + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """ + Fold subgraphs whose inputs are all parameters into constants. + + The partitioner configs require a static weight, so a convolution or a + linear whose weight is computed from parameters, for example under + torch.nn.utils.parametrizations.weight_norm, would otherwise be left + to the portable kernels together with the weight computation. + """ + # Quantization primitives are kept as well, so that the Q/DQ chain + # convert_pt2e or torchao's quantize_ leaves on a weight stays in the + # graph. A folded dequantize would hand the delegate a float weight. + skip_targets = set(self._CONSTANT_PROP_SKIP_TARGETS) + for node in exported_program.graph.nodes: + if ( + node.op == "call_function" + and getattr(node.target, "namespace", None) + in self._CONSTANT_PROP_SKIP_NAMESPACES + ): + skip_targets.add(node.target) + return constant_prop_pass(exported_program, custom_skip_targets=skip_targets) + def partition(self, exported_program): """ Override partition to add deprecation warning when called from to_backend. diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..15270d808ef 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -161,3 +161,104 @@ def example_inputs(self): fwd2_et = executorch_module.run_method("forward_2", example_inputs) self.assertTrue(torch.allclose(fwd1_eager, fwd1_et[0], 1e-3)) self.assertTrue(torch.allclose(fwd2_eager, fwd2_et[0], 1e-3)) + + def test_parametrized_weight_is_folded_before_partitioning(self): + """ + A weight computed from parameters (here weight_norm) is folded into a + constant before partitioning, so the convolution is delegated instead + of falling back to the portable kernels with the weight computation. + """ + + class ParametrizedConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.utils.parametrizations.weight_norm( + torch.nn.Conv1d(4, 4, 3) + ) + + def forward(self, x): + return self.conv(x) + + model = ParametrizedConv().eval() + example_inputs = (torch.randn(1, 4, 8),) + eager = model(*example_inputs) + + edge = to_edge_transform_and_lower( + export(model, example_inputs), partitioner=[XnnpackPartitioner()] + ) + call_functions = [ + node + for node in edge.exported_program().graph_module.graph.nodes + if node.op == "call_function" + ] + delegates = [ + node + for node in call_functions + if node.target == torch.ops.higher_order.executorch_call_delegate + ] + self.assertEqual(len(delegates), 1) + # The delegate call and the getitem on its output are all that is left. + self.assertEqual(len(call_functions), 2) + + executorch_module = _load_for_executorch_from_buffer( + edge.to_executorch().buffer + ) + self.assertTrue( + torch.allclose(executorch_module.forward(example_inputs)[0], eager, 1e-5) + ) + + def test_pre_decomposition_folding_keeps_quantization_primitives(self): + """ + Folding must not touch the Q/DQ chain that convert_pt2e leaves on a + quantized weight, or the weight would be dequantized at export time. + """ + from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, + ) + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + model = self.SimpleModel().eval() + example_inputs = (torch.randn(2, 10),) + quantizer = XNNPACKQuantizer() + quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True)) + prepared = prepare_pt2e(export(model, example_inputs).module(), quantizer) + prepared(*example_inputs) + converted = convert_pt2e(prepared) + + def quant_targets(ep): + return sorted( + str(node.target) + for node in ep.graph.nodes + if node.op == "call_function" + and "quantized_decomposed" in str(node.target) + ) + + class GroupwiseLinear(torch.nn.Module): + """A weight stored as int8 groups, the way 4-bit LLM exports do.""" + + def __init__(self): + super().__init__() + self.register_buffer( + "weight", torch.randint(-8, 8, (8, 16), dtype=torch.int8) + ) + self.register_buffer("scales", torch.rand(8, 2)) + self.register_buffer("zeros", torch.zeros(8, 2, dtype=torch.int8)) + + def forward(self, x): + weight = torch.ops.quantized_decomposed.dequantize_per_channel_group( + self.weight, self.scales, self.zeros, -8, 7, torch.int8, 8, x.dtype + ) + return torch.nn.functional.linear(x, weight) + + for model, example_inputs in ( + (converted, example_inputs), + (GroupwiseLinear(), (torch.randn(2, 16),)), + ): + exported = export(model, example_inputs) + before = quant_targets(exported) + self.assertGreater(len(before), 0) + after = quant_targets( + XnnpackPartitioner().transform_for_pre_decomposition(exported) + ) + self.assertEqual(before, after) diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 11640d875c0..ea8ee1ad3a9 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -146,6 +146,10 @@ def get_propagated_const_tensor_dict( node.op != "call_function" or node.target is memory.alloc or node.target in all_skip_targets + # Ops with side effects (RNG draws, mutation) have to run at + # runtime. `aten.rand` has no tensor inputs, so without this check + # it would be folded into a single frozen draw. + or node.is_impure() ): continue diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0b586bd44cd..54211a490a3 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2475,6 +2475,31 @@ def forward(self, x): # 1 constant: a (= self.w @ self.cst) self.assertEqual(1, len(pass_result.constants)) + def test_constant_prop_pass_skips_nondeterministic_ops(self) -> None: + """ + Ops that draw from the RNG take no tensor inputs, so they look constant + to the pass. They have to stay in the graph: folding one would freeze a + single random draw into the program. + """ + + class RandomAdd(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + torch.rand(4) + + x = torch.zeros(4) + edge = to_edge(export(RandomAdd(), (x,), strict=True)) + new_ep = constant_prop_pass(edge.exported_program()) + + rand_nodes = [ + node + for node in new_ep.graph.nodes + if node.target == exir_ops.edge.aten.rand.default + ] + self.assertEqual(len(rand_nodes), 1) + self.assertEqual(len(new_ep.constants), 0) + module = new_ep.module() + self.assertFalse(torch.equal(module(x), module(x))) + def test_constant_prop_pass_zero_stride_tensors(self) -> None: """ Test that constant propagation correctly handles tensors with zero strides