Convert aten.split to ttnn.split

tests/lowering/tensor_manipulation/test_split.py

@@ -0,0 +1,47 @@
+import torch
+import torch_ttnn
+import pytest
+import ttnn
+
+
+class SplitModule(torch.nn.Module):
+    def __init__(self, split_size_or_sections, dim):
+        super().__init__()
+        self.split_size_or_sections = split_size_or_sections
+        self.dim = dim
+
+    def forward(self, input):
+        return torch.split(input, self.split_size_or_sections, self.dim)
+
+
+SYMMETRY = pytest.mark.xfail(reason="ttnn.split only supports symmetric split to 2 tensors")
+
+
+@pytest.mark.parametrize(
+    "input_shape, split_size_or_sections, dim",
+    (
+        ((1, 4, 16, 32), 8, 2),
+        ((1, 4, 16, 32), 16, 3),
+        pytest.param((1, 4, 16, 32), 8, 3, marks=SYMMETRY),
+        pytest.param((1, 4, 16, 32), 4, 2, marks=SYMMETRY),
+        pytest.param((1, 4, 16, 32), (8, 8), 2, marks=SYMMETRY),
+    ),
+)
+def test_split(device, input_shape, split_size_or_sections, dim):
+    m = SplitModule(split_size_or_sections, dim)
+    input_tensor = torch.randn(input_shape, dtype=torch.bfloat16)
+    result_before = m.forward(input_tensor)
+    option = torch_ttnn.TorchTtnnOption(device=device)
+    # option.gen_graphviz = True
+    # The compilation is lazy, so we need to run forward once to trigger the compilation
+    m = torch.compile(m, backend=torch_ttnn.backend, options=option)
+    result_after = m.forward(input_tensor)
+    option._out_fx_graphs[0].print_tabular()
+
+    # Check the graph has been rewritten and contains ttnn ops
+    nodes = list(option._out_fx_graphs[0].nodes)
+    assert [node.target for node in nodes].count(ttnn.split) == 1
+
+    # Check inference result
+    for res_before, res_after in zip(result_before, result_after):
+        assert torch.allclose(res_before, res_after)

torch_ttnn/passes/lowering/to_tt_pass.py

@@ -576,6 +576,12 @@ def rewrite_node(node):
 
                 return None
 
+            if node.target == torch.ops.aten.split.Tensor:
+                tensor, split_size, dim = args
+                if 2 * split_size == tensor.meta["val"].size()[dim]:
+                    return g.call_function(ttnn.split, args=(tensor, 2, dim))
+                return None
+
             if node.target == torch.ops.aten.unsqueeze.default:
                 output_size = node.meta["val"].size()
                 output_size = list(output_size)