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

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,14 @@ 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],
+        args_bounds_check(args, 2, trt.MatrixOperation.NONE),
+        args_bounds_check(args, 3, trt.MatrixOperation.NONE),
     )
 
 

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

@@ -19,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")
@@ -30,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: