Add baddbmm operator support (#14)

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

Add support for torch.baddbmm. Add unit test to cover bmm and addmm, both are mapping to other converter insteand of having new acc_op converter.

Reviewed By: yinghai

Differential Revision: D34743279

fbshipit-source-id: efd417b2b494635c63f2ec58daa3fc568f72111a

Author: Shirong Wu

test/converters/acc_op/test_mm_ops.py

@@ -0,0 +1,66 @@
+from torch.testing._internal.common_fx2trt import AccTestCase
+from parameterized import parameterized
+from torch import nn
+import torch
+from fx2trt_oss.tracer.acc_tracer import acc_ops
+
+
+
+class TestBmmConverters(AccTestCase):
+    @parameterized.expand(
+        [
+            ("1", (10, 2, 13), (10, 13, 2)),
+        ]
+    )
+    def test_bmm(self, _, input_shape, other_shape, alpha=1, beta=1):
+        class Bmm(nn.Module):
+            def forward(self, input, other):
+                return torch.bmm(input, other)
+
+        inputs = [torch.randn(*input_shape), torch.randn(*other_shape)]
+
+        self.run_test(
+            Bmm(),
+            inputs,
+            expected_ops={acc_ops.matmul}
+        )
+
+    @parameterized.expand(
+        [
+            ("default", (2, 3), (2, 3), (3, 3)),
+            ("broadcast", (1, 1), (10, 7), (7, 13)),
+        ]
+    )
+    def test_addmm(self, _, input_shape, m1_shape, m2_shape, alpha=1, beta=1):
+        class Addmm(nn.Module):
+            def forward(self, input, m1, m2):
+                return torch.addmm(input, m1, m2, alpha=alpha, beta=beta)
+
+        inputs = [torch.randn(input_shape), torch.randn(*m1_shape), torch.randn(*m2_shape)]
+        test_implicit_batch_dim = len(input_shape) > 2
+
+        self.run_test(
+            Addmm(),
+            inputs,
+            expected_ops={acc_ops.matmul, acc_ops.add},
+            test_implicit_batch_dim = test_implicit_batch_dim,
+        )
+
+    @parameterized.expand(
+        [
+            ("default", (10, 2, 2), (10, 2, 3), (10, 3, 2)),
+            ("broadcast", (10, 2, 1), (10, 2, 3), (10, 3, 2)),
+        ]
+    )
+    def test_baddbmm(self, _, input_shape, m1_shape, m2_shape, alpha=1, beta=1):
+        class Baddbmm(nn.Module):
+            def forward(self, input, m1, m2):
+                return torch.baddbmm(input, m1, m2, alpha=alpha, beta=beta)
+
+        inputs = [torch.randn(*input_shape), torch.randn(*m1_shape), torch.randn(*m2_shape)]
+
+        self.run_test(
+            Baddbmm(),
+            inputs,
+            expected_ops={acc_ops.matmul, acc_ops.add}
+        )

test/tracer/test_acc_tracer.py

@@ -1314,6 +1314,49 @@ 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)
@@ -1716,7 +1759,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 = None
+        ph_in = ph_a = ph_b = mm = add = mm_mul = add_mul = reshape = None
         for node in traced.graph.nodes:
             if node.op == "placeholder":
                 if str(node.target) == "a":

tracer/acc_tracer/acc_ops.py

@@ -361,6 +361,60 @@ 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=[
@@ -400,6 +454,13 @@ 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"