Only rewrite addmm/mm to linear when the weight is a constant

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))

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)