Addressing review comments and checking in tests

apbose · apbose · commit 3f4a2e8045dd · 2023-09-20T15:13:08.000-07:00
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -352,22 +352,45 @@ def aten_ops_softmax(
     return impl.normalization.softmax(
         network, target, SourceIR.ATEN, name, args[0], args[1]
     )
+    
+def dynamic_unsupported_split(node: torch.fx.Node) -> bool:
+    # Validate that none of the inputs to the node have Dynamic shapes
+    assert isinstance(
+        node, torch.fx.Node
+    ), "Inputs to validator functions must be FX Nodes"
+
+    if isinstance(node.args[1], torch.fx.Node):
+        if getattr(node.args[1].meta["val"], "_has_symbolic_sizes_strides", True):
+            return False
+    return True
 
 
 @dynamo_tensorrt_converter(
-    torch.ops.aten.split.default, capability_validator=dynamic_unsupported
+    torch.ops.aten.split.Tensor, capability_validator=dynamic_unsupported_split
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.split.sizes, capability_validator=dynamic_unsupported_split
 )
 @dynamo_tensorrt_converter(
-    torch.ops.aten.split.sizes, capability_validator=dynamic_unsupported
+    torch.ops.aten.split_with_sizes.default,
+    capability_validator=dynamic_unsupported_split,
 )
 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[2])
+) -> Union[TRTTensor, Sequence[TRTTensor]]: 
+    return impl.split.split(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        split_size_or_sections=args[1],
+        dim=args_bounds_check(args, 2, 0),
+    )
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.where.self)  # type: ignore[misc]
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/split.py b/py/torch_tensorrt/dynamo/conversion/impl/split.py
@@ -20,8 +20,8 @@ def split(
     source_ir: Optional[SourceIR],
     name: str,
     input: TRTTensor,
-    split_size_or_sections: Union[int, List(int)],
-    dim: Optional[Any] = 0,
+    split_size_or_sections: Union[int, List[int]],
+    dim: int = 0,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
@@ -30,16 +30,12 @@ def split(
 
     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!"
+    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:
+    if isinstance(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:
@@ -48,12 +44,16 @@ def split(
     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]
+        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)
+        sum_split_sizes = sum(split_sizes)
+        if sum_split_sizes != input.shape[dim]:
+            raise RuntimeError(
+                f"split sizes don't add up to the tensor's size in the given dimension"
+            )
 
     if num_splits < 1:
         raise RuntimeError(
@@ -69,7 +69,7 @@ def split(
         start[dim] = offset
         if dynamic_shape:
             shape = get_shape_with_dynamic_shape(
-                network, shape, input, target, f"{name}_shape_{i}"
+                network, target, source_ir, f"{name}_shape_{i}", shape, input
             )
         layer = network.add_slice(
             input, start=start, shape=[] if dynamic_shape else shape, stride=stride
diff --git a/tests/py/dynamo/converters/test_split_aten.py b/tests/py/dynamo/converters/test_split_aten.py
@@ -3,15 +3,14 @@
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
 
 
 # FIXME: check about implicit and explicit batch
 class TestSplitConverterNoDim(DispatchTestCase):
     @parameterized.expand(
         [
             ("split_size_or_sections_no_dim", 2),
-            ("split_size_or_sections_list_no_dim", [1, 4]),
-            ("split_size_or_sections_list_no_dim_not_full_split", [1, 3]),
         ]
     )
     def test_split(self, _, split_size_or_tensor):
@@ -23,20 +22,62 @@ def forward(self, input):
                 out = torch.split(input, split_size_or_tensor)
                 return out
 
-        input = torch.arange(10).reshape(5, 2)
+        input = [torch.randn(10).reshape(5, 2)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.split.default},
+            expected_ops={torch.ops.aten.split.Tensor},
+            disable_passes=True,
         )
 
+    @parameterized.expand(
+        [
+            ("split_size_or_sections_list_no_dim_list", [1, 4]),
+        ]
+    )
+    def test_split_list(self, _, split_size_or_tensor):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.split(input, split_size_or_tensor)
+                return out
+
+        input = [torch.randn(10).reshape(5, 2)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.split_with_sizes.default},
+            disable_passes=True,
+        )
 
-class TestSplitConverterWithDim(DispatchTestCase):
     @parameterized.expand(
         [
-            ("split_size_or_sections_dim", 2, 1),
-            ("split_size_or_sections_list_dim", [1, 4], 1),
-            ("split_size_or_sections_list_dim_not_full_split", [1, 3], 1),
+            ("split_size_or_sections_list_no_dim_not_full_list", [1, 3]),
+        ]
+    )
+    def test_split_not_full_list(self, _, split_size_or_tensor):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.split(input, split_size_or_tensor)
+                return out
+
+        input = [torch.randn(10).reshape(5, 2)]
+        with self.assertRaises(RuntimeError):
+            self.run_test(
+                TestModule(),
+                input,
+                expected_ops={torch.ops.aten.split_with_sizes.default},
+                disable_passes=True,
+            )
+
+    @parameterized.expand(
+        [
+            ("split_size_or_sections_dims", 2, 1),
         ]
     )
     def test_split(self, _, split_size_or_tensor, dim):
@@ -45,25 +86,90 @@ def __init__(self):
                 super().__init__()
 
             def forward(self, input):
-                out = torch.split(split_size_or_tensor, dim)
+                out = torch.split(input, split_size_or_tensor, dim)
+                return out
+
+        input = [torch.randn(10).reshape(5, 2)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.split.Tensor},
+            disable_passes=True,
+        )
+
+    @parameterized.expand(
+        [
+            ("split_size_or_sections_list_dims", [1, 1], 1),
+        ]
+    )
+    def test_split_dim(self, _, split_size_or_tensor, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.split(input, split_size_or_tensor, dim)
                 return out
 
-        input = torch.arange(10).reshape(2, 5)
+        input = [torch.randn(10).reshape(5, 2)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.split.default},
+            expected_ops={torch.ops.aten.split_with_sizes.default},
+            disable_passes=True,
         )
 
+    @parameterized.expand(
+        [
+            ("split_size_or_sections_list_dims", [1, 1], 1),
+        ]
+    )
+    def test_split_dim_list(self, _, split_size_or_tensor, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.split(input, split_size_or_tensor, dim)
+                return out
+
+        input = [torch.randn(10).reshape(5, 2)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.split_with_sizes.default},
+            disable_passes=True,
+        )
 
-class TestSplitConverterDynamicShape(DispatchTestCase):
     @parameterized.expand(
         [
-            ("select_split_size_or_sections_dim", 2, 1),
-            ("select_split_size_or_sections_list_dim", [1, 4], 1),
+            ("split_size_or_sections_list_dims_not_full_list", [1, 1], 1),
         ]
     )
-    def test_split(self, _, split_size_or_tensor, dim):
+    def test_split_dim_list(self, _, split_size_or_tensor, dim):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                out = torch.split(input, split_size_or_tensor, dim)
+                return out
+
+        input = [torch.randn(15).reshape(5, 3)]
+        with self.assertRaises(RuntimeError):
+            self.run_test(
+                TestModule(),
+                input,
+                expected_ops={torch.ops.aten.split_with_sizes.default},
+                disable_passes=True,
+            )
+
+    @parameterized.expand(
+        [
+            ("select_split_size_or_sections_dim_dynamic_shape", 2, 1),
+        ]
+    )
+    def test_split_dynamic(self, _, split_size_or_tensor, dim):
         class TestModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -82,17 +188,16 @@ def forward(self, input):
         self.run_test_with_dynamic_shape(
             TestModule(),
             input_specs,
-            expected_ops={torch.ops.aten.split.default},
+            expected_ops={torch.ops.aten.split.Tensor},
+            disable_passes=True,
         )
 
-
-class TestSplitSymIntConverterImplicitBatch(DispatchTestCase):
     @parameterized.expand(
         [
             ("select_chunk_dim", 6, 0),
         ]
     )
-    def test_chunk(self, _, chunk, dim):
+    def test_split_dynamic(self, _, chunk, dim):
         class TestModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -102,11 +207,13 @@ def forward(self, input):
                 return out
 
         input = [torch.randn(11)]
-        self.run_test(
-            TestModule(),
-            input,
-            expected_ops={torch.ops.aten.split.default},
-        )
+        with self.assertRaises(UnsupportedOperatorException):
+            self.run_test(
+                TestModule(),
+                input,
+                expected_ops={torch.ops.aten.split.Tensor},
+                disable_passes=True,
+            )
 
 
 if __name__ == "__main__":
