converter reorg and slice

converter reorg slice op

Correcting linting error and slice changes

Correcting the slice operation

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -30,6 +30,7 @@
 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
+from torch_tensorrt.fx.converters.impl.slice import slice_op
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -673,6 +674,27 @@ def aten_ops_sym_numel(
     return reduce_layer.get_output(0)
 
 
+@tensorrt_converter(torch.ops.aten.slice.Tensor)
+def aten_ops_slice(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return slice_op(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args[2],
+        args[3],
+        args[4],
+    )
+
+
 @tensorrt_converter(torch.ops.aten.sym_size)
 def aten_ops_sym_size(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/impl/slice/__init__.py

@@ -0,0 +1 @@
+from .ops import *

py/torch_tensorrt/fx/converters/impl/slice/base.py

@@ -0,0 +1,45 @@
+import operator
+import warnings
+from typing import Optional, cast
+
+import numpy as np
+
+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,
+    has_dynamic_shape,
+    set_layer_name,
+)
+
+from torch_tensorrt.fx.converters.impl.shape import get_shape_with_dynamic_shape
+
+
+def slice(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    start: Shape,
+    shape: Shape,
+    stride: Shape,
+) -> TRTTensor:
+    dynamic_shape = has_dynamic_shape(input.shape)
+    if dynamic_shape:
+        shape = get_shape_with_dynamic_shape(
+            network, target, source_ir, name, shape, input
+        )
+    layer = network.add_slice(
+        input,
+        start=start,
+        shape=[] if dynamic_shape else shape,
+        stride=stride,
+    )
+    if dynamic_shape:
+        layer.set_input(2, shape)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)

py/torch_tensorrt/fx/converters/impl/slice/ops.py

@@ -0,0 +1,64 @@
+import operator
+import warnings
+from typing import Optional, cast
+import math
+
+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
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    set_layer_name,
+    get_positive_dim,
+    has_dynamic_shape,
+)
+from torch_tensorrt.fx.converters.impl.shape import get_shape_with_dynamic_shape
+from torch_tensorrt.fx.converters.impl.slice.base import slice
+
+
+def slice_op(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: int,
+    start: int,
+    stop: int,
+    step: int,
+) -> 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 chunk on dynamic shape dimension!"
+    start_int = cast(int, start)
+    stop_int = cast(int, stop)
+    step_int = cast(int, step)
+    start = [0] * len(input.shape)
+    start[dim] = start_int
+    stride = [1] * len(start)
+    stride[dim] = step_int
+    output_shape = list(input.shape)
+    output_shape[dim] = math.ceil((stop_int - start_int) / step_int)
+
+    return slice(network, target, source_ir, name, input, start, output_shape, stride)

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

@@ -0,0 +1,89 @@
+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 TestSelectConverterImplicitBatch(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_start_stop_step", 0, 0, 7, 2),
+        ]
+    )
+    def test_slice(self, _, dim, start, stop, step):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(input, dim, start, stop, step)
+                return out
+
+        input = [torch.randn(10, 2, 3, 1)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.slice.Tensor},
+        )
+
+
+class TestSelectConverterExplicitBatch(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_start_stop_step", 1, 0, 7, 2),
+            ("select_dim_start_stop_step_exact", 1, 0, 10, 2),
+        ]
+    )
+    def test_slice(self, _, dim, start, stop, step):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(input, dim, start, stop, step)
+                return out
+
+        input = [torch.randn(10, 10, 3, 1)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.slice.Tensor},
+            test_explicit_precision=True,
+        )
+
+
+class TestSelectConverterDynamicShape(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_dim_start_stop_step", 1, 0, 7, 2),
+            ("select_dim_start_stop_step", 1, 0, 10, 2),
+        ]
+    )
+    def test_slice(self, _, dim, start, stop, step):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(input, dim, start, stop, step)
+                return out
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(1, 10, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 10, 1), (1, 10, 10), (1, 10, 10))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(),
+            input_specs,
+            expected_ops={torch.ops.aten.slice.Tensor},
+        )
+
+
+if __name__ == "__main__":
+    run_tests()