converter reorg and matmul

Matmul issue fixes and lint error check

moving matmul to individual file

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -31,6 +31,7 @@
 from torch_tensorrt.fx.converters.impl.squeeze import squeeze
 from torch_tensorrt.fx.converters.impl.select import select
 from torch_tensorrt.fx.converters.impl.slice import slice_op
+from torch_tensorrt.fx.converters.impl.matmul import matrix_multiply
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -239,7 +240,6 @@ def aten_ops_hardtanh(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-
     return activation.hardtanh(
         network, target, SourceIR.ATEN, name, args[0], args[1], args[2]
     )
@@ -262,6 +262,18 @@ def aten_ops_gelu(
     )
 
 
+@tensorrt_converter(torch.ops.aten.matmul)
+@tensorrt_converter(torch.ops.aten.mm.default)
+def aten_ops_matmul(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return matrix_multiply(network, target, SourceIR.ATEN, name, args[0], args[1])
+
+
 @tensorrt_converter(torch.ops.aten.fmod.Tensor)
 def aten_ops_fmod(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/impl/elementwise/ops.py

@@ -14,6 +14,8 @@
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     get_trt_tensor,
+    broadcast,
+    set_layer_name,
 )
 
 from torch_tensorrt.fx.converters.impl.elementwise.base import (

py/torch_tensorrt/fx/converters/impl/matmul.py

@@ -0,0 +1,53 @@
+import operator
+import warnings
+from typing import Optional, cast, Any
+
+import numpy as np
+
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, Shape
+from torch_tensorrt.fx.utils import torch_dtype_from_trt
+
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    get_trt_tensor,
+    broadcast,
+    set_layer_name,
+)
+
+
+def matrix_multiply(
+    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=torch_dtype_from_trt(input.dtype)
+        )
+
+    input_matrix_op = other_matrix_op = trt.MatrixOperation.NONE
+    preset_diff = 0
+
+    if len(input.shape) == 1:
+        preset_diff -= 1
+        input_matrix_op = trt.MatrixOperation.VECTOR
+
+    if len(other.shape) == 1:
+        preset_diff += 1
+        other_matrix_op = trt.MatrixOperation.VECTOR
+
+    input, other = broadcast(
+        network, input, other, f"{name}_input", f"{name}_other", preset_diff
+    )
+    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)

py/torch_tensorrt/fx/test/converters/aten_op/test_matmul_aten.py

@@ -0,0 +1,114 @@
+import unittest
+
+import torch
+import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
+from parameterized import param, parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+import torch
+import torch.nn as nn
+import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
+
+
+class TestMatMulConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2_2", (2, 3), (3, 2)),
+            ("2_2", (2, 3), (3, 1)),
+            # FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
+            # (2,3), (3,) torch.ops.aten.mv.default
+            # Following cases use torch.ops.aten.bmm.defauly
+            # ("4_3", (3,1,3,2), (2,2,3)),
+            # ("3_4", (3,1,3,2), (2,2,3)),
+            # ("3_4", (2, 2, 3), (3, 1, 3, 3)),
+            # ("4_2", (1, 2, 2, 3), (3, 2)),
+        ]
+    )
+    def test_matmul_other_constant(self, _, input_shape, other_shape):
+        class MatMul(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.other = nn.Parameter(torch.randn(*other_shape))
+
+            def forward(self, input):
+                return torch.matmul(input, self.other)
+
+        inputs = [torch.randn(*input_shape)]
+
+        self.run_test(
+            MatMul(),
+            inputs,
+            expected_ops={torch.ops.aten.mm.default},
+            test_explicit_batch_dim=(len(input_shape) >= 1),
+        )
+
+    @parameterized.expand(
+        [
+            ("2_2", (2, 3), (3, 2)),
+            ("1_2", (1, 3), (3, 2)),
+            # FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
+            # (2,3), (3,) torch.ops.aten.mv.default
+            # Following cases use torch.ops.aten.bmm.defauly
+            # ("4_3", (3,1,3,2), (2,2,3)),
+            # ("3_4", (3,1,3,2), (2,2,3)),
+            # ("3_4", (2, 2, 3), (3, 1, 3, 3)),
+            # ("4_2", (1, 2, 2, 3), (3, 2)),
+        ]
+    )
+    def test_matmul_input_constant(self, _, input_shape, other_shape):
+        class MatMul(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.input = nn.Parameter(torch.randn(*input_shape))
+
+            def forward(self, other):
+                return torch.matmul(self.input, other)
+
+        inputs = [torch.randn(*other_shape)]
+
+        self.run_test(
+            MatMul(),
+            inputs,
+            expected_ops={torch.ops.aten.mm.default},
+            test_explicit_batch_dim=True
+            # test_explicit_batch_dim=(len(other_shape) <= 2),
+        )
+
+    @parameterized.expand(
+        [
+            ("2_2", (2, 3), (3, 2)),
+            # ("2_3", (2, 3), (2, 3, 4)),
+            # ("4_4", (2, 2, 2, 3), (2, 1, 3, 2)),
+            # ("4_2", (2, 1, 2, 3), (3, 2)),
+            # ("2_1", (2, 3), (3,)),
+            # ("1_2", (3,), (3, 2)),
+            # ("1_1", (3,), (3,)),
+        ]
+    )
+    def test_matmul(self, _, input_shape, other_shape):
+        class MatMul(nn.Module):
+            def forward(self, input, other):
+                return torch.matmul(input, other)
+
+        inputs = [torch.randn(*input_shape), torch.randn(*other_shape)]
+        test_explicit_batch_dim = not (
+            input_shape[0] == other_shape[0]
+            and len(input_shape) > 2
+            and len(other_shape) > 2
+        )
+        self.run_test(
+            MatMul(),
+            inputs,
+            expected_ops={torch.ops.aten.mm.default},
+            test_explicit_batch_dim=test_explicit_batch_dim,
+        )
+
+    # FIXME: dynamic shape is giving bmm
+
+
+if __name__ == "__main__":
+    run_tests()