revert linear converter

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -3579,3 +3579,23 @@ def aten_ops_nonzero(
         name,
         args[0],
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.linear.default, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.linear, supports_dynamic_shapes=True)
+def aten_ops_linear(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.linear.linear(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        weight=args[1],
+        bias=args_bounds_check(args, 2, None),
+    )

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -12,6 +12,7 @@
     embedding,
     full,
     grid,
+    linear,
     matmul,
     nccl_ops,
     normalization,

py/torch_tensorrt/dynamo/conversion/impl/linear.py

@@ -0,0 +1,54 @@
+from typing import Optional, Union
+
+import numpy as np
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+from torch_tensorrt.dynamo.conversion import impl
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import SourceIR, get_trt_tensor
+from torch_tensorrt.fx.types import TRTTensor
+
+
+def linear(
+    ctx: ConversionContext,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    weight: Union[TRTTensor, torch.Tensor, np.ndarray],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+) -> TRTTensor:
+    # Process weight terms
+    if not isinstance(weight, (TRTTensor, torch.Tensor, np.ndarray)):
+        raise RuntimeError(
+            f"Linear layer {name} has weight of type {type(weight)}, Expect Union[TRTTensor, torch.Tensor, np.ndarray],"
+        )
+    elif isinstance(weight, (torch.Tensor, np.ndarray)):
+        weight = get_trt_tensor(ctx, weight, f"{name}_weight")
+
+    # Process bias terms
+    if bias is not None and not isinstance(bias, (TRTTensor, torch.Tensor, np.ndarray)):
+        raise RuntimeError(
+            f"Linear layer {name} has bias of type {type(bias)}, Expect Union[TRTTensor, torch.Tensor, np.ndarray],"
+        )
+    elif isinstance(bias, (torch.Tensor, np.ndarray)):
+        bias = get_trt_tensor(ctx, bias, f"{name}_bias")
+
+    # add IMatrixMultiplyLayer
+    out = impl.matmul.matrix_multiply(
+        ctx,
+        target,
+        source_ir,
+        name,
+        input,
+        weight,
+        input_matrix_op=trt.MatrixOperation.NONE,
+        other_matrix_op=trt.MatrixOperation.TRANSPOSE,
+    )
+
+    if bias is not None:
+        # add bias
+        out = impl.elementwise.add(ctx, target, source_ir, name, out, bias)
+
+    return out

py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py

@@ -171,6 +171,7 @@
     aten.upsample_bilinear2d.vec,
     aten.upsample_trilinear3d.vec,
     aten.upsample_bicubic2d.vec,
+    aten.linear.default,
 }
 
 

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -10,6 +10,7 @@
 from .constant_folding import constant_fold
 from .fuse_distributed_ops import fuse_distributed_ops
 from .fuse_prims_broadcast import fuse_prims_broadcast
+from .lower_linear import lower_linear
 from .pass_manager import DynamoPassManager
 from .remove_assert_nodes import remove_assert_nodes
 from .remove_detach import remove_detach
@@ -28,6 +29,7 @@
     accumulate_fp32_matmul,
     remove_num_users_is_0_nodes,
     complex_graph_detection,
+    lower_linear,
 ]
 
 pre_lowering_pass_list = [

py/torch_tensorrt/dynamo/lowering/passes/lower_linear.py

@@ -0,0 +1,42 @@
+import logging
+
+import torch
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
+    clean_up_graph_after_modifications,
+)
+from torch_tensorrt.dynamo.utils import get_metadata, set_metadata
+
+logger = logging.getLogger(__name__)
+
+
+def lower_linear(
+    gm: torch.fx.GraphModule, settings: CompilationSettings
+) -> torch.fx.GraphModule:
+    """Replace aten.linear with an equivalent implementation which can be easily converted to TRT"""
+    orig_op = torch.ops.aten.addmm.default
+    replacement_op = torch.ops.aten.linear.default
+
+    # Original graph
+    def orig(
+        input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
+    ) -> torch.Tensor:
+        W_T = torch.ops.aten.permute.default(weight, [1, 0])
+        out = orig_op(bias, input, W_T)
+        return out
+
+    # Replacement graph
+    def replacement(
+        input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
+    ) -> torch.Tensor:
+        return replacement_op(input, weight, bias)
+
+    metadata = get_metadata(gm, orig_op)
+    replaced_nodes = torch.fx.subgraph_rewriter.replace_pattern(gm, orig, replacement)
+
+    if len(replaced_nodes) > 0:
+        gm = clean_up_graph_after_modifications(gm)
+        set_metadata(gm, replacement_op, metadata)
+        logger.debug(f"Graph after lowering linear:\n{gm.graph}")
+
+    return gm

tools/perf/perf_run.py

@@ -282,7 +282,7 @@ def run_dynamo(model, input_tensors, params, precision, batch_size):
         cache_built_engines=params.get("cache_built_engines", False),
         reuse_cached_engines=params.get("reuse_cached_engines", False),
         use_python_runtime=params.get("use_python_runtime", False),
-        optimization_level=params.get("optimization_level", 5),
+        optimization_level=params.get("optimization_level", 3),
     )
     end_compile = timeit.default_timer()
     compile_time_s = end_compile - start_compile
@@ -441,21 +441,26 @@ def run_tensorrt(
     if params["is_trt_engine"]:
         serialized_engine = model
     else:
-        # Export an ONNX model and convert to TRT
-        torch.onnx.export(model.eval().cuda(), tuple(input_tensors), "./tmp.onnx")
+        if params["onnx"]:
+            onnx_path = params["onnx"]
+        else:
+            # Export an ONNX model and convert to TRT
+            onnx_path = "./onnx-trt.onnx"
+            exp_program = torch.export.export(model.eval().cuda(), tuple(input_tensors))
+            torch.onnx.export(exp_program, tuple(input_tensors), onnx_path)
         builder = trt.Builder(logger)
         network = builder.create_network(
             1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
         )
         parser = trt.OnnxParser(network, logger)
-        success = parser.parse_from_file("./tmp.onnx")
+        success = parser.parse_from_file(onnx_path)
         if not success:
             raise ValueError("ONNX conversion failed")
 
         config = builder.create_builder_config()
         if precision == "fp16":
             config.set_flag(trt.BuilderFlag.FP16)
-        config.builder_optimization_level = params.get("optimization_level", 5)
+        config.builder_optimization_level = params.get("optimization_level", 3)
         start_compile = timeit.default_timer()
         serialized_engine = builder.build_serialized_network(network, config)
         end_compile = timeit.default_timer()
@@ -561,7 +566,7 @@ def run(
                 print("int8 precision expects calibration cache file for inference")
                 return False
 
-        if (model is None) and (backend in ("tensorrt", "ts_trt", "all")):
+        if (model is None) and (backend in ("ts_trt", "all")):
             warnings.warn(
                 f"Requested backend {backend} without specifying a TorchScript Model, "
                 + "skipping this backend"
@@ -585,7 +590,7 @@ def run(
                 batch_size,
             )
             run_tensorrt(
-                model,
+                model_torch,
                 input_tensors,
                 params,
                 precision,
@@ -606,7 +611,7 @@ def run(
             )
         elif backend == "tensorrt":
             run_tensorrt(
-                model,
+                model_torch,
                 input_tensors,
                 params,
                 precision,
@@ -641,6 +646,12 @@ def run(
         default="",
         help="Name of torch model file",
     )
+    arg_parser.add_argument(
+        "--onnx",
+        type=str,
+        default="",
+        help="ONNX model file which helps bypass the step of exporting ONNX from torchscript model. If this argument is provided, the ONNX will be directly converted to TRT engine",
+    )
     arg_parser.add_argument(
         "--inputs",
         type=str,
@@ -683,7 +694,7 @@ def run(
     arg_parser.add_argument(
         "--optimization_level",
         type=int,
-        default=5,
+        default=3,
         help="Builder optimization level for TensorRT",
     )
     arg_parser.add_argument(
@@ -767,7 +778,9 @@ def run(
         )
 
     backends = parse_backends(params["backends"])
-    if ("dynamo" in backends or "torch_compile" in backends) and (model_torch is None):
+    if any(
+        backend in ["dynamo", "torch_compile", "tensorrt"] for backend in backends
+    ) and (model_torch is None):
         raise ValueError(
             "No Pytorch model (nn.Module) is provided for torchdynamo compilation. Please provide a pytorch model using --model_torch argument"
         )