aten::split converter

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1,19 +1,19 @@
 import logging
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
+import tensorrt as trt
 import torch
 from torch.fx.node import Argument, Node, Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     cast_int_int_div_trt_tensor,
     cast_trt_tensor,
+    dynamic_unsupported,
 )
 from torch_tensorrt.fx.converters import acc_ops_converters
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-import tensorrt as trt
-
 from .converter_registry import dynamo_tensorrt_converter
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -279,6 +279,22 @@ def aten_ops_unsqueeze(
     )
 
 
+@dynamo_tensorrt_converter(
+    torch.ops.aten.split.default, capability_validator=dynamic_unsupported
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.split.sizes, capability_validator=dynamic_unsupported
+)
+def aten_ops_split(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.split(network, target, SourceIR.ATEN, name, args[0], args[1], args[3])
+
+
 @dynamo_tensorrt_converter(torch.ops.aten._softmax.default)
 def aten_ops_softmax(
     network: TRTNetwork,

py/torch_tensorrt/dynamo/conversion/impl/split.py

@@ -0,0 +1,82 @@
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+
+import numpy as np
+import torch
+import torch_tensorrt as trt
+from torch import Tensor
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
+from torch_tensorrt.fx.converters.converter_utils import (
+    has_dynamic_shape,
+    set_layer_name,
+)
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+
+def acc_ops_split(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    split_size_or_sections: Union[int, List(int)],
+    dim: Optional[Any] = 0,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"split received input {input} that is not part " "of the TensorRT region!"
+        )
+
+    dim = cast(int, dim)
+    dynamic_shape = has_dynamic_shape(input.shape)
+    if network.has_implicit_batch_dimension:
+        assert dim != 0, "Can't split on batch dim when it's implicit!"
+        dim -= 1
+    else:
+        if dynamic_shape > 0:
+            # Check whether slice target dim is dynamic shape dim
+            assert input.shape[dim] != -1, "Can't chunk on dynamic shape dimension!"
+
+    split_sizes = []
+    if type(split_size_or_sections) == int:
+        split_sizes.append(cast(int, split_size_or_sections))
+    else:
+        for split_size_or_section in split_size_or_sections:
+            split_sizes.append(cast(int, split_size_or_section))
+
+    start = [0] * len(input.shape)
+    stride = [1] * len(start)
+    offset = 0
+
+    if len(split_sizes) == 1:
+        num_splits = input.shape[dim] + split_sizes[0] - 1 // split_sizes[0]
+        split_sizes = [split_sizes[0]] * num_splits
+    else:
+        num_splits = len(split_sizes)
+
+    if num_splits < 1:
+        raise RuntimeError(
+            f"Invalid split: {input.shape[dim]} with split_size={split_sizes}"
+        )
+
+    max_offset = input.shape[dim]
+    # add slice layers
+    output = []
+    for i in range(num_splits):
+        shape = list(input.shape)
+        shape[dim] = min(split_sizes[i], cast(int, max_offset - offset))
+        start[dim] = offset
+        if dynamic_shape:
+            shape = get_shape_with_dynamic_shape(
+                network, shape, input, target, f"{name}_shape_{i}"
+            )
+        layer = network.add_slice(
+            input, start=start, shape=[] if dynamic_shape else shape, stride=stride
+        )
+        if dynamic_shape:
+            layer.set_input(2, shape)
+        offset += split_sizes[i]
+        set_layer_name(layer, target, f"{name}_{i}")
+        output.append(layer.get_output(0))
+    return output

tests/py/dynamo/converters/test_split_aten.py

@@ -0,0 +1,113 @@
+import torch
+from harness import DispatchTestCase
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+
+# FIXME: check about implicit and explicit batch
+class TestSliceConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_split_size_or_sections_no_dim", 2),
+            ("select_split_size_or_sections_list_no_dim", [1, 4]),
+            ("select_split_size_or_sections_list_no_dim_not_full_split", [1, 3]),
+        ]
+    )
+    def test_slice(self, _, split_size_or_tensor):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(input, split_size_or_tensor)
+                return out
+
+        input = torch.arange(10).reshape(5, 2)
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.slice.Tensor},
+        )
+
+
+class TestSliceConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_split_size_or_sections_dim", 2, 1),
+            ("select_split_size_or_sections_list_dim", [1, 4], 1),
+            ("select_split_size_or_sections_list_dim_not_full_split", [1, 3], 1),
+        ]
+    )
+    def test_slice(self, _, split_size_or_tensor, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(split_size_or_tensor, dim)
+                return out
+
+        input = torch.arange(10).reshape(2, 5)
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.slice.Tensor},
+        )
+
+
+class TestSliceConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_split_size_or_sections_dim", 2, 1),
+            ("select_split_size_or_sections_list_dim", [1, 4], 1),
+        ]
+    )
+    def test_slice(self, _, split_size_or_tensor, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.slice.Tensor(input, split_size_or_tensor, dim)
+                return out
+
+        input_specs = [
+            Input(
+                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},
+        )
+
+
+class TestSplitSymIntConverterImplicitBatch(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("select_chunk_dim", 6, 0),
+        ]
+    )
+    def test_chunk(self, _, chunk, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.ops.aten.chunk(input, chunk, dim)
+                return out
+
+        input = [torch.randn(11)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.split},
+        )
+
+
+if __name__ == "__main__":
+    run_tests()