[XNNPACK] Delegate GEMMs whose weight is computed from parameters

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