diff --git a/backends/samsung/_passes/enn_pass_manager.py b/backends/samsung/_passes/enn_pass_manager.py index 874aa7115c9..e52f66ea986 100644 --- a/backends/samsung/_passes/enn_pass_manager.py +++ b/backends/samsung/_passes/enn_pass_manager.py @@ -74,7 +74,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 358cbb7ac14..823229f645b 100644 --- a/backends/transforms/addmm_mm_to_linear.py +++ b/backends/transforms/addmm_mm_to_linear.py @@ -4,15 +4,21 @@ # 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 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 +logger = logging.getLogger(__name__) + def get_shape(input_node: torch.fx.Node): """ @@ -32,6 +38,84 @@ 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 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. 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, +) + + +@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: + """ + 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 == "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 + if isinstance(arg, torch.fx.Node) + ) + 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, 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 + ): + return True + # 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 + ) + + def get_dqlinear_input(node: torch.fx.Node): ops = exir_ops.edge node_to_backtrack = node @@ -99,7 +183,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 +211,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 +231,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 +250,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..272d8de89b3 --- /dev/null +++ b/backends/transforms/test/test_addmm_mm_to_linear.py @@ -0,0 +1,245 @@ +# 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 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 + + +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_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_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_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_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 + # 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_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()) + + 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) + + 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() + # 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) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index 520c18441a9..e8a6de1e306 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -189,9 +189,9 @@ def preprocess( # noqa: C901 program = apply_passes( program, [ - AddmmToLinearTransform(), + AddmmToLinearTransform(program), FuseBatchNormPass(program), - AddmmToLinearTransform(), + AddmmToLinearTransform(program), InsertDtypePromotionPass(), FusePatternsPass(), Conv1dAsConv2dPass(), diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 8259aae4e98..e6140b4b555 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/backends/xnnpack/_passes/convert_to_linear.py @@ -13,9 +13,9 @@ 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.backends.xnnpack.utils.quant_utils import is_dequant 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 ( @@ -27,7 +27,7 @@ logger.setLevel(logging.WARNING) -class ConvertToLinearPass(ExportPass): +class ConvertToLinearPass(XNNPACKPass): linear_modules = [ torch.nn.Linear, torch.nn.functional.linear, @@ -189,7 +189,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() diff --git a/backends/xnnpack/test/passes/test_convert_to_linear.py b/backends/xnnpack/test/passes/test_convert_to_linear.py index 0e7bc7d01c4..c70a0ba4212 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,40 @@ 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". 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__() + 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) + + 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 6ec9e320eb8..ba88be56611 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -555,7 +555,15 @@ 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()]) + # 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( + { + name: [ConvertToLinearPass(self.edge_manager.exported_program(name))] + for name in self.edge_manager.methods + } + ) self.export_program = self.edge_manager.to_executorch( ExecutorchBackendConfig(