[XNNPACK] Skip linear conversion when the weight transpose was folded

backends/xnnpack/_passes/convert_to_linear.py

@@ -12,6 +12,7 @@
 from executorch.backends.transforms import get_shape
 from executorch.backends.transforms.addmm_mm_to_linear import (
     apply_addmm_mm_to_linear_transform,
+    is_transpose,
 )
 from executorch.backends.xnnpack.utils.quant_utils import is_dequant
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -74,6 +75,33 @@ def get_arg(node: torch.fx.Node, arg: str):
             map_ = {"input": 0, "weight": 1}
             return None if arg == "bias" else node.args[map_[arg]]
 
+    @staticmethod
+    def can_convert(node: torch.fx.Node, src_partition: SourcePartition) -> bool:
+        """
+        Whether a linear can be rebuilt from node. A transpose must lie between
+        the weight that find() resolves to and the node, within the partition.
+        A bmm with a bias is left alone: the bias add that follows it is not
+        part of the bmm's linear (find_bias_for_mm is not used for bmm), so
+        rebuilding would drop the bias.
+        """
+        if node.target == exir_ops.edge.aten.bmm.default and any(
+            n.target == exir_ops.edge.aten.add.Tensor for n in src_partition.nodes
+        ):
+            return False
+        # Stop where find() stops, so the walk does not leave the partition.
+        boundary = src_partition.input_nodes + src_partition.params
+        weight = ConvertToLinearPass.get_arg(node, "weight")
+        while (
+            isinstance(weight, torch.fx.Node)
+            and weight not in boundary
+            and weight.op == "call_function"
+            and not is_dequant(weight)
+        ):
+            if is_transpose(weight):
+                return True
+            weight = weight.args[0]
+        return False
+
     def find_bias_for_mm(self, src_partition: SourcePartition, mm_node: torch.fx.Node):
         """
         For linear decomposed with mm + add, find bias in src partition
@@ -165,7 +193,7 @@ def call(self, graph_module: torch.fx.GraphModule):
                 for src_partitions in src_partition_dict.values()
                 for src_partition in src_partitions
                 for node in src_partition.nodes
-                if node.target in self.targets
+                if node.target in self.targets and self.can_convert(node, src_partition)
             }
 
             # No more [add]mm target in source partitions

backends/xnnpack/test/BUCK

@@ -60,6 +60,7 @@ fbcode_target(_kind = runtime.python_test,
         "//executorch/exir:pass_base",
         "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib",
         "//executorch/exir/dialects:lib",
+        "//executorch/exir/passes:constant_prop_pass",
     ],
     external_deps = [
         "libtorch",

backends/xnnpack/test/passes/test_convert_to_linear.py

@@ -4,11 +4,16 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import itertools
 import unittest
 
 import torch
+from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import (
+    FuseCascadedTransposeOrPermuteOps,
+)
 from executorch.backends.xnnpack._passes.convert_to_linear import ConvertToLinearPass
 from executorch.backends.xnnpack.test.tester import RunPasses, Tester
+from executorch.exir.passes.constant_prop_pass import constant_prop_pass
 
 
 class TestConvertToLinear(unittest.TestCase):
@@ -42,6 +47,127 @@ def test_fp32_convert_to_linear(self):
                 .run_method_and_compare_outputs()
             )
 
+    def test_fp32_folded_weight_transpose_is_not_converted_to_linear(self):
+        for in_size, out_size, bias in [(4, 8, True), (6, 6, True), (6, 6, False)]:
+            torch._dynamo.reset()
+            linear = torch.nn.Linear(in_size, out_size, bias=bias)
+            tester = Tester(linear, (torch.randn(4, in_size),)).export().to_edge()
+            # Folds the weight transpose, so the [add]mm reads a weight that
+            # linear would need transposed back.
+            constant_prop_pass(tester.get_artifact().exported_program())
+            mm = "addmm" if bias else "mm"
+            (
+                tester.run_passes(self.PassStage)
+                .check_count(
+                    {f"executorch_exir_dialects_edge__ops_aten_{mm}_default": 1}
+                )
+                .check_not(["executorch_exir_dialects_edge__ops_aten_linear_default"])
+                .partition()
+                .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+                .check_not([f"executorch_exir_dialects_edge__ops_aten_{mm}_default"])
+                .to_executorch()
+                .serialize()
+                .run_method_and_compare_outputs()
+            )
+
+    def test_fp32_only_folded_linear_is_not_converted(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(6, 6)
+
+            def forward(self, x, weight):
+                # The weight is an input, so its transpose is not folded.
+                return torch.nn.functional.linear(self.linear(x), weight)
+
+        inputs = (torch.randn(4, 6), torch.randn(6, 6))
+        tester = Tester(Model(), inputs).export().to_edge()
+        constant_prop_pass(tester.get_artifact().exported_program())
+        (
+            tester.run_passes(self.PassStage)
+            .check_count(
+                {
+                    "executorch_exir_dialects_edge__ops_aten_addmm_default": 1,
+                    "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+                }
+            )
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_transpose_upstream_of_partition_is_not_converted(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = torch.nn.Parameter(torch.randn(6, 6))
+
+            def forward(self, x):
+                # The transposes around the add are outside the linear.
+                return torch.nn.functional.linear(x, (self.a.t() + 1).t())
+
+        x = torch.randn(4, 6)
+        tester = Tester(Model(), (x,)).export().to_edge()
+        # Cancels the permute pair at the weight, so the mm reads the add
+        # directly and the only permute left is upstream of the partition.
+        FuseCascadedTransposeOrPermuteOps()(
+            tester.get_artifact().exported_program().graph_module
+        )
+        (
+            tester.check_count(
+                {"executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1}
+            )
+            .run_passes(self.PassStage)
+            .check_count({"executorch_exir_dialects_edge__ops_aten_mm_default": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_linear_default"])
+            .run_method_and_compare_outputs(inputs=(x,))
+        )
+
+    def test_fp32_weight_with_non_node_arg_is_not_converted(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, w1, w2):
+                return torch.nn.functional.linear(x, torch.cat([w1, w2]).t())
+
+        inputs = (torch.randn(4, 6), torch.randn(3, 6), torch.randn(3, 6))
+        tester = Tester(Model(), inputs).export().to_edge()
+        # Cancels the permute pair at the weight, so the mm reads the cat, whose
+        # first arg is a list.
+        FuseCascadedTransposeOrPermuteOps()(
+            tester.get_artifact().exported_program().graph_module
+        )
+        (
+            tester.run_passes(self.PassStage)
+            .check_count({"executorch_exir_dialects_edge__ops_aten_mm_default": 1})
+            .check_not(["executorch_exir_dialects_edge__ops_aten_linear_default"])
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_bmm_in_linear_partition(self):
+        class Model(torch.nn.Module):
+            def __init__(self, bias):
+                super().__init__()
+                self.linear = torch.nn.Linear(8, 6, bias=bias)
+
+            def forward(self, x):
+                # A non-contiguous input makes the linear decompose to bmm.
+                return self.linear(x.transpose(1, 2))
+
+        x = torch.randn(3, 8, 4)
+        for bias, fold in itertools.product([False, True], [False, True]):
+            torch._dynamo.reset()
+            tester = Tester(Model(bias), (x,)).export().to_edge()
+            if fold:
+                constant_prop_pass(tester.get_artifact().exported_program())
+            tester.run_passes(self.PassStage)
+            # The bias add of a bmm is not found, so it stays a bmm.
+            if fold or bias:
+                tester.check_count(
+                    {"executorch_exir_dialects_edge__ops_aten_bmm_default": 1}
+                ).check_not(["executorch_exir_dialects_edge__ops_aten_linear_default"])
+            else:
+                tester.check_count(
+                    {"executorch_exir_dialects_edge__ops_aten_linear_default": 1}
+                ).check_not(["executorch_exir_dialects_edge__ops_aten_bmm_default"])
+            tester.run_method_and_compare_outputs(inputs=(x,))
+
     def test_fp32_view_after_two_linears_with_dynamic_shapes(self):
         class Model(torch.nn.Module):
             def __init__(self):

extension/llm/export/builder.py

@@ -552,9 +552,6 @@ def to_executorch(
         # If there are Linear operations left in the graph, let's execute
         # them with the optimized op_linear rather than materializing a
         # transpose followed by a regular op_mm.
-        # TODO: ConvertToLinearPass is not a sound pass and must be called before
-        # const propagation.  It requires fixing:
-        # https://github.com/pytorch/executorch/issues/10499
         self.edge_manager.transform([ConvertToLinearPass()])
 
         self.export_program = self.edge_manager.to_executorch(