Fix memory planning error

Differential Revision: D88770975

Pull Request resolved: pytorch#16165

Author: DrJessop

backends/cadence/aot/replace_ops.py

@@ -15,13 +15,12 @@
 import math
 import operator
 from operator import neg
-from typing import cast, Dict, Iterable, Optional, Sequence
+from typing import cast, Dict, Optional, Sequence
 
 import torch
 import torch.fx
 from executorch.backends.cadence.aot.compiler_utils import (
     get_zero_point,
-    is_node_with_op,
     quantize_tensor_multiplier,
 )
 from executorch.backends.cadence.aot.fuse_ops import (
@@ -52,17 +51,6 @@
 }
 
 
-def contains_placeholder_or_param(nodes: Iterable[torch.fx.Node]) -> bool:
-    """
-    Return true if any of the node in the incoming nodes list is a placeholder
-    or parameter
-    """
-    return any(
-        is_node_with_op(node, "placeholder") or is_node_with_op(node, "get_attr")
-        for node in nodes
-    )
-
-
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
 class ReplaceLogicalNotBooleanWhereWithWherePass(RemoveOrReplacePassInterface):
     """
@@ -367,6 +355,9 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
 
         graph = node.graph
 
+        fit_bias = beta == 1.0
+        fit_mat2 = False
+
         # Handle transpose: if mat2 is a transpose op, extract the original tensor
         transposed_mat2 = False
         if (
@@ -376,6 +367,10 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
             # mat2 is already transposed, so we use the input to the transpose
             mat2 = cast(torch.fx.Node, mat2.args[0])
             transposed_mat2 = True
+            fit_mat2 = alpha == 1.0
+
+        if not (fit_bias and fit_mat2):
+            return False
 
         # Multiply bias by beta if needed
         if beta != 1.0:

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -1007,9 +1007,10 @@ def test_replace_linear_with_fully_connected(self) -> None:
             0,
         )
 
-    @expand([[1.0, 1.0], [2.0, 3.0]])
     @torch.no_grad()
-    def test_replace_addmm_with_linear(self, alpha: float, beta: float) -> None:
+    def test_replace_addmm_with_linear(self) -> None:
+        alpha = 1.0
+        beta = 1.0
         M, K, N = 14, 12, 10
         builder = GraphBuilder()
         x_input = torch.randn(N, dtype=torch.float32)
@@ -1054,6 +1055,37 @@ def test_replace_addmm_with_linear(self, alpha: float, beta: float) -> None:
             0,
         )
 
+    @torch.no_grad()
+    def test_no_replace_addmm_with_linear(self) -> None:
+        alpha = 2.0
+        beta = 3.0
+        M, K, N = 14, 12, 10
+        builder = GraphBuilder()
+        x_input = torch.randn(N, dtype=torch.float32)
+        y_input = torch.randn([M, K], dtype=torch.float32)
+        z_input = torch.randn([N, K], dtype=torch.float32)
+        x = builder.placeholder("x", x_input)
+        y = builder.placeholder("y", y_input)
+        z = builder.placeholder("z", z_input)
+        permute_copy = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default,
+            args=(z, [1, 0]),
+        )
+        addmm = builder.call_operator(
+            op=exir_ops.edge.aten.addmm.default,
+            args=(x, y, permute_copy),
+            kwargs={"beta": beta, "alpha": alpha},
+        )
+        builder.output([addmm])
+        original_gm = builder.get_graph_module()
+
+        gm = cast(
+            PassResult, ReplacePermuteWithTransposePass()(original_gm)
+        ).graph_module
+
+        pass_result = cast(PassResult, ReplaceAddMMWithLinearPass()(gm))
+        self.assertFalse(pass_result.modified)
+
     @torch.no_grad()
     def test_replace_mm_with_addmm(self) -> None:
         M, K, N = 14, 48, 24