[XNNPACK] Delegate GEMMs whose weight is computed from parameters

A convolution, linear, addmm or mm whose weight or bias is computed from
parameters, such as a convolution under weight_norm or spectral_norm, was
declined by the XNNPACK partitioner and ran on the portable kernels together
with the computation of its weight.

The GEMM configs now accept such a weight or bias when it is computed with
pure ops from parameters, unmutated buffers or constants, the computation
feeds only that GEMM, and the folded tensor is no larger than the parameters
it replaces. The computation is partitioned with the GEMM, and
FoldComputedParamsPass folds it into a constant at the start of XNNPACK
preprocess, inside the delegate, so the top-level program is not modified.

AddmmConfig now also rejects addmm with non-default alpha or beta, which the
XNNPACK lowering ignored.

Fixes #22078

This change was authored with Claude Code.

Author: JakeStevens

backends/xnnpack/_passes/__init__.py

@@ -26,6 +26,9 @@
 )
 from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
+from executorch.backends.xnnpack._passes.fold_computed_params import (
+    FoldComputedParamsPass,
+)
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import (
@@ -74,6 +77,7 @@ def __init__(
         if not passes:
             # All the XNNPACK passes
             self.passes = [
+                FoldComputedParamsPass,
                 RewriteFp16SiluPass,
                 XNNPACKRemoveCloneOpsTransform,
                 # TODO - remove this pass once we have a better support for dim_order ops lowering

backends/xnnpack/_passes/fold_computed_params.py

@@ -0,0 +1,70 @@
+# 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 torch
+from executorch.backends.transforms.utils import create_constant_placeholder
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.backends.xnnpack.utils.utils import COMPUTED_PARAM_TAG, get_param_tensor
+from executorch.exir.pass_base import PassResult
+from torch.export.graph_signature import InputKind
+
+
+class FoldComputedParamsPass(XNNPACKPass):
+    """
+    Folds the nodes that the partitioner tagged with COMPUTED_PARAM_TAG, which
+    compute a GEMM weight or bias from params, into a param holding the result.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+        names = {node.name for node in graph.nodes}
+        values = {}
+        for node in graph.nodes:
+            # Checking a GEMM tags the nodes computing its weight or bias before
+            # the GEMM is accepted, so they are folded only next to that GEMM.
+            if not names & node.meta.get(COMPUTED_PARAM_TAG, set()):
+                continue
+            args, kwargs = torch.fx.node.map_arg(
+                (node.args, node.kwargs),
+                lambda n: (
+                    values[n]
+                    if n in values
+                    else get_param_tensor(self.exported_program, n)
+                ),
+            )
+            with torch.no_grad():
+                values[node] = node.target(*args, **kwargs)
+
+        if not values:
+            return PassResult(graph_module, False)
+
+        taken = (
+            names
+            | self.exported_program.state_dict.keys()
+            | self.exported_program.constants.keys()
+        )
+        for node, value in values.items():
+            if all(user in values for user in node.users):
+                continue
+            name = f"{node.name}_folded"
+            while name in taken:
+                name += "_"
+            taken.add(name)
+            with graph.inserting_before(next(iter(graph.nodes))):
+                param = create_constant_placeholder(
+                    exp_program=self.exported_program,
+                    graph=graph,
+                    kind=InputKind.PARAMETER,
+                    name=name,
+                    data=value,
+                )
+            node.replace_all_uses_with(param)
+
+        for node in reversed(values):
+            graph.erase_node(node)
+
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/xnnpack/partition/config/gemm_configs.py

@@ -29,6 +29,7 @@
 )
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
+    get_static_param_deps,
     is_depthwise_conv,
     is_getitem,
     is_node,
@@ -153,10 +154,11 @@ def _get_weight_deps(
         if precision == ConfigPrecisionType.FP32:
             # First find the weight
             weight_node = get_input_node(node, self.weight_idx)
-            if not is_param_node(ep, weight_node):
+            weight_deps = get_static_param_deps(ep, weight_node, node)
+            if weight_deps is None:
                 why(node, "Expected weight to be a static param")
                 return (False, [])
-            gemm_deps.append(weight_node)
+            gemm_deps.extend(weight_deps)
 
             return (True, gemm_deps)
         else:
@@ -248,10 +250,11 @@ def _get_bias_deps(
         if len(node.all_input_nodes) > 2 and self.bias_idx is not None:
             bias_node = get_input_node(node, self.bias_idx)
             if bias_node:
-                if not is_param_node(ep, bias_node):
+                bias_deps = get_static_param_deps(ep, bias_node, node)
+                if bias_deps is None:
                     why(node, "Expected bias to be a static param")
                     return (False, [])
-                gemm_deps.append(bias_node)
+                gemm_deps.extend(bias_deps)
 
         return (True, gemm_deps)
 
@@ -479,6 +482,12 @@ def __init__(self, **kwargs):
         self.src_partitions = None
         self.linear_modules = [torch.nn.functional.linear, torch.nn.Linear]
 
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        if node.kwargs.get("alpha", 1) != 1 or node.kwargs.get("beta", 1) != 1:
+            why(node, "XNNPACK does not support addmm with alpha or beta")
+            return False
+        return super().check_constraints(node, ep)
+
     def _get_weight_deps(
         self, node: torch.fx.Node, ep: ExportedProgram, precision: ConfigPrecisionType
     ) -> Tuple[bool, List[torch.fx.Node]]:

backends/xnnpack/test/passes/test_fold_computed_params.py

@@ -0,0 +1,235 @@
+# 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 unittest
+
+import torch
+import torch.nn.functional as F
+from executorch.backends.xnnpack.partition.config.generic_node_configs import MulConfig
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+from executorch.backends.xnnpack.test.tester import Tester, ToEdgeTransformAndLower
+from torch.nn.utils.parametrizations import weight_norm
+
+
+class TestFoldComputedParams(unittest.TestCase):
+    def setUp(self):
+        torch._dynamo.reset()
+
+    class WeightNormConv1d(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = weight_norm(
+                torch.nn.Conv1d(8, 8, 5, padding=2, groups=2), dim=2
+            )
+
+        def forward(self, x):
+            return self.conv(x)
+
+    class ComputedLinear(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 8))
+            self.bias = torch.nn.Parameter(torch.randn(4))
+
+        def forward(self, x):
+            return F.linear(x, self.weight * 2, self.bias + 1)
+
+    class FoldedNameCollision(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 8))
+            # The name the folded weight would get by default.
+            self.aten_mul_tensor_folded = torch.nn.Parameter(torch.randn(4))
+
+        def forward(self, x):
+            return F.linear(x, self.weight * 2, self.aten_mul_tensor_folded)
+
+    class MatmulWeightConv1d(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(3, 3))
+
+        def forward(self, x):
+            return F.conv1d(x, (self.weight @ self.weight).reshape(3, 3, 1))
+
+    class NestedGemmWeightConv1d(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.a = torch.nn.Parameter(torch.randn(3, 3))
+            self.b = torch.nn.Parameter(torch.randn(3, 3))
+
+        def forward(self, x):
+            return F.conv1d(x, (self.a @ self.b.flip([1])).reshape(3, 3, 1))
+
+    class ScaledConv1d(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 4, 3))
+
+        def forward(self, x):
+            return F.conv1d(x, self.weight * 2)
+
+    class ScaledAddmm(torch.nn.Module):
+        def __init__(self, computed):
+            super().__init__()
+            self.computed = computed
+            self.weight = torch.nn.Parameter(torch.randn(8, 4))
+            self.bias = torch.nn.Parameter(torch.randn(4))
+
+        def forward(self, x):
+            weight, bias = self.weight, self.bias
+            if self.computed:
+                weight, bias = weight * 2, bias + 1
+            return torch.addmm(bias, x, weight, alpha=2, beta=3)
+
+    class SourceUsedElsewhere(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 8))
+
+        def forward(self, x):
+            return F.linear(x, self.weight * 2), self.weight.sum()
+
+    class SourceRetainedByGemm(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(16))
+
+        def forward(self, x):
+            return F.linear(x, (self.weight * 2).unsqueeze(1), self.weight)
+
+    class WeightComputationFeedsActivation(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(16, 16))
+
+        def forward(self, x):
+            w = self.weight * 2
+            return w @ (w * 2) + x
+
+    class LowRankWeight(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.a = torch.nn.Parameter(torch.randn(16, 2))
+            self.b = torch.nn.Parameter(torch.randn(2, 16))
+
+        def forward(self, x):
+            return F.linear(x, self.a @ self.b)
+
+    class MutatedBufferSource(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.register_buffer("weight", torch.randn(4, 4, 3))
+
+        def forward(self, x):
+            y = F.conv1d(x, self.weight * 2)
+            self.weight.copy_(x.mean() * torch.ones(4, 4, 3))
+            return y
+
+    def _test_fully_delegated(self, module, inputs):
+        (
+            Tester(module.eval(), inputs)
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_weight_norm_conv1d(self):
+        self._test_fully_delegated(self.WeightNormConv1d(), (torch.randn(1, 8, 16),))
+
+    def test_fp32_computed_linear_weight_and_bias(self):
+        self._test_fully_delegated(self.ComputedLinear(), (torch.randn(2, 8),))
+
+    def test_fp32_folded_name_does_not_replace_param(self):
+        self._test_fully_delegated(self.FoldedNameCollision(), (torch.randn(2, 8),))
+
+    def test_fp32_computed_weight_of_rejected_gemm_not_folded(self):
+        # Conv1d with two dynamic dims is rejected after its weight was checked,
+        # while the matmul computing the weight is still delegated on its own.
+        dims = {
+            0: torch.export.Dim("batch", min=1, max=4),
+            2: torch.export.Dim("length", min=4, max=16),
+        }
+        for module in (self.MatmulWeightConv1d(), self.NestedGemmWeightConv1d()):
+            with self.subTest(module=type(module).__name__):
+                (
+                    Tester(
+                        module.eval(), (torch.randn(2, 3, 8),), dynamic_shapes=(dims,)
+                    )
+                    .export()
+                    .to_edge_transform_and_lower()
+                    .check(
+                        ["executorch_exir_dialects_edge__ops_aten_convolution_default"]
+                    )
+                    .to_executorch()
+                    .serialize()
+                    .run_method_and_compare_outputs()
+                )
+
+    def test_fp32_scaled_addmm_not_delegated(self):
+        for computed in (False, True):
+            with self.subTest(computed=computed):
+                (
+                    Tester(self.ScaledAddmm(computed).eval(), (torch.randn(2, 8),))
+                    .export()
+                    .to_edge_transform_and_lower()
+                    .check(["executorch_exir_dialects_edge__ops_aten_addmm_default"])
+                    .to_executorch()
+                    .serialize()
+                    .run_method_and_compare_outputs()
+                )
+
+    def test_fp32_weight_computed_by_prior_delegate_not_folded(self):
+        (
+            Tester(self.ScaledConv1d().eval(), (torch.randn(1, 4, 8),))
+            .export()
+            .to_edge_transform_and_lower(
+                ToEdgeTransformAndLower(
+                    [XnnpackPartitioner(configs=[MulConfig]), XnnpackPartitioner()]
+                )
+            )
+            .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_computed_weight_not_folded(self):
+        for module, inputs, gemm in (
+            (self.SourceUsedElsewhere(), (torch.randn(2, 8),), "aten_mm_default"),
+            (self.LowRankWeight(), (torch.randn(2, 16),), "aten_mm_default"),
+            (self.SourceRetainedByGemm(), (torch.randn(2, 1),), "aten_addmm_default"),
+            (
+                self.WeightComputationFeedsActivation(),
+                (torch.randn(16, 16),),
+                "aten_mm_default",
+            ),
+            (
+                self.MutatedBufferSource(),
+                (torch.randn(1, 4, 8),),
+                "aten_convolution_default",
+            ),
+        ):
+            with self.subTest(module=type(module).__name__):
+                (
+                    Tester(module.eval(), inputs)
+                    .export()
+                    .to_edge_transform_and_lower()
+                    .check([f"executorch_exir_dialects_edge__ops_{gemm}"])
+                )
+
+    def test_fp32_externally_stored_weight_not_folded(self):
+        tester = Tester(self.WeightNormConv1d().eval(), (torch.randn(1, 8, 16),))
+        for node in tester.export().get_artifact().graph.nodes:
+            if node.op == "placeholder" and node.name.startswith("p_"):
+                node.meta["custom"] = {"delegate_constant_tag": "weights.ptd"}
+        tester.to_edge_transform_and_lower().check(
+            ["executorch_exir_dialects_edge__ops_aten_convolution_default"]
+        )