converter: support bmm converter for dynamo

Author: bowang007

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -182,6 +182,17 @@ def aten_ops_matmul(
         network, target, SourceIR.ATEN, name, args[0], args[1]
     )
 
+@dynamo_tensorrt_converter(torch.ops.aten.bmm.default)
+def aten_ops_bmm(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.matmul.bmm(
+        network, target, SourceIR.ATEN, name, args[0], args[1]
+    )
 
 @dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)
 def aten_ops_layernorm(

py/torch_tensorrt/dynamo/conversion/impl/matmul.py

@@ -48,3 +48,37 @@ def matrix_multiply(
     layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
     set_layer_name(layer, target, name)
     return layer.get_output(0)
+
+def bmm(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    other: TRTTensor,
+) -> TRTTensor:
+    if not isinstance(input, trt.tensorrt.ITensor):
+        input = get_trt_tensor(network, input, f"{name}_input")
+    if not isinstance(other, trt.tensorrt.ITensor):
+        other = get_trt_tensor(
+            network,
+            other,
+            f"{name}_other",
+            dtype=unified_dtype_converter(input.dtype, Frameworks.TORCH),
+        )
+
+    input_matrix_op = other_matrix_op = trt.MatrixOperation.NONE
+    preset_diff = 0
+
+    if len(input.shape) != 3:
+        raise RuntimeError(f"Expected 3-dimensional tensor, but got ")
+
+    if len(other.shape) != 3:
+        raise RuntimeError(f"Expected 3-dimensional tensor, but got")
+    
+    if (input.shape[0] != other.shape[0]):
+        raise RuntimeError("expected input tensors to have same batch size.")
+    
+    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)

tests/py/dynamo/converters/test_bmm.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.dynamo.test_utils import DispatchTestCase
+from torch_tensorrt import Input
+
+class TestBmmConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("10_3_5", (10, 3, 4), (9, 4, 5)),
+        ]
+    )
+    def test_bmm(self, _, input_shape, mat2_shape):
+        class BMM(nn.Module):
+            def __init__(self):
+                super().__init__()
+                
+            def forward(self, input, mat2):
+                return torch.bmm(input, mat2)
+            
+        inputs = [torch.randn(*input_shape), torch.randn(*mat2_shape)]
+
+
+        self.run_test(
+            BMM(), 
+            inputs,
+            expected_ops={},
+        )
+        
+
+if __name__ == "__main__":
+    run_tests()