diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 91481da60ce..76345d9a66f 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/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 diff --git a/backends/xnnpack/test/BUCK b/backends/xnnpack/test/BUCK index 9392ae12a50..19532b9e2c8 100644 --- a/backends/xnnpack/test/BUCK +++ b/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", diff --git a/backends/xnnpack/test/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index 67fb6819920..8f21bd20ac9 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/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): diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index 6ec9e320eb8..2880cf1cc29 100644 --- a/extension/llm/export/builder.py +++ b/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(