[FX] Changes done internally at Facebook (#1625)

Co-authored-by: Wei Wei <wwei6@fb.com>

Author: frank-wei

.circleci/config.yml

@@ -263,7 +263,7 @@ commands:
     parameters:
       torch-build:
         type: string
-        default: "2.0.0.dev20230120+cu117"
+        default: "2.0.0.dev20230129+cu117"
       torch-build-index:
         type: string
         default: "https://download.pytorch.org/whl/nightly/cu117"
@@ -1026,7 +1026,7 @@ parameters:
   # Nightly platform config
   torch-build:
     type: string
-    default: "2.0.0.dev20230120+cu117"
+    default: "2.0.0.dev20230129+cu117"
   torch-build-index:
     type: string
     default: "https://download.pytorch.org/whl/nightly/cu117"

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -298,8 +298,6 @@ def aten_ops_sub(
     return acc_ops_converters.acc_ops_sub(network, target, None, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten._unsafe_view.default)
-@tensorrt_converter(torch.ops.aten._reshape_alias.default)
 @tensorrt_converter(torch.ops.aten.view.default)
 def aten_ops_reshape(
     network: TRTNetwork,
@@ -308,11 +306,33 @@ def aten_ops_reshape(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-        "acc_out_ty": acc_utils.build_raw_tensor_meta(shape=args[1]),
-    }
-    return acc_ops_converters.acc_ops_reshape(network, target, None, kwargs_new, name)
+    input_val = args[0]
+    # for case where input_val is TRTensor
+    input_val = get_trt_tensor(network, input_val, f"{name}_input_val")
+    shape = args[1]
+
+    layer = network.add_shuffle(input_val)
+
+    if all(isinstance(s, int) for s in shape):
+        layer.reshape_dims = tuple(shape)
+    else:
+        # Convert all the dimensions to trt Tensors.
+        trt_shape = []
+
+        for i, s in enumerate(shape):
+            if isinstance(s, TRTTensor):
+                trt_shape.append(s)
+            else:
+                a = get_trt_tensor(network, s, f"{name}_{i}")
+                trt_shape.append(a)
+
+        shape_layer = network.add_concatenation(inputs=trt_shape)
+        shape_layer.axis = 0
+        shape_layer.name = f"{name}_output_shape"
+        layer.set_input(1, shape_layer.get_output(0))
+
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)
 
 
 @tensorrt_converter(torch.ops.aten.cat.default)
@@ -345,3 +365,104 @@ def aten_ops_expand(
     return acc_ops_converters.acc_ops_expand_tensor(
         network, target, None, kwargs_new, name
     )
+
+
+@tensorrt_converter(operator.floordiv)
+def aten_ops_operator_floordiv(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input": args[0],
+        "other": args[1],
+    }
+    return acc_ops_converters.acc_ops_floor_div(network, target, None, kwargs_new, name)
+
+
+@tensorrt_converter(operator.mul)
+def aten_ops_operator_mul(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input": args[0],
+        "other": args[1],
+    }
+    return acc_ops_converters.acc_ops_mul(network, target, None, kwargs_new, name)
+
+
+@tensorrt_converter(operator.add)
+def aten_ops_operator_add(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input": args[0],
+        "other": args[1],
+    }
+    return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
+
+
+@tensorrt_converter(operator.sub)
+def aten_ops_operator_sub(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input": args[0],
+        "other": args[1],
+    }
+    return acc_ops_converters.acc_ops_sub(network, target, None, kwargs_new, name)
+
+
+@tensorrt_converter(torch.ops.aten.sym_numel)
+def aten_ops_sym_numel(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    shape_layer = network.add_shape(args[0])
+    set_layer_name(shape_layer, target, "_shape_layer")
+    reduce_layer = network.add_reduce(
+        shape_layer.get_output(0),
+        trt.ReduceOperation.PROD,
+        axes=get_axes_for_reduce_op(0, False),
+        keep_dims=True,
+    )
+    set_layer_name(reduce_layer, target, "_reduce_layer")
+    return reduce_layer.get_output(0)
+
+
+@tensorrt_converter(torch.ops.aten.sym_size)
+def aten_ops_sym_size(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    shape_layer = network.add_shape(args[0])
+    ind = args[1]
+    set_layer_name(shape_layer, target, "_shape_layer")
+    slice_layer = network.add_slice(
+        input=shape_layer.get_output(0),
+        start=[ind],
+        shape=[1],
+        stride=[1],
+    )
+    set_layer_name(slice_layer, target, "_slice_layer")
+    return slice_layer.get_output(0)

py/torch_tensorrt/fx/converters/convolution.py

@@ -1,8 +1,9 @@
 # @manual=//deeplearning/trt/python:py_tensorrt
+import logging
+
 import numpy as np
 import tensorrt as trt
 import torch
-import logging
 
 from ..converter_registry import tensorrt_converter
 

py/torch_tensorrt/fx/lower.py

@@ -7,6 +7,7 @@
 import torch
 import torch.fx as fx
 import torch.nn as nn
+import torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer as aten_tracer
 from torch.fx.passes.splitter_base import SplitResult
 
 from .fx2trt import TRTInterpreter, TRTInterpreterResult
@@ -18,8 +19,7 @@
 
 from .tracer.acc_tracer import acc_tracer
 from .trt_module import TRTModule
-from .utils import LowerPrecision, proxytensor_trace
-
+from .utils import LowerPrecision
 
 logger = logging.getLogger(__name__)
 
@@ -259,7 +259,9 @@ def create(
             return cls(
                 lower_pass_manager_builder=LowerPassManagerBuilder(
                     lower_setting=lower_setting,
-                    trace_func=lambda module, inputs: proxytensor_trace(module, inputs),
+                    trace_func=lambda module, inputs: aten_tracer.opt_trace(
+                        module, inputs
+                    ),
                     split_func=split_func,
                     lower_func=default_lower_pass(interpreter_builder),
                 )
@@ -308,14 +310,6 @@ def do_lower(module: nn.Module, inputs: Input) -> nn.Module:
                 pm = self.lower_pass_manager_builder.build_trt_lower_pipeline(
                     inputs, additional_inputs
                 )
-            if lower_setting.is_aten:
-                pm = self.lower_pass_manager_builder.build_aten2trt_lower_pipeline(
-                    inputs, additional_inputs
-                )
-            else:
-                pm = self.lower_pass_manager_builder.build_trt_lower_pipeline(
-                    inputs, additional_inputs
-                )
             lower_result = pm(module)
             return lower_result
 

py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py

@@ -127,6 +127,31 @@ def graph_optimization_pass(self) -> PassManager:
 
         return PassManager.build_from_passlist(passes)
 
+    def graph_optimization_pass_aten(self) -> PassManager:
+        passes = []
+
+        for p in self.lower_setting.customized_fuse_pass.passes:
+            passes.append(wrapper(p, self._input))
+        for p in self.lower_setting.lower_basic_fuse_pass.passes:
+            passes.append(wrapper(p, self._input))
+        # TODO fix this pass for aten graph
+        # if (
+        #     hasattr(self.lower_setting, "lower_precision")
+        #     and self.lower_setting.lower_precision is LowerPrecision.FP16
+        # ) or (
+        #     hasattr(self.lower_setting, "precision")
+        #     and self.lower_setting.precision is LowerPrecision.FP16
+        # ):
+        #     passes.append(wrapper(fix_clamp_numerical_limits_to_fp16, self._input))
+
+        passes.append(
+            inplace_wrapper(lambda m: FUSE_PASSES_POST_OBSERVER.observe(m, self._input))
+        )
+        # TODO we most likely do not need it for aten
+        # passes.append(fix_reshape_batch_dim)
+
+        return PassManager.build_from_passlist(passes)
+
     def _split_pass(self) -> PassManager:
         passes = [
             partial(
@@ -259,8 +284,7 @@ def build_aten2trt_lower_pipeline(
         passes.append(
             wrapper(self._trace_func, self._input),
         )
-        passes.append(self._default_replace_mutable_op_pass())
-        passes.append(self.graph_optimization_pass())
+        passes.append(self.graph_optimization_pass_aten())
         passes.append(self._split_pass())
         passes.append(self._trt_lower_pass())
 

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

@@ -73,7 +73,7 @@ def forward(self, x):
             # param("ceil_mode", 1, ceil_mode=True),
         ]
     )
-    @unittest.skip("PT tracer issue")
+    @unittest.skip("PT2 tracer issue")
     def test_max_pool3d(
         self,
         test_name,
@@ -95,6 +95,7 @@ def forward(self, x):
         inputs = [torch.randn(1, 3, 32, 32, 32)]
         self.run_test(TestModule(), inputs, expected_ops={})
 
+    @unittest.skip("PT2 tracer issue")
     def test_max_pool3d_with_dynamic_shape(self):
         class TestModule(torch.nn.Module):
             def __init__(self):
@@ -118,7 +119,7 @@ def forward(self, x):
     @parameterized.expand(
         [
             ("default", 1),
-            param("stride", 2, stride=()),
+            # param("stride", 2, stride=()),  #PT2 tracer issue
         ]
     )
     def test_stride_none_max_pool2d(
@@ -147,7 +148,7 @@ def forward(self, x):
             param("stride", 2, stride=()),
         ]
     )
-    @unittest.skip("PT tracer issue")
+    @unittest.skip("PT2 tracer issue")
     def test_stride_none_max_pool3d(
         self,
         test_name,
@@ -209,6 +210,7 @@ def forward(self, x):
             param("stride", 2, stride=()),
         ]
     )
+    @unittest.skip("PT2 tracer issue")
     def test_stride_none_max_pool3d_with_dynamic_shape(
         self,
         test_name,