fix: improve matmul dynamo converter

fix bug

add args input_matrix_op and other_matrix_op, support aten.mv.default

minor fix

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -171,6 +171,7 @@ def aten_ops_gelu(
 
 @dynamo_tensorrt_converter(torch.ops.aten.matmul)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.mm.default)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.mv.default)  # type: ignore[misc]
 def aten_ops_matmul(
     network: TRTNetwork,
     target: Target,
@@ -179,7 +180,12 @@ def aten_ops_matmul(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return impl.matmul.matrix_multiply(
-        network, target, SourceIR.ATEN, name, args[0], args[1]
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
     )
 
 

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

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.fx.converters.converter_utils import (
@@ -10,8 +11,6 @@
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
-import tensorrt as trt
-
 
 def matrix_multiply(
     network: TRTNetwork,
@@ -20,6 +19,8 @@ def matrix_multiply(
     name: str,
     input: TRTTensor,
     other: TRTTensor,
+    input_matrix_op: trt.MatrixOperation = trt.MatrixOperation.NONE,
+    other_matrix_op: trt.MatrixOperation = trt.MatrixOperation.NONE,
 ) -> TRTTensor:
     if not isinstance(input, trt.tensorrt.ITensor):
         input = get_trt_tensor(network, input, f"{name}_input")
@@ -31,7 +32,6 @@ def matrix_multiply(
             dtype=unified_dtype_converter(input.dtype, Frameworks.TORCH),
         )
 
-    input_matrix_op = other_matrix_op = trt.MatrixOperation.NONE
     preset_diff = 0
 
     if len(input.shape) == 1:
@@ -46,5 +46,5 @@ def matrix_multiply(
         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)
+    set_layer_name(layer, target, name, source_ir)
     return layer.get_output(0)