Changes done internally at Facebook (#1194)

6703b98dff0695d91026f057b951dba1355825fa Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.prod
c822345d6d673e1653c2208435e34ab400bada3d Jason Park <jasonjk@fb.com> Add support for generic torch ops to be used in training.
e5758602a0592d6c2b71d6d66a0398c4dd9b5e20 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for repeat interleave
c13c633f04df162500eed477c0569eb2b81eb070 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for reduce ops
863476cf43b210922b88585b8f196dd84fbebb56 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_op.convolution
68dff39793e5c30c20010919a855bb3d984015d7 Ruichao Xiao <xiaoruichao@fb.com> [fbcode][GPU][DHEN]fuse split squeeze cat as reshape
f8b920769507ebd2ff02419b4aece25451298a95 Ruichao Xiao <xiaoruichao@fb.com> [fbcode][DHEN][GPU] reorder and merge cats whose input is a sublist of another cat
5b6a8d2d6be979983a52ac96225fefb510c3817c Andrew Or <andrewor@fb.com> [Quant][fx] Rename convert_to_reference to convert_to_reference_fx
996a0e080b8a8bc0b292a7c2ac92f41f6db33a2e Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_op.expand
084631fe74b304fbb9481ca15fd452a3714fb1b8 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_op.to_dtype
b3195e76329ccddbb5c4640cfa884d0e457d2d34 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for std
a5d964e62bdf769cf8c2e67321138b33e1f524a7 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_op.tile
3d33d45b2fc7f10f25c22946ba474b227e4b6529 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for squeeze
09085abf63d7e7732e2cd66e600e8afc6d58964f Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_op.topk
65edc7ea12899e9bd2af42c890a64de853d9b7fe Huamin Li <huaminli@fb.com> temporarily skip gelu tests
d11e521f9b90554ca86912a49920afa4406bb40d Shirong Wu <shirong@fb.com> Suppress accuracy check for remove_reshape_with_batch_size_change
6d948298b2327d229e010a34f1c221b11d2eb504 Ankur Singla <ankursingla@fb.com> [GPULowering] Suppress accuracy check for fuse_unsqueeze_cat_sum
e780b647fc9571b77d9f41c963041a6ac3d66f33 Janet Yang <qxy11@fb.com> Lower xrayvideo2022 to fx2trt
433c7207fef16b1fdff985546ea969c39fa83e7c generatedunixname89002005287564 <generatedunixname89002005287564@fb.com> [Codemod][Remove @noautodeps and @autodeps-skip tags] deeplearning/trt 1/2
66fdb65cffa925660c77b4758388399db3cbfe48 Scott Wolchok <swolchok@fb.com> [fx2ait] Minor Python cleanup in acc_ops_getitem
188132ecb2c19bcbf83cb2dc381f6e3798629f87 generatedunixname89002005324833 <generatedunixname89002005324833@fb.com> [AutoAccept][Codemod][FBSourceBuckFormatLinter] Daily `arc lint --take BUCKFORMAT`
4536bae4686dd01f2149541ea7fb330e178a4969 Wei Wei <wwei6@fb.com> [fx2trt] support sub
064602e666f86c110d931cd90a8536112a19b4ad Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.interpolate
9dfd0ee0cecb1975e3f53c44de237d67ca443ec5 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for unary_ops
39b9efad8d5d82463a2016d135c0cf277de1c3c6 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for unsqueeze
2bb17667d1dabc95391950426fc1f921eb3d0959 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.split
64dfb7b096686cb2fd33197340dc72f30d525456 Shirong Wu <shirong@fb.com> Group LN trt plugin
438f670e28df59b0734baa092a514fba3d75eb4f Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.avgpool
df0fe32dae4343827bd9b37b72daae761b02f228 Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops masked fill
44fe735d3493ea2d05a56b49093e4a23dd63a98e Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shaope support for acc_ops.pad
4f931acca706d8ce79045ceafef2ea0486609149 Wei Wei <wwei6@fb.com> [fx2trt] torch.max dynamic shape test
bf6f6cbe217d26a95ca9122574adf7de3966db9e Shreyansh Prajapati <shreyanshp@fb.com> Change the name of the test from full_reduce to dim_reduce
1c5680ed107d9206f3514eff4069a3f6c870ba8c Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.type_as
33e4c175a4f5fec78ac0b1c8eb262ca777c7aaba Shreyansh Prajapati <shreyanshp@fb.com> Test dynamic shape support for acc_ops.min
f37be34bcef9716080b8bafbd1f4ad72e412c44c Wei Wei <wwei6@fb.com> [fx2trt] plugin for grid_sample
57b5cc6a0f4839686ae360361a3a13b424794ee7 generatedunixname89002005367269 <generatedunixname89002005367269@fb.com> [AutoAccept][Codemod][FBSourceBlackLinter] Daily `arc lint --take BLACK`
eb741cc5e5a7babdc94e72d411670905f54da3e0 Shreyansh Prajapati <shreyanshp@fb.com> Updated the dynamic shape support for narrow op
521c36b96a14741ae89d7af6cbb658120bcec2ea Shreyansh Prajapati <shreyanshp@fb.com> Removing the comment for 4 dims dynamic shape support after analysis
e947343375967fe9efb0a16fdb9f63bff1449328 Shreyansh Prajapati <shreyanshp@fb.com> Updated the pad test for dynamic batch for analysis
3d64087014e91bc301a315eae43683b1aa2b66bc Oleg Khabinov <khabinov@fb.com> [trt_bc] Some improvements
dfd937a56fa01aca88a89b46176befdac4c202c4 Shreyansh Prajapati <shreyanshp@fb.com> Updated the test for as_strided op for analysis
11d76d0420dcaa4bb8890dcdeb86b6e534af831c Bangsheng Tang <bangsheng@fb.com> [gpu][infer] replace fx2trt_layer_norm with fbgemm layer_norm
932046ff6ea6dead114c0222b23ca3854690cffa Wei Wei <wwei6@fb.com> [fx2trt] bridge the dynamic batch and fixed shape

Author: Wei

examples/fx/quantized_resnet_test.py

@@ -8,7 +8,7 @@
 import torchvision.models as models
 from torch.ao.quantization.quantize_fx import (
     convert_fx,
-    convert_to_reference,
+    convert_to_reference_fx,
     prepare_fx,
 )
 from torch.fx.experimental.normalize import NormalizeArgs
@@ -52,7 +52,7 @@ def build_int8_trt(rn18):
     prepared = prepare_fx(rn18, {"": qconfig}, data)
     for _ in range(10):
         prepared(data)
-    quantized_rn18 = convert_to_reference(prepared)
+    quantized_rn18 = convert_to_reference_fx(prepared)
     ref_res = quantized_rn18(data)
     print("quantized model:", quantized_rn18)
 

py/torch_tensorrt/_compile.py

@@ -117,7 +117,7 @@ def compile(module: Any, ir="default", inputs=[], enabled_precisions=set([_enums
         else:
             raise ValueError(f"Precision {enabled_precisions} not supported on FX")
 
-        return lower_to_trt(module, inputs, lower_precision=lower_precision, max_batch_size=inputs[0].size(0), explicit_batch_dimension=True)
+        return lower_to_trt(module, inputs, lower_precision=lower_precision, max_batch_size=inputs[0].size(0), explicit_batch_dimension=True, dynamic_batch=False)
     else:
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -591,9 +591,33 @@ def acc_ops_batch_norm(
     )
     power = np.ones_like(scale)
 
+    # For BatchNorm1d, reshape 1d to 2d
+    output_shape = input_val.shape
+    if not network.has_implicit_batch_dimension and len(input_val.shape) < 4:
+        assert (
+            len(get_dynamic_dims(input_val.shape)) <= 1
+        ), "BatchNorm1D with more than one dynamic dims is not currently supported."
+        reshape_layer = network.add_shuffle(input_val)
+        if len(input_val.shape) == 2:
+            reshape_layer.reshape_dims = (input_val.shape[0], input_val.shape[1], 1, 1)
+        else:  # len(input_val.shape) == 3
+            reshape_layer.reshape_dims = (
+                input_val.shape[0],
+                input_val.shape[1],
+                input_val.shape[2],
+                1,
+            )
+        set_layer_name(reshape_layer, target, f"{name}_reshape_2d")
+        input_val = reshape_layer.get_output(0)
     layer = network.add_scale(input_val, trt.ScaleMode.CHANNEL, bias, scale, power)
     set_layer_name(layer, target, name)
 
+    # For BatchNorm1d, reshape output back to 1d
+    if not network.has_implicit_batch_dimension and len(output_shape) < 4:
+        reshape_output_layer = network.add_shuffle(layer.get_output(0))
+        reshape_output_layer.reshape_dims = tuple(output_shape)
+        set_layer_name(reshape_output_layer, target, f"{name}_reshape_1d")
+        layer = reshape_output_layer
     return layer.get_output(0)
 
 
@@ -614,7 +638,18 @@ def acc_ops_layer_norm(network, target, args, kwargs, name):
     eps_field = trt.PluginField(
         "eps", np.array([kwargs["eps"]], dtype=np.float32), trt.PluginFieldType.FLOAT32
     )
-    field_collection = trt.PluginFieldCollection([gamma_field, beta_field, eps_field])
+    try:
+        normalized_shape = np.array(kwargs["normalized_shape"], dtype=np.int32)
+    except TypeError:
+        print("Unable to convert normalized_shape to a field, fall back to []")
+        normalized_shape = np.array([], dtype=np.int32)
+
+    normalized_shape_filed = trt.PluginField(
+        "normalized_shape", normalized_shape, trt.PluginFieldType.INT32
+    )
+    field_collection = trt.PluginFieldCollection(
+        [gamma_field, beta_field, eps_field, normalized_shape_filed]
+    )
 
     try:
         if network.has_implicit_batch_dimension:
@@ -2838,11 +2873,7 @@ def num_slice_types(slices):
         """
         Gather the number of slice in getitem slices.
         """
-        num_slice = 0
-        for s in slices:
-            if isinstance(s, slice) or isinstance(s, int):
-                num_slice += 1
-        return num_slice
+        return sum(1 for s in slices if isinstance(s, slice) or isinstance(s, int))
 
     def slice_to_trt_params(py_slice, dim_size):
         """
@@ -2878,9 +2909,9 @@ def slice_to_trt_params(py_slice, dim_size):
     new_slices = []
     for s in slices:
         if s == Ellipsis:
-            while num_ellipsis > 0:
+            # pass explicit start to guard against negative num_ellipsis
+            for _ in range(0, num_ellipsis):
                 new_slices.append(slice(None, None, None))
-                num_ellipsis -= 1
         else:
             new_slices.append(s)
     slices = new_slices

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -41,6 +41,11 @@ def get_trt_plugin(
     Returns:
         A TensorRT plugin that can be added to TensorRT network as Plugin layer.
     """
+    # print the registered plugins
+    # PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
+    # for plugin_creator in PLUGIN_CREATORS:
+    #     print(plugin_creator.name)
+
     plugin_registry = trt.get_plugin_registry()
     plugin_creator = plugin_registry.get_plugin_creator(
         plugin_name, version, plugin_namespace
@@ -214,7 +219,6 @@ def create_constant(
 
     if dtype:
         value = value.to(dtype)
-
     constant = network.add_constant(value.shape, to_numpy(value))
     constant.name = name
     return constant.get_output(0)

py/torch_tensorrt/fx/input_tensor_spec.py

@@ -6,15 +6,17 @@
 from .utils import get_dynamic_dims
 
 
-def generate_input_specs(
-    inputs, lower_setting, additional_inputs=None, fixed_shape=False
-):
+def generate_input_specs(inputs, lower_setting, additional_inputs=None):
     # AIT lower setting doesn't have explicit_batch_dimension field and
     # we just return None.
     if not hasattr(lower_setting, "explicit_batch_dimension"):
         return None
 
-    if not lower_setting.explicit_batch_dimension or fixed_shape:
+    # dynamic_batch is TRT only flag. It does not exist in AIT lower setting
+    if (
+        not lower_setting.explicit_batch_dimension
+        or lower_setting.dynamic_batch is False
+    ):
         return InputTensorSpec.from_tensors(inputs)
 
     # If we don't have additional inputs, we assume the first dimension

py/torch_tensorrt/fx/lower.py

@@ -36,7 +36,7 @@ def lower_to_trt(
     timing_cache_prefix="",
     save_timing_cache=False,
     cuda_graph_batch_size=-1,
-    dynamic_batch=False,
+    dynamic_batch=True,
 ) -> nn.Module:
     """
     Takes in original module, input and lowering setting, run lowering workflow to turn module

py/torch_tensorrt/fx/lower_setting.py

@@ -86,4 +86,4 @@ class LowerSetting(LowerSettingBasic):
     cuda_graph_batch_size: int = -1
     preset_lowerer: str = ""
     opt_profile_replica: int = 1
-    dynamic_batch: bool = False
+    dynamic_batch: bool = True

py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py

@@ -1,3 +1,4 @@
+import datetime
 from functools import partial, wraps
 from typing import Any, Callable, Optional, Sequence
 
@@ -142,6 +143,9 @@ def lower_func(split_result: SplitResult) -> nn.Module:
 
                 # Only acc submodules will be lowered.
                 if not submod_name.startswith(split_result.non_acc_submodule_prefix):
+                    print("Now lowering submodule", submod_name)
+                    lowering_start_time = datetime.datetime.now()
+
                     self.lower_setting.input_specs = generate_input_specs(
                         submod_inputs,
                         self.lower_setting,
@@ -156,6 +160,10 @@ def lower_func(split_result: SplitResult) -> nn.Module:
                     LOWER_SPLIT_POST_OBSERVER.observe(
                         submod_name, lowered_module, submod_inputs
                     )
+                    print(
+                        f"Lowering submodule {submod_name} elapsed time",
+                        datetime.datetime.now() - lowering_start_time,
+                    )
 
             return split_result.split_module
 

py/torch_tensorrt/fx/test/converters/acc_op/test_as_strided.py

@@ -3,7 +3,9 @@
 import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
-from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
+from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase
+
+# from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
 
 
 class TestConverter(AccTestCase):
@@ -30,7 +32,9 @@ def forward(self, x):
             test_implicit_batch_dim=False,
         )
 
-    # Testing with shape (-1, -1, -1, -1) results into error: RuntimeError: setStorage: sizes [2, 3], strides [1, 2], storage offset 0, and itemsize 8 requiring a storage size of 48 are out of bounds for storage of size 1
+    # Testing with shape (-1, 3) results into error:
+    # RuntimeError: setStorage: sizes [2, 3], strides [1, 2], storage offset 0, and itemsize 8 requiring a storage size of 48 are out of bounds for storage of size 16
+
     """
     def test_as_strided_with_dynamic_shape_four_dimensions(self):
         class Stride(nn.Module):
@@ -39,9 +43,9 @@ def forward(self, x):
 
         input_specs = [
             InputTensorSpec(
-                shape=(-1, -1, -1, -1),
+                shape=(-1, 3),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 1), (3, 3, 3, 3), (5, 5, 5, 5))],
+                shape_ranges=[((1, 3), (2, 3), (2, 3))],
             ),
         ]
 

py/torch_tensorrt/fx/test/converters/acc_op/test_avgpool.py

@@ -39,6 +39,48 @@ def forward(self, x):
         inputs = [torch.randn(1, 3, 224)]
         self.run_test(TestModule(), inputs, expected_ops={acc_ops.avg_pool1d})
 
+    @parameterized.expand(
+        [
+            ("default", 1),
+            ("kernal_size", 3),
+            ("stride", 1, 2),
+            ("tuple_parameters", 2, (1,), (1,)),
+            param("padding", 2, padding=1),
+            param("ceil_mode", 1, ceil_mode=True),
+            param("include_pad", 2, padding=1, count_include_pad=False),
+        ]
+    )
+    def test_avg_pool1d_with_dynamic_shape(
+        self,
+        test_name="default",
+        kernel_size=1,
+        stride=1,
+        padding=0,
+        ceil_mode=False,
+        count_include_pad=True,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.avg_pool = torch.nn.AvgPool1d(
+                    kernel_size, stride, padding, ceil_mode, count_include_pad
+                )
+
+            def forward(self, x):
+                return self.avg_pool(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, 3, 3),
+                dtype=torch.float32,
+                shape_ranges=[((1, 3, 3), (3, 3, 3), (3, 3, 3))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={acc_ops.avg_pool1d}
+        )
+
     def test_avg_pool2d_with_dynamic_shape_four_dimensions(
         self,
         test_name="default",
@@ -218,38 +260,6 @@ def forward(self, x):
             TestModule(), input_specs, expected_ops={acc_ops.avg_pool2d}
         )
 
-    # Testing with (-1, -1, -1, -1) results in error: RuntimeError: ShapeProp error for: node=%avg_pool1d : [#users=1] = call_function[target=torch.avg_pool1d](args = (%x, (1,), (1,), (0,), False, True), kwargs = {}) with meta={}
-    """
-    def test_avg_pool1d_with_dynamic_shape_four_dimensions(
-        self,
-        test_name="default",
-        kernel_size=1,
-        stride=1,
-        padding=0,
-        ceil_mode=False,
-        count_include_pad=True,
-    ):
-        class TestModule(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.avg_pool = torch.nn.AvgPool1d(
-                    kernel_size, stride, padding, ceil_mode, count_include_pad
-                )
-
-            def forward(self, x):
-                return self.avg_pool(x)
-
-        input_specs = [
-            InputTensorSpec(
-                shape=(-1, -1, -1, -1),
-                dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 1), (3, 3, 3, 3), (5, 5, 5, 5))],
-            ),
-        ]
-
-        self.run_test(TestModule(), input_specs, expected_ops={acc_ops.avg_pool1d})
-    """
-
 
 if __name__ == "__main__":
     run_tests()