Fold parameter-only subgraphs before XNNPACK partitioning

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

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,13 +23,29 @@
     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)
 logger = logging.getLogger(__name__)
 
 
 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.

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)

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
 

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