Back out "Add baddbmm operator support" (#26)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/26

Original commit changeset: efd417b2b494

Original Phabricator Diff: D34743279 (https://github.com/pytorch/fx2trt/commit/4f75c6f42213c841880de5358d32a9cee09beb4f)

Reviewed By: frank-wei

Differential Revision: D35002287

fbshipit-source-id: 7929aae49ef1daa0e116ba930f25a40c9b1388aa

Author: Shirong Wu

test/converters/acc_op/test_mm_ops.py

@@ -1314,49 +1314,6 @@ def test_bmm(self):
             acc_ops.matmul, lambda x: torch.bmm(x, x), input_shape=(2, 4, 4)
         )
 
-    def test_baddbmm_with_alpha_beta(self):
-        class TestModule(torch.nn.Module):
-            def forward(self, input: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-                return torch.baddbmm(input, a, b, alpha = 1.2, beta = 1.1)
-
-        m = TestModule()
-        input, a, b = torch.randn(10, 2, 3), torch.randn(10, 2, 4), torch.randn(10, 4, 3)
-        traced = acc_tracer.trace(m, [input, a, b])
-
-        ph_in = ph_a = ph_b = mm = add = mm_mul = add_mul = None
-        for node in traced.graph.nodes:
-            if node.op == "placeholder":
-                if str(node.target) == "a":
-                    ph_a = node
-                elif str(node.target) == "b":
-                    ph_b = node
-                else:
-                    self.assertTrue(str(node.target) == "input")
-                    ph_in = node
-            elif node.op == "call_function":
-                if node.target == acc_ops.matmul:
-                    self.assertEqual(node.kwargs["input"], ph_a)
-                    self.assertEqual(node.kwargs["other"], ph_b)
-                    mm = node
-                elif node.target == acc_ops.add:
-                    self.assertEqual(node.kwargs["input"], mm_mul)
-                    self.assertEqual(node.kwargs["other"], add_mul)
-                    add = node
-                elif mm_mul:
-                    self.assertEqual(node.kwargs["input"], ph_in)
-                    self.assertEqual(node.kwargs["other"], 1.1)
-                    add_mul = node
-
-                else:
-                    self.assertEqual(node.kwargs["input"], mm)
-                    self.assertEqual(node.kwargs["other"], 1.2)
-                    mm_mul = node
-            elif node.op == "output":
-                self.assertEqual(add, node.args[0])
-            else:
-                self.fail(f"Unexpected node: {node.format_node()}")
-        torch.testing.assert_allclose(m(input, a, b), traced(input, a, b))
-
     def test_tile(self):
         return self._make_acc_op_function_test(
             acc_ops.tile, lambda x: torch.tile(x, (2, 1, 2)), input_shape=(1, 2)
@@ -1759,7 +1716,7 @@ def forward(
         input, a, b = torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2)
         traced = acc_tracer.trace(m, [input, a, b])
 
-        ph_in = ph_a = ph_b = mm = add = mm_mul = add_mul = reshape = None
+        ph_in = ph_a = ph_b = mm = add = mm_mul = add_mul = None
         for node in traced.graph.nodes:
             if node.op == "placeholder":
                 if str(node.target) == "a":

test/tracer/test_acc_tracer.py

@@ -361,60 +361,6 @@ def softmax(*, input, dim, dtype):
     return torch.nn.functional.softmax(input=input, dim=dim, dtype=dtype)
 
 
-@register_custom_acc_mapper_fn(
-    op_and_target=("call_function", torch.baddbmm),
-    arg_replacement_tuples=[
-        ("input", "input"),
-        ("batch1", "mat1"),
-        ("batch2", "mat2"),
-        ("beta", "beta"),
-        ("alpha", "alpha"),
-    ],
-)
-def baddbmm_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
-    """
-    Mapping from torch.addmm to acc_ops.mm -> acc_ops.add, if alpha or beta is not 1
-    then we also insert acc_ops.mul to the right place.
-    """
-    with node.graph.inserting_before(node):
-        mm_kwargs = {"input": node.kwargs["mat1"], "other": node.kwargs["mat2"]}
-        mm_node = node.graph.create_node(
-            "call_function", matmul, kwargs=mm_kwargs, name=f"{node.name}_mm"
-        )
-        mm_node.meta = node.meta.copy()
-
-        if node.kwargs["alpha"] != 1:
-            mul_kwargs = {"input": mm_node, "other": node.kwargs["alpha"]}
-            mm_node = node.graph.create_node(
-                "call_function", mul, kwargs=mul_kwargs, name=f"{mm_node.name}_mul"
-            )
-        mm_node.meta = node.meta.copy()
-
-        input_node = node.kwargs["input"]
-        if node.kwargs["beta"] != 1:
-            mul_kwargs = {"input": input_node, "other": node.kwargs["beta"]}
-            new_input_node = node.graph.create_node(
-                "call_function", mul, kwargs=mul_kwargs, name=f"{node.name}_input_mul"
-            )
-            assert isinstance(input_node, torch.fx.Node)
-            new_input_node.meta = input_node.meta.copy()
-            input_node = new_input_node
-
-        # broadcast input to target shape
-        if input_node.meta["tensor_rank"] < mm_node.meta["tensor_rank"]:
-            rank = input_node.meta["tensor_rank"]
-            raise RuntimeError(
-                f"Unable to broadcast input with dimension {rank} on batch size dimension. "
-            )
-
-        add_kwargs = {"input": mm_node, "other": input_node}
-        add_node = node.graph.create_node(
-            "call_function", add, kwargs=add_kwargs, name=f"{node.name}_add"
-        )
-        add_node.meta = node.meta.copy()
-        return add_node
-
-
 @register_custom_acc_mapper_fn(
     op_and_target=("call_function", torch.addmm),
     arg_replacement_tuples=[
@@ -454,13 +400,6 @@ def addmm_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
             new_input_node.meta = input_node.meta.copy()
             input_node = new_input_node
 
-        # broadcast input to target shape
-        if input_node.meta["tensor_rank"] < mm_node.meta["tensor_rank"]:
-            rank = input_node.meta["tensor_rank"]
-            raise RuntimeError(
-                f"Unable to broadcast input with dimension {rank} on batch size dimension. "
-            )
-
         add_kwargs = {"input": mm_node, "other": input_node}
         add_node = node.graph.create_node(
             "call_function", add, kwargs=add_kwargs, name=f"{node.name}_add"