Converter reorg and select operation

select operation correction and linting changes

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -29,6 +29,7 @@
 from torch_tensorrt.fx.converters.impl.normalization import layer_norm
 from torch_tensorrt.fx.converters.impl.normalization import softmax
 from torch_tensorrt.fx.converters.impl.squeeze import squeeze
+from torch_tensorrt.fx.converters.impl.select import select
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -626,6 +627,17 @@ def aten_ops_operator_add(
     return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
 
 
+@tensorrt_converter(torch.ops.aten.select.int)
+def aten_ops_select(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return select(network, target, SourceIR.ATEN, name, args[0], args[1], args[2])
+
+
 @tensorrt_converter(operator.sub)
 def aten_ops_operator_sub(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/impl/select.py

@@ -0,0 +1,70 @@
+import operator
+import warnings
+from typing import Union, Callable, Any, Optional, cast
+
+import numpy as np
+
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, Shape
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    get_positive_dim,
+    has_dynamic_shape,
+    to_numpy,
+)
+from torch_tensorrt.fx.converters.impl.shape import get_shape_with_dynamic_shape
+
+
+def select(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: Shape,
+    index: Shape,
+) -> TRTTensor:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"slice_tensor received input {input} that is not part "
+            "of the TensorRT region!"
+        )
+
+    ranks = len(input.shape) + (1 if network.has_implicit_batch_dimension else 0)
+    dim = get_positive_dim(cast(int, dim), ranks)
+    dynamic_shape = has_dynamic_shape(input.shape)
+    if network.has_implicit_batch_dimension:
+        if dim == 0:
+            raise RuntimeError(
+                f"We do not support slice_tensor at batch dim when it's implicit, got {dim}!"
+            )
+        dim = dim - 1
+    else:
+        if dynamic_shape:
+            # Check whether slice target dim is dynamic shape dim
+            assert input.shape[dim] != -1, "Can't select on negative shape dimension!"
+    index = index
+
+    if index >= input.shape[dim]:
+        raise RuntimeError(
+            f"cannot have index greater than the dimension length! {input.shape[dim]}"
+        )
+    output_shape = list(input.shape)
+    output_shape[dim] = 1
+    if dynamic_shape > 0:
+        output_shape = get_shape_with_dynamic_shape(
+            network, target, source_ir, name, output_shape, input
+        )
+    index_value = torch.tensor(index, dtype=torch.int32)
+    indices_tensor = network.add_constant(
+        index_value.shape, to_numpy(index_value)
+    ).get_output(0)
+    layer = network.add_gather(input, indices_tensor, dim)
+    out = layer.get_output(0)
+    if len(out.shape) != 1:
+        layer = network.add_shuffle(out)
+    return layer.get_output(0)

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

@@ -0,0 +1,83 @@
+import unittest
+
+import torch
+import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
+from parameterized import param, parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestSelectConverterOne(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_index", 1, 0),
+        ]
+    )
+    def test_select(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                return torch.select(input, dim, index)
+
+        input = [torch.randn(1, 2)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.select.int},
+            test_explicit_precision=True,
+        )
+
+
+class TestSelectConverterTwo(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_index", 1, 0),
+        ]
+    )
+    def test_select(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                return torch.select(input, dim, index)
+
+        input = [torch.randn(4, 4, 4, 4)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.select.int},
+            test_explicit_precision=True,
+        )
+
+
+class TestSelectConverterWithDynamicShape(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_index", 1, 0),
+        ]
+    )
+    def test_select_with_dynamic_shape(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                return torch.select(input, dim, index)
+
+        input_spec = [
+            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_spec, expected_ops={torch.ops.aten.select.int}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()