Changes done internally at Facebook (#1390)

7ae5990ba20126da1e0a93ad0887cb1892ff48cd Janet Yang <[email protected]> Pass to remove for _validate_and_get_n_vectors
d269be2fc7d84738a642d1d53eb44e6886a28d0c Alex Beloi <[email protected]> [fx] add deferred weights (xl_weight) and tracing for xl_embedding_bag
6f233bc9c72d90a908db0548c9d2dbe853895137 Alex Beloi <[email protected]> [fx] fix out of bounds indices/offsets for embedding_bag ops with xl_weight
3ca3b21c6a85ab9a6e9de503d0f13ee713a7b67c Janet Yang <[email protected]> Support div, torch.norm
52955d93d25e857510ed1b765220e8e5b0b0bb08 Janet Yang <[email protected]> Pass to replace sum(elmtwise(X))/numel(X) w/ mean(elmtwise(X))
89c56ef76a7a329f244a013ac5ccb099cb00c3c0 Janet Yang <[email protected]> Support scalar clamp, fixes for nan_to_num and benchmark
48071d8da1dc66fffceb0b42ea386079f1fb9709 Wei Wei <[email protected]> [ads] bug fix in push_down_parrallel_split_ops
afdc533da031a64e162bb08c8629ff38739e24f8 Wei Wei <[email protected]> [fx2trt] disable dispatch trace leaf node test
9905612fd8e6e2e79dc2f2bd1fa5b5d7fd5c98c3 Shirong Wu <[email protected]> Add number constrain for fuse group ln
d160a7a5e554d37c142e13f100bf4d8739ced232 Wei Wei <[email protected]> add option to remove passes
c22f691e6eae1b06ecd301eb6285b32d5dc9717c Mike Iovine <[email protected]> [fx2trt] Support dict inputs in acc tracer
8c05a3c57b1f5c63108b979ef8c61411525d0b1f Mike Iovine <[email protected]> [fx2trt] Support namedtuple access in acc tracer getattr
ff2000594e3f3ff75e0074edf9c38b5609128bbd Janet Yang <[email protected]> Generalize remove split ops more
1580805d827eb40c941e769b0b99e7c6a3ed6f89 Wei Wei <[email protected]> [fx2trt] add reshape unit test
d6a975462071a3747d18edcbe87a3b143b3ece88 Archie Sravankumar <[email protected]> Added FX tracing for `log_softmax`
6943ac0e322077b36a03c50c4c9065de6cd32837 Sungmin Cho <[email protected]> Add replace_mutable_op lower pass
baab27b81b1275de92fdaf760a158ce951564d33 Donglin Xia <[email protected]> Register avg_pool3d for acc_op in acc_op.py
ae4c4e2c3c18d78542140fcc30e1c24f7c647ef3 Wei Wei <[email protected]> [aten2trt] init check-in
fc94c5e110d5552349b2634662eae41f9f0b8933 Wei Wei <[email protected]> [ads] fix a bug in fuse_parallel_linear
87ef03338c9a25c5a610a2eb590345e8935f8d75 Wei Wei <[email protected]> [aten2trt] add binary ops
fca64a5b09749284fc6028b510078257fd4717b1 Shirong Wu <[email protected]> Fix dper pass
2bb168517ace7e638cffc7a241b1cbf528790b92 Mike Iovine <[email protected]> [fx2trt] Add acc normalization blocklist
8c912e085cf8722d572698286020ae1ce055023d Zhijing Li (Accelerator Enablement) <[email protected]> Skip unstable test_conv_add_standalone_module
137a3977ffeb03d0387e8a95ff2f32f3d15b3de8 Wei Wei <[email protected]> [aten2trt] resnet support
f06174dbb190df4ea488ca99a81d4884b5ed3aa2 wwei6 <[email protected]> [fx2trt] compile
817c1f0b6278ce0ad04dd88d43d21e7390e3baea wwei6 <[email protected]> [aten2trt] init check-in
92ce42c16f34804584a7e553eddf897c9fa4f65e wwei6 <[email protected]> [aten2trt] binary op

Author: Wei

py/torch_tensorrt/fx/README.md

@@ -2,3 +2,20 @@ FX2TRT is merged as FX module in Torch-TensorRT
 
 - The user guide is in [link](../../../docsrc/tutorials/getting_started_with_fx_path.rst#installation)
 - The examples are moved to [link](../../../examples/fx)
+
+* Method 1. Follow the instrucions for Torch-TensorRT
+* Method 2. To install FX path only (Python path) and avoid the C++ build for torchscript path
+`
+    $ conda create --name python_env python=3.8
+    $ conda activate python_env
+    # Recommend to install PyTorch 1.12 and later
+    $ conda install pytorch torchvision torchtext cudatoolkit=11.3 -c pytorch-nightly
+    # Install TensorRT python package
+    $ pip3 install nvidia-pyindex
+    $ pip3 install nvidia-tensorrt==8.2.4.2
+    $ git clone https://github.com/pytorch/TensorRT.git
+    $ cd TensorRT/py && python setup.py install --fx-only && cd ..
+    $ pyton -c "import torch_tensorrt.fx"
+    # Test an example by
+    $ python py/torch_tensorrt/fx/example/lower_example.py
+`

py/torch_tensorrt/fx/converters/__init__.py

@@ -13,6 +13,7 @@
     from .transformation import *  # noqa: F401 F403
     from .quantization import *  # noqa: F401 F403
     from .acc_ops_converters import *  # noqa: F401 F403
+    from .aten_ops_converters import *  # noqa: F401 F403
 
     TRT_LOGGER = trt.Logger()
     trt.init_libnvinfer_plugins(TRT_LOGGER, "")

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -21,11 +21,63 @@
 from ..utils import get_dynamic_dims, torch_dtype_from_trt, torch_dtype_to_trt
 
 from .converter_utils import *  # noqa: F403
-
+from torch_tensorrt.fx.passes.lower_basic_pass import (
+    trt_transposed_linear,
+    trt_transposed_matmul,
+)
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
+@tensorrt_converter(trt_transposed_matmul)
+def trt_transposed_matmul_converter(network, target, args, kwargs, name):
+    lhs, rhs, lhs_transposed, rhs_transposed = args
+
+    if isinstance(lhs, torch.nn.Parameter):
+        lhs = get_trt_tensor(network, lhs, f"{name}_lhs")
+    if isinstance(rhs, torch.nn.Parameter):
+        rhs = get_trt_tensor(network, rhs, f"{name}_rhs")
+    layer = network.add_matrix_multiply(
+        lhs,
+        trt.MatrixOperation.TRANSPOSE if lhs_transposed else trt.MatrixOperation.NONE,
+        rhs,
+        trt.MatrixOperation.TRANSPOSE if rhs_transposed else trt.MatrixOperation.NONE,
+    )
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)
+
+
+@tensorrt_converter(trt_transposed_linear)
+def trt_transposed_linear_converter(network, target, args, kwargs, name):
+    input, weight, bias = args
+
+    weight = get_trt_tensor(network, weight.t(), f"{name}_weight")
+    bias = get_trt_tensor(network, bias.reshape(1, -1), f"{name}_bias")
+
+    input, weight = broadcast(
+        network,
+        input,
+        weight,
+        f"{input.name}_broadcast",
+        f"{weight.name}_broadcast",
+    )
+    layer = network.add_matrix_multiply(
+        input,
+        trt.MatrixOperation.TRANSPOSE,
+        weight,
+        trt.MatrixOperation.NONE,
+    )
+    set_layer_name(layer, target, f"{name}_mm")
+    return add_binary_elementwise_layer(
+        network,
+        layer.get_output(0),
+        bias,
+        trt.ElementWiseOperation.SUM,
+        target,
+        f"{name}_add",
+    )
+
+
 @tensorrt_converter(acc_ops.conv1d)
 def acc_ops_conv1d(
     network: TRTNetwork,
@@ -1975,7 +2027,10 @@ def acc_ops_max_poolnd(
             f"MaxPool2d received input {input_val} that is not part "
             "of the TensorRT region!"
         )
-    extend_len = 2 if target == acc_ops.max_pool2d else 3
+    if target not in (acc_ops.max_pool2d, acc_ops.max_pool3d):
+        extend_len = 2 if len(kwargs["kernel_size"]) == 2 else 3
+    else:
+        extend_len = 2 if target == acc_ops.max_pool2d else 3
     kernel_size = extend_attr_to_tuple(kwargs["kernel_size"], extend_len)
     stride = extend_attr_to_tuple(kwargs["stride"], extend_len)
     padding = extend_attr_to_tuple(kwargs["padding"], extend_len)
@@ -2259,8 +2314,11 @@ def acc_ops_adaptive_avg_poolnd(
             f"AdaptiveAvgPool2d received input {input_val} that is not part "
             "of the TensorRT region!"
         )
+    if target not in (acc_ops.adaptive_avg_pool3d, acc_ops.adaptive_avg_pool2d):
+        extend_len = 2 if len(kwargs["output_size"]) == 2 else 3
+    else:
+        extend_len = 2 if target == acc_ops.adaptive_avg_pool2d else 3
 
-    extend_len = 2 if target == acc_ops.adaptive_avg_pool2d else 3
     assert all(
         input_val.shape[-(i + 1)] != -1 for i in range(extend_len)
     ), "AdaptiveAvgPool2d and AdaptiveAvgPool3d currently doesn't support dynamic shapes for last two dims."
@@ -2747,7 +2805,10 @@ def acc_ops_linear(
 
     if isinstance(kwargs["weight"], torch.Tensor):
         weight = get_trt_tensor(network, kwargs["weight"].t(), f"{name}_weight")
-        weight_op = trt.MatrixOperation.NONE
+        if target is not acc_ops.linear:
+            weight_op = trt.MatrixOperation.TRANSPOSE
+        else:
+            weight_op = trt.MatrixOperation.NONE
     else:
         assert isinstance(
             kwargs["weight"], TRTTensor
@@ -2782,17 +2843,26 @@ def acc_ops_linear(
     return res
 
 
-def add_clamp(network, input, val, op):
-    acc_ops_clamp_shape = (1,) * len(input.shape)  # broadcast all dimensions
-    acc_ops_clamp_tensor = (
-        val
-        * torch.ones(acc_ops_clamp_shape, dtype=torch_dtype_from_trt(input.dtype))
-        .cpu()
-        .numpy()
-    )
-    acc_ops_clamp_trt = network.add_constant(acc_ops_clamp_shape, acc_ops_clamp_tensor)
-    layer = network.add_elementwise(input, acc_ops_clamp_trt.get_output(0), op)
-
+def add_clamp(network, input, val, op, name):
+    if not len(input.shape):
+        # clamping scalar
+        acc_ops_clamp_trt = get_trt_tensor(
+            network,
+            squeeze_left(torch.tensor([val], dtype=torch_dtype_from_trt(input.dtype))),
+            f"{name}_clamp_{val}",
+        )
+    else:
+        acc_ops_clamp_shape = (1,) * len(input.shape)  # broadcast all dimensions
+        acc_ops_clamp_tensor = (
+            val
+            * torch.ones(acc_ops_clamp_shape, dtype=torch_dtype_from_trt(input.dtype))
+            .cpu()
+            .numpy()
+        )
+        acc_ops_clamp_trt = network.add_constant(
+            acc_ops_clamp_shape, acc_ops_clamp_tensor
+        ).get_output(0)
+    layer = network.add_elementwise(input, acc_ops_clamp_trt, op)
     return layer
 
 
@@ -2816,13 +2886,13 @@ def acc_ops_clamp(
 
     if min_val is not None:
         clamp_min_layer = add_clamp(
-            network, input_val, min_val, trt.ElementWiseOperation.MAX
+            network, input_val, min_val, trt.ElementWiseOperation.MAX, name
         )
         set_layer_name(clamp_min_layer, target, f"{name}_clamp_min")
         input_val = clamp_min_layer.get_output(0)
     if max_val is not None:
         clamp_max_layer = add_clamp(
-            network, input_val, max_val, trt.ElementWiseOperation.MIN
+            network, input_val, max_val, trt.ElementWiseOperation.MIN, name
         )
         set_layer_name(clamp_max_layer, target, f"{name}_clamp_max")
         input_val = clamp_max_layer.get_output(0)