Implementation of slice and select operations

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -572,3 +572,19 @@ def aten_ops_sigmoid(
         "input": args[0],
     }
     return add_sigmoid(network, target, kwargs_new, name)
+
+
+@tensorrt_converter(torch.ops.aten.select)
+def aten_ops_select(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input": args[0],
+        "dim": args[1],
+        "index": args[2],
+    }
+    return add_select(network, target.kwargs_new, name)

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -543,3 +543,27 @@ def type_cast(
     layer_i.set_output_type(0, cast_type)
     set_layer_name(layer_i, target, f"{name}_dtype_change")
     return layer_i.get_output(0)
+
+
+def to_numpy(tensor: Optional[torch.Tensor]) -> Optional[np.ndarray]:
+    """
+    Convert a PyTorch Tensor to a Numpy Array. If the tensor is
+    quantized it will be dequantized first.
+
+    Args:
+        tensor (Optional[torch.Tensor]): A PyTorch tensor or None.
+
+    Returns:
+        A Numpy array.
+    """
+
+    if tensor is None:
+        return tensor
+
+    assert isinstance(
+        tensor, torch.Tensor
+    ), f"to_numpy can only be called on None or a torch.Tensor, got: {tensor}"
+    if tensor.is_quantized:
+        tensor = tensor.dequantize()
+
+    return tensor.cpu().detach().contiguous().numpy()

py/torch_tensorrt/fx/converters/operator.py

@@ -22,6 +22,7 @@
 from .converter_utils import prepend_ones
 from .converter_utils import has_dynamic_shape
 from .converter_utils import get_shape_with_dynamic_shape
+from .converter_utils import to_numpy
 
 from ..types import (
     Shape,
@@ -278,30 +279,6 @@ def trunc_div(
     return output
 
 
-def to_numpy(tensor: Optional[torch.Tensor]) -> Optional[np.ndarray]:
-    """
-    Convert a PyTorch Tensor to a Numpy Array. If the tensor is
-    quantized it will be dequantized first.
-
-    Args:
-        tensor (Optional[torch.Tensor]): A PyTorch tensor or None.
-
-    Returns:
-        A Numpy array.
-    """
-
-    if tensor is None:
-        return tensor
-
-    assert isinstance(
-        tensor, torch.Tensor
-    ), f"to_numpy can only be called on None or a torch.Tensor, got: {tensor}"
-    if tensor.is_quantized:
-        tensor = tensor.dequantize()
-
-    return tensor.cpu().detach().contiguous().numpy()
-
-
 def trt_dtype_to_torch_dtype(trt_dtype):
     table = {
         trt.bool: torch.bool,
@@ -1050,3 +1027,44 @@ def add_expand(network, target, kwargs, name):
     layer = network.add_slice(input_val, start=start, shape=shape, stride=stride)
     set_layer_name(layer, target, name)
     return layer.get_output(0)
+
+
+def add_select(network, target, kwargs, name):
+    input_val = kwargs["input"]
+    if not isinstance(input_val, TRTTensor):
+        raise RuntimeError(
+            f"slice_tensor received input {input_val} that is not part "
+            "of the TensorRT region!"
+        )
+
+    ranks = len(input_val.shape) + (1 if network.has_implicit_batch_dimension else 0)
+    dim = get_positive_dim(cast(int, kwargs["dim"]), ranks)
+    dynamic_shape = has_dynamic_shape(input_val.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_val.shape[dim] != -1
+            ), "Can't select on negative shape dimension!"
+    index = kwargs[2]
+    if index >= input_val.shape[dim]:
+        raise RuntimeError(
+            f"cannot have index greater than the dimension length! {input_val.shape[dim]}"
+        )
+    output_shape = list(input_val.shape)
+    output_shape[dim] = 1
+    if dynamic_shape > 0:
+        output_shape = get_shape_with_dynamic_shape(
+            network, output_shape, input_val, target, name
+        )
+    layer = network.add_gather(input_val, dim, index)
+    out = layer.getOutput(0)
+    if len(out.shape) != 1:
+        layer = network.add_shuffle(out)
+    return layer.getOutput(0)

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

@@ -0,0 +1,56 @@
+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 TestSelectConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_index", 2, 1),
+        ]
+    )
+    def test_select(self, _, dim_test, index_test):
+        class TestModule(torch.nn.Module):
+            def __init__(self, dim, index):
+                super().__init__()
+                self.dim = dim
+                self.index = index
+
+            def forward(self, input):
+                return torch.select(input, self.dim, self.index)
+
+        input = [torch.randn(1, 3, 32)]
+        self.run_test(
+            TestModule(dim_test, index_test),
+            input,
+            expected_ops={torch.ops.aten.select},
+            test_explicit_precision=True,
+        )
+
+    # def test_select_with_dynamic_shape(self, _, dim_test, index_test):
+    #     class TestModule(torch.nn.Module):
+    #         def __init__(self, dim, index):
+    #             super().__init__()
+    #             self.dim = dim
+    #             self.index = index
+    #         def forward(self, input):
+    #             return torch.select(input, self.dim, self.index)
+
+    #     input_spec = [
+    #         InputTensorSpec(
+    #             shape=(-1, 3, 32),
+    #             dtype=torch.float32,
+    #             shape_ranges=[((1, 3, 3), (3, 3, 3), (32, 32, 32))],
+    #         ),
+    #     ]
+    #     self.run_test_with_dynamic_shape(
+    #         TestModule(dim_test, index_test), input_spec, expected_ops={torch.ops.aten.select}
+    #     )
+
+
+if __name__ == "__main__":
+    run_tests()