From dcda28cdc8b4722dcb109e12925d3c8c94e37cb4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 1 Sep 2026 11:38:05 +0200 Subject: [PATCH 01/10] Only rewrite addmm/mm to linear when the weight is a constant `replace_addmm_mm_with_linear` turns `mm(x, transpose(w))` into `linear(x, w)` whenever the second operand is fed through a transpose. It never checks where `w` comes from, so a matmul against a tensor computed at runtime is rewritten into a linear too. That rewrite is not sound. `mm`/`addmm` place no constraint on their second operand, but backends prepack a linear's weight while building their delegate graph, which is only possible for a constant. The Vulkan runtime aborts when it gets anything else: Exception raised from toTensorRef at backends/vulkan/runtime/graph/containers/Value.h:266: (isTensorRef()) is false! Expected value to have type TensorRef, got TENSOR instead. reached through linear_packed_weight -> prepack_fp_linear_weight -> PrepackNode, at delegate build time. A model that multiplies an activation by a runtime-derived matrix fails to load with no indication of which op is responsible. Guard both rewrites on the operand actually being a parameter, buffer or lifted constant, and thread the owning program through so that check can be made. Callers that pass no program keep the previous behaviour for placeholders and only lose the rewrite for runtime-computed operands, which is the case that was broken anyway. The pass is used by the Vulkan and Samsung ENN backends. --- backends/transforms/addmm_mm_to_linear.py | 60 ++++++++++++-- .../test/test_addmm_mm_to_linear.py | 80 +++++++++++++++++++ 2 files changed, 134 insertions(+), 6 deletions(-) create mode 100644 backends/transforms/test/test_addmm_mm_to_linear.py diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index 358cbb7ac14..f6ae4cfae5e 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -4,12 +4,14 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from typing import Optional + import torch from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult - from executorch.exir.sym_util import eval_shape, eval_shape_upper_bound - +from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param +from torch.export import ExportedProgram _int64_max_dim_val = torch.iinfo(torch.int64).max - 1 @@ -32,6 +34,33 @@ def get_shape(input_node: torch.fx.Node): return upper_bound_shape +def is_constant_tensor( + node: torch.fx.Node, exported_program: Optional[ExportedProgram] +) -> bool: + """ + Whether `node` produces a tensor whose contents are known at build time. + + `mm`/`addmm` place no constraint on their second operand, but a `linear` + node's weight does carry one: backends prepack it while building the + delegate graph. Rewriting to `linear` is therefore only valid when the + operand really is a constant. + """ + if node.op == "get_attr": + return True + if node.op != "placeholder": + return False + if exported_program is None: + # Without the owning program a lifted parameter cannot be told apart + # from a user input. Placeholders were always rewritten before, so keep + # accepting them rather than regressing callers that pass no program. + return True + return ( + is_param(exported_program, node) + or is_buffer(exported_program, node) + or is_lifted_tensor_constant(exported_program, node) + ) + + def get_dqlinear_input(node: torch.fx.Node): ops = exir_ops.edge node_to_backtrack = node @@ -99,7 +128,9 @@ def replace_linear_view_copy_input_output(graph: torch.fx.Graph) -> torch.fx.Gra return graph -def replace_addmm_mm_with_linear(graph: torch.fx.Graph) -> torch.fx.Graph: +def replace_addmm_mm_with_linear( + graph: torch.fx.Graph, exported_program: Optional[ExportedProgram] = None +) -> torch.fx.Graph: """ Replace calls to addmm/mm with linear node Reason is that it simplifies the downstream logic of lowering to just linear node. @@ -125,6 +156,9 @@ def replace_addmm_mm_with_linear(graph: torch.fx.Graph) -> torch.fx.Graph: # Skip this node as it appears to be a standalone `addmm` continue weight_node = weight_t_node.args[0] + if not is_constant_tensor(weight_node, exported_program): + # A runtime-computed operand is a matmul, not a linear + continue args = (node.args[1], weight_node, node.args[0]) linear_node = graph.create_node( "call_function", ops.aten.linear.default, args @@ -142,6 +176,9 @@ def replace_addmm_mm_with_linear(graph: torch.fx.Graph) -> torch.fx.Graph: # Skip this node as it appears to be a standalone `mm` continue weight_node = weight_t_node.args[0] + if not is_constant_tensor(weight_node, exported_program): + # A runtime-computed operand is a matmul, not a linear + continue args = (node.args[0], weight_node) linear_node = graph.create_node( "call_function", ops.aten.linear.default, args @@ -158,13 +195,24 @@ def replace_addmm_mm_with_linear(graph: torch.fx.Graph) -> torch.fx.Graph: return graph -def apply_addmm_mm_to_linear_transform(graph: torch.fx.Graph) -> torch.fx.Graph: - graph = replace_addmm_mm_with_linear(graph) +def apply_addmm_mm_to_linear_transform( + graph: torch.fx.Graph, exported_program: Optional[ExportedProgram] = None +) -> torch.fx.Graph: + graph = replace_addmm_mm_with_linear(graph, exported_program) graph = replace_linear_view_copy_input_output(graph) return graph class AddmmToLinearTransform(ExportPass): + def __init__(self, exported_program: Optional[ExportedProgram] = None) -> None: + super().__init__() + # Backends that run this pass through a pass manager which threads the + # owning program set this attribute themselves; see + # backends/vulkan/vulkan_preprocess.py. + self._exported_program = exported_program + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - graph_module.graph = apply_addmm_mm_to_linear_transform(graph_module.graph) + graph_module.graph = apply_addmm_mm_to_linear_transform( + graph_module.graph, self._exported_program + ) return PassResult(graph_module, True) diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py new file mode 100644 index 00000000000..371c298061b --- /dev/null +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -0,0 +1,80 @@ +# 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 +from executorch.backends.transforms.addmm_mm_to_linear import AddmmToLinearTransform +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +def count_targets(graph: torch.fx.Graph, target) -> int: + return sum(1 for n in graph.nodes if n.op == "call_function" and n.target == target) + + +class TestAddmmToLinearTransform(unittest.TestCase): + def _transform(self, model, example_inputs): + edge = to_edge(torch.export.export(model, example_inputs, strict=True)) + program = edge.exported_program() + transform = AddmmToLinearTransform(program) + return transform(program.graph_module).graph_module.graph + + def test_constant_weight_is_rewritten_to_linear(self): + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc = torch.nn.Linear(8, 4) + + def forward(self, x): + return self.fc(x) + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.addmm.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + + def test_computed_weight_stays_a_matmul(self): + # `w` is produced at runtime, so the transposed matmul below is not a + # linear: backends prepack a linear's weight while building their graph + # and cannot do that for a value that only exists during execution. + class Model(torch.nn.Module): + def forward(self, x, w): + return torch.mm(x, (w * 2.0).t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8), torch.randn(4, 8))) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) + + def test_computed_bias_operand_stays_an_addmm(self): + class Model(torch.nn.Module): + def forward(self, x, w, b): + return torch.addmm(b, x, (w * 2.0).t()) + + graph = self._transform( + Model().eval(), + (torch.randn(2, 8), torch.randn(4, 8), torch.randn(4)), + ) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.addmm.default), 1) + + def test_user_input_weight_is_not_rewritten(self): + class Model(torch.nn.Module): + def forward(self, x, w): + return torch.nn.functional.linear(x, w) + + edge = to_edge( + torch.export.export( + Model().eval(), (torch.randn(2, 8), torch.randn(4, 8)), strict=True + ) + ) + program = edge.exported_program() + graph = AddmmToLinearTransform(program)(program.graph_module).graph_module.graph + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + + +if __name__ == "__main__": + unittest.main() From 3ebf9e22b26058d8c19163febc7de58fefdc29d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Fri, 11 Sep 2026 15:10:02 +0200 Subject: [PATCH 02/10] Exclude mutated buffers, and thread the program through every caller is_buffer() also accepts a buffer the program writes back to. That is state, not a build-time constant: a delegate that prepacked it would keep returning the value it was built with. Reject buffers named in buffers_to_mutate. Accepting placeholders when no program was passed left user-input weights unprotected for the callers that omit it. Vulkan, Samsung and XNNPACK all have the owning program in scope, so pass it (ConvertToLinearPass becomes an XNNPACKPass, which the pass manager already hands the program), and make the no-program case reject placeholders instead. --- backends/samsung/_passes/enn_pass_manager.py | 2 +- backends/transforms/addmm_mm_to_linear.py | 25 ++++++--- .../test/test_addmm_mm_to_linear.py | 54 +++++++++++++++++++ backends/vulkan/vulkan_preprocess.py | 4 +- backends/xnnpack/_passes/convert_to_linear.py | 8 +-- 5 files changed, 81 insertions(+), 12 deletions(-) diff --git a/backends/samsung/_passes/enn_pass_manager.py b/backends/samsung/_passes/enn_pass_manager.py index 9c143f716b8..b5512f207c2 100644 --- a/backends/samsung/_passes/enn_pass_manager.py +++ b/backends/samsung/_passes/enn_pass_manager.py @@ -72,7 +72,7 @@ def transform_for_preprocess_pass(self, exported_program: ExportedProgram): self.add_pass(FoldQDQPass()) self.add_pass(Conv1dToConv2d(exported_program)) self.add_pass(FuseBatchNormWithConvPass(exported_program)) - self.add_pass(AddmmToLinearTransform()) + self.add_pass(AddmmToLinearTransform(exported_program)) self.add_pass(ReplaceOpsWithScalar()) self.add_pass(RemoveGetItemPass()) self.add_pass(InsertQDQPass(exported_program)) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index f6ae4cfae5e..97a12784cc4 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -34,6 +34,13 @@ def get_shape(input_node: torch.fx.Node): return upper_bound_shape +def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) -> bool: + """Whether `node` is a buffer the program writes back to.""" + signature = exported_program.graph_signature + name = signature.inputs_to_buffers.get(node.name) + return name is not None and name in set(signature.buffers_to_mutate.values()) + + def is_constant_tensor( node: torch.fx.Node, exported_program: Optional[ExportedProgram] ) -> bool: @@ -51,13 +58,19 @@ def is_constant_tensor( return False if exported_program is None: # Without the owning program a lifted parameter cannot be told apart - # from a user input. Placeholders were always rewritten before, so keep - # accepting them rather than regressing callers that pass no program. + # from a user input, and prepacking a user input would bake in whatever + # the first call passed. Every in-tree caller threads the program, so + # take the safe side and leave the matmul alone. + return False + if is_param(exported_program, node) or is_lifted_tensor_constant( + exported_program, node + ): return True - return ( - is_param(exported_program, node) - or is_buffer(exported_program, node) - or is_lifted_tensor_constant(exported_program, node) + # A buffer the program mutates is state, not a build-time constant: + # prepacking it would freeze the value the delegate was built with and go + # on returning it after the buffer has moved on. + return is_buffer(exported_program, node) and not _is_mutated_buffer( + node, exported_program ) diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index 371c298061b..08d1c1ef159 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -75,6 +75,60 @@ def forward(self, x, w): graph = AddmmToLinearTransform(program)(program.graph_module).graph_module.graph self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + def test_mutated_buffer_weight_is_not_rewritten(self): + # The buffer is a weight by shape but not by lifetime: forward() writes + # to it, so a delegate that prepacked it while building its graph would + # keep returning the value it was built with. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.randn(4, 8)) + + def forward(self, x): + # Read first, then write: functionalization leaves the matmul + # reading the placeholder itself, which is the case the pass + # has to reject. + y = torch.mm(x, self.w.t()) + self.w.add_(1.0) + return y + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) + + def test_frozen_buffer_weight_is_rewritten(self): + # The mirror of the case above: a buffer nothing writes to is as + # constant as a parameter, and has to keep being rewritten. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.randn(4, 8)) + + def forward(self, x): + return torch.mm(x, self.w.t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + + def test_no_program_leaves_placeholders_alone(self): + # Without the owning program a lifted parameter and a user input look + # the same, so nothing is rewritten. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc = torch.nn.Linear(8, 4) + + def forward(self, x): + return self.fc(x) + + edge = to_edge( + torch.export.export(Model().eval(), (torch.randn(2, 8),), strict=True) + ) + program = edge.exported_program() + graph = AddmmToLinearTransform()(program.graph_module).graph_module.graph + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + if __name__ == "__main__": unittest.main() diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index f7d6955ce26..d56405e4ed0 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -185,9 +185,9 @@ def preprocess( # noqa: C901 program = apply_passes( program, [ - AddmmToLinearTransform(), + AddmmToLinearTransform(program), FuseBatchNormPass(program), - AddmmToLinearTransform(), + AddmmToLinearTransform(program), InsertDtypePromotionPass(), FusePatternsPass(), FuseClampPass(), diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 2cef71bf927..33fd3e76da7 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/backends/xnnpack/_passes/convert_to_linear.py @@ -13,8 +13,8 @@ from executorch.backends.transforms.addmm_mm_to_linear import ( apply_addmm_mm_to_linear_transform, ) +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass from torch.fx.passes.infra.pass_base import PassResult from torch.fx.passes.utils.source_matcher_utils import ( @@ -26,7 +26,7 @@ logger.setLevel(logging.WARNING) -class ConvertToLinearPass(ExportPass): +class ConvertToLinearPass(XNNPACKPass): linear_modules = [ torch.nn.Linear, torch.nn.functional.linear, @@ -185,7 +185,9 @@ def call(self, graph_module: torch.fx.GraphModule): break # fall back to linear transform - graph_module.graph = apply_addmm_mm_to_linear_transform(graph_module.graph) + graph_module.graph = apply_addmm_mm_to_linear_transform( + graph_module.graph, self.exported_program + ) graph_module.recompile() From d63c99b86f125b9c72e4c3454d023b783fe1e54f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 15 Sep 2026 16:29:58 +0200 Subject: [PATCH 03/10] Treat a dequantized constant as a constant weight A quantized weight reaches mm/addmm through a dequantize node, which is a call_function, so is_constant_tensor read it as a runtime value and the rewrite to linear was skipped. That silently disables downstream quantized-matmul fusion, which matches on linear. Exporting all-mpnet-base-v2 int8 for Vulkan produced 72 dequantize_per_channel plus 72 float addmm instead of 72 _weight_int8pack_mm, dequantizing the weights at runtime and aborting at load because the Vulkan runtime has no dequantize_per_channel kernel. Accept a dequantize node when everything it dequantizes is itself constant. Both mpnet int8 Vulkan exports are then byte-identical to the published artifacts again. --- backends/transforms/addmm_mm_to_linear.py | 20 +++++++++++++++ .../test/test_addmm_mm_to_linear.py | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index 97a12784cc4..f4c3c67f769 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -41,6 +41,19 @@ def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) - return name is not None and name in set(signature.buffers_to_mutate.values()) +# Dequantized constants are still known at build time. Quantized weights +# reach mm/addmm through one of these, and backends rely on the rewrite to +# `linear` to fuse the pair back into a quantized matmul, so they have to +# count as constant here. +_DEQUANT_TARGETS = ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + exir_ops.edge.quantized_decomposed.dequantize_per_token.default, + exir_ops.edge.torchao.dequantize_affine.default, +) + + def is_constant_tensor( node: torch.fx.Node, exported_program: Optional[ExportedProgram] ) -> bool: @@ -54,6 +67,13 @@ def is_constant_tensor( """ if node.op == "get_attr": return True + if node.op == "call_function" and node.target in _DEQUANT_TARGETS: + # Constant only if everything it dequantizes is itself constant. + return all( + is_constant_tensor(arg, exported_program) + for arg in node.args + if isinstance(arg, torch.fx.Node) + ) if node.op != "placeholder": return False if exported_program is None: diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index 08d1c1ef159..8c16ec229b0 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -37,6 +37,31 @@ def forward(self, x): self.assertEqual(count_targets(graph, exir_ops.edge.aten.addmm.default), 0) self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + def test_dequantized_constant_weight_is_rewritten_to_linear(self): + # A quantized weight reaches addmm through a dequantize node, so the + # operand is a call_function rather than a placeholder. It is still + # known at build time, and backends fuse the dequantize+linear pair + # back into a quantized matmul, so the rewrite must still happen. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "w", torch.randint(-128, 127, (4, 8), dtype=torch.int8) + ) + self.register_buffer("scale", torch.rand(4) + 0.1) + self.register_buffer("zero", torch.zeros(4, dtype=torch.int8)) + self.register_buffer("bias", torch.randn(4)) + + def forward(self, x): + w = torch.ops.quantized_decomposed.dequantize_per_channel.default( + self.w, self.scale, self.zero, 0, -128, 127, torch.int8 + ) + return torch.addmm(self.bias, x, w.t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.addmm.default), 0) + def test_computed_weight_stays_a_matmul(self): # `w` is produced at runtime, so the transposed matmul below is not a # linear: backends prepack a linear's weight while building their graph From 0b768edfb41e43dea254cbaa8e69a29055c84721 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Wed, 16 Sep 2026 08:46:27 +0200 Subject: [PATCH 04/10] Give the LLM builder's ConvertToLinearPass its exported program ConvertToLinearPass became an XNNPACKPass in this stack, so its constructor now takes the exported program. This call site was still constructing it with no arguments and raised TypeError: XNNPACKPass.__init__() missing 1 required positional argument: 'exported_program' Every llama-family export goes through here, which is why it took out the llama runner, lora, qnn and phi_4_mini jobs along with the export_llama_lib unit tests. --- extension/llm/export/builder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index 6ec9e320eb8..ba19bc0dd96 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -555,7 +555,9 @@ def to_executorch( # 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.edge_manager.transform( + [ConvertToLinearPass(self.edge_manager.exported_program())] + ) self.export_program = self.edge_manager.to_executorch( ExecutorchBackendConfig( From ff2480d8bdad64d042f0f05574256fd464bc1b6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Wed, 16 Sep 2026 18:54:49 +0200 Subject: [PATCH 05/10] Bind ConvertToLinearPass to each method's own program EdgeProgramManager.transform() runs its passes over every method, so passing a single program bound one pass to every graph, and exported_program() with no argument looks up "forward". A multimethod export has no such method, so test-lora-multimethod failed with KeyError: 'forward' before any rewriting happened. Use the per-method dict form of transform() so each pass gets the program of the graph it rewrites. --- .../test/passes/test_convert_to_linear.py | 44 +++++++++++++++++++ extension/llm/export/builder.py | 8 +++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/backends/xnnpack/test/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index 0e7bc7d01c4..a7bf75a2af2 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/backends/xnnpack/test/passes/test_convert_to_linear.py @@ -7,8 +7,10 @@ import unittest import torch +from executorch import exir from executorch.backends.xnnpack._passes.convert_to_linear import ConvertToLinearPass from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir.dialects._ops import ops as exir_ops class TestConvertToLinear(unittest.TestCase): @@ -41,3 +43,45 @@ def test_fp32_convert_to_linear(self): ) .run_method_and_compare_outputs() ) + + def test_multimethod_program_without_a_forward_method(self): + # A multimethod export has no "forward". The pass needs the exported + # program of the method it is rewriting, so binding one program for + # every method, or letting exported_program() default to "forward", + # raises KeyError before any rewriting happens. + class Linear(torch.nn.Module): + def __init__(self, out_features: int): + super().__init__() + self.fc = torch.nn.Linear(8, out_features) + + def forward(self, x): + return self.fc(x) + + inputs = (torch.randn(2, 8),) + programs = { + name: torch.export.export(Linear(out).eval(), inputs, strict=True) + for name, out in (("base_forward", 4), ("lora_forward", 6)) + } + # aten.linear is not a core ATen op, so the verifier has to be told to + # accept what this pass deliberately produces. + config = exir.EdgeCompileConfig( + _core_aten_ops_exception_list=[torch.ops.aten.linear.default] + ) + edge = exir.to_edge(programs, compile_config=config) + + with self.assertRaises(KeyError): + edge.exported_program() + + edge = edge.transform( + { + name: [ConvertToLinearPass(edge.exported_program(name))] + for name in edge.methods + }, + compile_config=config, + ) + + for name in ("base_forward", "lora_forward"): + graph = edge.exported_program(name).graph + targets = [n.target for n in graph.nodes if n.op == "call_function"] + self.assertIn(exir_ops.edge.aten.linear.default, targets) + self.assertNotIn(exir_ops.edge.aten.addmm.default, targets) diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index ba19bc0dd96..ba88be56611 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -555,8 +555,14 @@ def to_executorch( # TODO: ConvertToLinearPass is not a sound pass and must be called before # const propagation. It requires fixing: # https://github.com/pytorch/executorch/issues/10499 + # transform() runs the passes over every method, so each method needs a + # pass bound to its own program. Defaulting to exported_program() would + # look up "forward", which a multimethod export does not have. self.edge_manager.transform( - [ConvertToLinearPass(self.edge_manager.exported_program())] + { + name: [ConvertToLinearPass(self.edge_manager.exported_program(name))] + for name in self.edge_manager.methods + } ) self.export_program = self.edge_manager.to_executorch( From 5f2d4706bdb6abf56f90d720f9d008077f88f628 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 12 Oct 2026 15:09:26 +0200 Subject: [PATCH 06/10] Count groupwise-dequantized constants as constant weights dequantize_per_channel_group (4-bit groupwise) was missing from the dequant allowlist, so those mm/addmm weights stopped being rewritten. --- backends/transforms/addmm_mm_to_linear.py | 1 + .../test/test_addmm_mm_to_linear.py | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index f4c3c67f769..79c47402549 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -49,6 +49,7 @@ def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default, exir_ops.edge.quantized_decomposed.dequantize_per_token.default, exir_ops.edge.torchao.dequantize_affine.default, ) diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index 8c16ec229b0..d83867c5cf9 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -62,6 +62,28 @@ def forward(self, x): self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) self.assertEqual(count_targets(graph, exir_ops.edge.aten.addmm.default), 0) + def test_groupwise_dequantized_constant_weight_is_rewritten_to_linear(self): + # 4-bit groupwise (8da4w-style) weights are dequantized per channel + # group, and have to be rewritten like any other dequantized constant. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "w", torch.randint(-8, 7, (4, 8), dtype=torch.int8) + ) + self.register_buffer("scale", torch.rand(4, 2) + 0.1) + self.register_buffer("zero", torch.zeros(4, 2, dtype=torch.int8)) + + def forward(self, x): + w = torch.ops.quantized_decomposed.dequantize_per_channel_group.default( + self.w, self.scale, self.zero, -8, 7, torch.int8, 4, torch.float32 + ) + return torch.mm(x, w.t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + def test_computed_weight_stays_a_matmul(self): # `w` is produced at runtime, so the transposed matmul below is not a # linear: backends prepack a linear's weight while building their graph From 983592c0932effaed8cb0ce3267b64c2e82a777d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 12 Oct 2026 15:10:17 +0200 Subject: [PATCH 07/10] Follow quantize ops when checking for a constant weight An unfolded weight reaches mm as dq(q(param)); a quantize node whose inputs are all constant is constant too. --- backends/transforms/addmm_mm_to_linear.py | 15 +++++-- .../test/test_addmm_mm_to_linear.py | 40 +++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index 79c47402549..1038b673420 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -44,14 +44,21 @@ def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) - # Dequantized constants are still known at build time. Quantized weights # reach mm/addmm through one of these, and backends rely on the rewrite to # `linear` to fuse the pair back into a quantized matmul, so they have to -# count as constant here. -_DEQUANT_TARGETS = ( +# count as constant here. Quantize ops are listed too because a weight that +# was never constant-folded reaches the matmul as dq(q(param)). +_QDQ_TARGETS = ( exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default, exir_ops.edge.quantized_decomposed.dequantize_per_token.default, exir_ops.edge.torchao.dequantize_affine.default, + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor, + exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + exir_ops.edge.quantized_decomposed.quantize_per_channel_group.default, + exir_ops.edge.quantized_decomposed.quantize_per_token.default, + exir_ops.edge.torchao.quantize_affine.default, ) @@ -68,8 +75,8 @@ def is_constant_tensor( """ if node.op == "get_attr": return True - if node.op == "call_function" and node.target in _DEQUANT_TARGETS: - # Constant only if everything it dequantizes is itself constant. + if node.op == "call_function" and node.target in _QDQ_TARGETS: + # Constant only if every tensor it (de)quantizes is itself constant. return all( is_constant_tensor(arg, exported_program) for arg in node.args diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index d83867c5cf9..49427f4fbcd 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -84,6 +84,46 @@ def forward(self, x): self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + def test_unfolded_quantized_constant_weight_is_rewritten_to_linear(self): + # Without constant folding the weight reaches mm as dq(q(param)). The + # quantize node only reads constants, so the operand is still constant. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(4, 8)) + self.register_buffer("scale", torch.rand(4) + 0.1) + self.register_buffer("zero", torch.zeros(4, dtype=torch.int64)) + + def forward(self, x): + q = torch.ops.quantized_decomposed.quantize_per_channel.default( + self.w, self.scale, self.zero, 0, -128, 127, torch.int8 + ) + w = torch.ops.quantized_decomposed.dequantize_per_channel.default( + q, self.scale, self.zero, 0, -128, 127, torch.int8 + ) + return torch.mm(x, w.t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8),)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 1) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 0) + + def test_quantized_computed_weight_stays_a_matmul(self): + # The quantize node reads a value produced at runtime, so dq(q(.)) is + # runtime data too. + class Model(torch.nn.Module): + def forward(self, x, w): + q = torch.ops.quantized_decomposed.quantize_per_tensor.default( + w * 2.0, 0.1, 0, -128, 127, torch.int8 + ) + w = torch.ops.quantized_decomposed.dequantize_per_tensor.default( + q, 0.1, 0, -128, 127, torch.int8 + ) + return torch.mm(x, w.t()) + + graph = self._transform(Model().eval(), (torch.randn(2, 8), torch.randn(4, 8))) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) + def test_computed_weight_stays_a_matmul(self): # `w` is produced at runtime, so the transposed matmul below is not a # linear: backends prepack a linear's weight while building their graph From 552860262361eb9fb6ba68c8b5dc7c401d0c16e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 12 Oct 2026 15:10:43 +0200 Subject: [PATCH 08/10] Warn once when a placeholder weight is skipped for lack of a program Without the ExportedProgram the pass leaves parameter weights alone; log it so out-of-tree callers notice the missing rewrite. --- backends/transforms/addmm_mm_to_linear.py | 14 ++++++++++++++ .../transforms/test/test_addmm_mm_to_linear.py | 6 +++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index 1038b673420..bdc07b12c92 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import functools +import logging from typing import Optional import torch @@ -15,6 +17,8 @@ _int64_max_dim_val = torch.iinfo(torch.int64).max - 1 +logger = logging.getLogger(__name__) + def get_shape(input_node: torch.fx.Node): """ @@ -62,6 +66,15 @@ def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) - ) +@functools.cache +def _warn_no_exported_program() -> None: + logger.warning( + "addmm/mm to linear: no ExportedProgram was passed, so matmuls whose " + "weight is a placeholder (including parameters) are not rewritten to " + "linear. Pass the owning program to enable the rewrite." + ) + + def is_constant_tensor( node: torch.fx.Node, exported_program: Optional[ExportedProgram] ) -> bool: @@ -89,6 +102,7 @@ def is_constant_tensor( # from a user input, and prepacking a user input would bake in whatever # the first call passed. Every in-tree caller threads the program, so # take the safe side and leave the matmul alone. + _warn_no_exported_program() return False if is_param(exported_program, node) or is_lifted_tensor_constant( exported_program, node diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index 49427f4fbcd..cc545f8b758 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -7,6 +7,7 @@ import unittest import torch +from executorch.backends.transforms import addmm_mm_to_linear from executorch.backends.transforms.addmm_mm_to_linear import AddmmToLinearTransform from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops @@ -213,7 +214,10 @@ def forward(self, x): torch.export.export(Model().eval(), (torch.randn(2, 8),), strict=True) ) program = edge.exported_program() - graph = AddmmToLinearTransform()(program.graph_module).graph_module.graph + # The skip is silent otherwise, so it is logged (once per process). + addmm_mm_to_linear._warn_no_exported_program.cache_clear() + with self.assertLogs(addmm_mm_to_linear.logger, level="WARNING"): + graph = AddmmToLinearTransform()(program.graph_module).graph_module.graph self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) From f32eb8c87179b67e7ff41948592a712b21adce7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 12 Oct 2026 15:11:18 +0200 Subject: [PATCH 09/10] Tidy the addmm/mm to linear tests Rename the computed-weight addmm test, cover a dequant of runtime data, and drop an assertion that only tested EdgeProgramManager's default method. --- .../test/test_addmm_mm_to_linear.py | 22 ++++++++++++++++++- .../test/passes/test_convert_to_linear.py | 9 ++------ 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/backends/transforms/test/test_addmm_mm_to_linear.py b/backends/transforms/test/test_addmm_mm_to_linear.py index cc545f8b758..272d8de89b3 100644 --- a/backends/transforms/test/test_addmm_mm_to_linear.py +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -125,6 +125,26 @@ def forward(self, x, w): self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) + def test_dequantized_runtime_input_stays_a_matmul(self): + # A dequantize is only constant when what it reads is: here it reads a + # value produced at runtime, so the matmul must stay a matmul. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("scale", torch.rand(4) + 0.1) + self.register_buffer("zero", torch.zeros(4, dtype=torch.int8)) + + def forward(self, x, w_int): + w = torch.ops.quantized_decomposed.dequantize_per_channel.default( + w_int + 1, self.scale, self.zero, 0, -128, 127, torch.int8 + ) + return torch.mm(x, w.t()) + + w_int = torch.randint(-128, 126, (4, 8), dtype=torch.int8) + graph = self._transform(Model().eval(), (torch.randn(2, 8), w_int)) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) + self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) + def test_computed_weight_stays_a_matmul(self): # `w` is produced at runtime, so the transposed matmul below is not a # linear: backends prepack a linear's weight while building their graph @@ -137,7 +157,7 @@ def forward(self, x, w): self.assertEqual(count_targets(graph, exir_ops.edge.aten.linear.default), 0) self.assertEqual(count_targets(graph, exir_ops.edge.aten.mm.default), 1) - def test_computed_bias_operand_stays_an_addmm(self): + def test_computed_addmm_weight_stays_an_addmm(self): class Model(torch.nn.Module): def forward(self, x, w, b): return torch.addmm(b, x, (w * 2.0).t()) diff --git a/backends/xnnpack/test/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index a7bf75a2af2..c70a0ba4212 100644 --- a/backends/xnnpack/test/passes/test_convert_to_linear.py +++ b/backends/xnnpack/test/passes/test_convert_to_linear.py @@ -45,10 +45,8 @@ def test_fp32_convert_to_linear(self): ) def test_multimethod_program_without_a_forward_method(self): - # A multimethod export has no "forward". The pass needs the exported - # program of the method it is rewriting, so binding one program for - # every method, or letting exported_program() default to "forward", - # raises KeyError before any rewriting happens. + # A multimethod export has no "forward". Each method's pass has to be + # bound to that method's own exported program. class Linear(torch.nn.Module): def __init__(self, out_features: int): super().__init__() @@ -69,9 +67,6 @@ def forward(self, x): ) edge = exir.to_edge(programs, compile_config=config) - with self.assertRaises(KeyError): - edge.exported_program() - edge = edge.transform( { name: [ConvertToLinearPass(edge.exported_program(name))] From eb610e43a2e2c229d45aa27aedc258b29890455b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 12 Oct 2026 15:11:29 +0200 Subject: [PATCH 10/10] Check mutated buffers without building a set --- backends/transforms/addmm_mm_to_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/transforms/addmm_mm_to_linear.py b/backends/transforms/addmm_mm_to_linear.py index bdc07b12c92..823229f645b 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -42,7 +42,7 @@ def _is_mutated_buffer(node: torch.fx.Node, exported_program: ExportedProgram) - """Whether `node` is a buffer the program writes back to.""" signature = exported_program.graph_signature name = signature.inputs_to_buffers.get(node.name) - return name is not None and name in set(signature.buffers_to_mutate.values()) + return name is not None and name in signature.buffers_to_mutate.values() # Dequantized constants are still known at build time. Quantized weights