From 835b2b3642fc8635516e9015fa22b09844f4492c 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] 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()