From 605171ea2540a55d52abea44f7955cf295bd3387 Mon Sep 17 00:00:00 2001 From: sohailshk <122166523+sohailshk@users.noreply.github.com> Date: Mon, 9 Nov 2026 07:14:51 +0530 Subject: [PATCH 1/3] [XNNPACK] Skip linear conversion when the weight transpose was folded ConvertToLinearPass rebuilds linear from an [add]mm by taking the weight from before its transpose. When constant propagation or a transpose removal pass has already folded that transpose into the constant, the [add]mm reads the transposed weight directly and the rebuilt linear transposes it a second time: square weights give wrong values and other shapes make the pass fail. Only convert nodes whose weight still comes through a transpose, and leave the rest as [add]mm. Fixes #10499 --- backends/xnnpack/_passes/convert_to_linear.py | 4 ++++ backends/xnnpack/test/BUCK | 1 + .../test/passes/test_convert_to_linear.py | 19 +++++++++++++++++++ 3 files changed, 24 insertions(+) diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 91481da60ce..49b733285b3 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 @@ -165,7 +166,10 @@ 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 + # linear takes the weight from before its transpose, which + # constant propagation can fold away. if node.target in self.targets + and is_transpose(self.get_arg(node, "weight")) } # 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..60d9b72da02 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/backends/xnnpack/test/passes/test_convert_to_linear.py @@ -9,6 +9,7 @@ import torch 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 +43,24 @@ 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"]) + .run_method_and_compare_outputs() + ) + def test_fp32_view_after_two_linears_with_dynamic_shapes(self): class Model(torch.nn.Module): def __init__(self): From 22b3f546b5659c29e2b17c6f157a3ae9b17668dd Mon Sep 17 00:00:00 2001 From: sohailshk <122166523+sohailshk@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:57:03 +0530 Subject: [PATCH 2/3] Check the weight that linear takes for a transpose, not just the direct argument A bmm weight arrives through view_copy and expand_copy, so the direct argument check skipped bmm that was still convertible. A bmm whose bias add follows a view_copy is left alone, since the pass cannot find that bias and would drop it. Also drop the stale TODO about #10499 and test lowering of the folded graph. --- backends/xnnpack/_passes/convert_to_linear.py | 25 ++++++-- .../test/passes/test_convert_to_linear.py | 58 +++++++++++++++++++ extension/llm/export/builder.py | 3 - 3 files changed, 79 insertions(+), 7 deletions(-) diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 49b733285b3..b743ce4ff61 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/backends/xnnpack/_passes/convert_to_linear.py @@ -75,6 +75,24 @@ 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, and a bmm must have no + bias, since find_bias_for_mm only finds an add right after an mm. + """ + 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 + weight = ConvertToLinearPass.get_arg(node, "weight") + while 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 @@ -161,15 +179,14 @@ def call(self, graph_module: torch.fx.GraphModule): graph_module.graph, self.linear_modules ) + # linear takes the weight from before its transpose, which constant + # propagation can fold away. src_node_dict = { node: src_partition for src_partitions in src_partition_dict.values() for src_partition in src_partitions for node in src_partition.nodes - # linear takes the weight from before its transpose, which - # constant propagation can fold away. - if node.target in self.targets - and is_transpose(self.get_arg(node, "weight")) + 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/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index 60d9b72da02..d95c331c786 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/backends/xnnpack/test/passes/test_convert_to_linear.py @@ -4,6 +4,7 @@ # 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 @@ -58,9 +59,66 @@ def test_fp32_folded_weight_transpose_is_not_converted_to_linear(self): {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_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( From 59bdd6dd04fc303c284bafc15c61c4913da7a601 Mon Sep 17 00:00:00 2001 From: sohailshk <122166523+sohailshk@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:14:00 +0530 Subject: [PATCH 3/3] Stop the weight walk in can_convert at the partition boundary can_convert followed args[0] past the partition inputs, so a transpose upstream of the linear (for example F.linear(x, (a.t() + 1).t()) once the permute pair at the weight is cancelled) counted as the weight transpose, and create_linear then transposed the weight a second time. A weight such as a cat, whose first arg is a list, raised AttributeError. Stop at the same nodes as find(): partition inputs and params, non-call_function nodes, dequants, and anything that is not a Node. Add tests for both cases and reword the docstring. --- backends/xnnpack/_passes/convert_to_linear.py | 17 +++++-- .../test/passes/test_convert_to_linear.py | 49 +++++++++++++++++++ 2 files changed, 61 insertions(+), 5 deletions(-) diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index b743ce4ff61..76345d9a66f 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/backends/xnnpack/_passes/convert_to_linear.py @@ -79,15 +79,24 @@ def get_arg(node: torch.fx.Node, arg: str): 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, and a bmm must have no - bias, since find_bias_for_mm only finds an add right after an mm. + 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 weight.op == "call_function" and not is_dequant(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] @@ -179,8 +188,6 @@ def call(self, graph_module: torch.fx.GraphModule): graph_module.graph, self.linear_modules ) - # linear takes the weight from before its transpose, which constant - # propagation can fold away. src_node_dict = { node: src_partition for src_partitions in src_partition_dict.values() diff --git a/backends/xnnpack/test/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index d95c331c786..8f21bd20ac9 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/backends/xnnpack/test/passes/test_convert_to_linear.py @@ -8,6 +8,9 @@ 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 @@ -91,6 +94,52 @@ def forward(self, x, weight): .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):