From e0a961b921922b906ee30ed3af52d3099c7c6864 Mon Sep 17 00:00:00 2001
From: Sean Nijjar <sean.nijjar@gmail.com>
Date: Thu, 26 Sep 2024 15:18:19 +0000
Subject: [PATCH] #0: fix CCL nightly tests

---
 .../operations/test_reduce_scatter_nightly.py | 79 ++++++++++++++++++-
 1 file changed, 78 insertions(+), 1 deletion(-)

diff --git a/tests/ttnn/unit_tests/operations/test_reduce_scatter_nightly.py b/tests/ttnn/unit_tests/operations/test_reduce_scatter_nightly.py
index 098bee32ad6a..9ab877af87cf 100644
--- a/tests/ttnn/unit_tests/operations/test_reduce_scatter_nightly.py
+++ b/tests/ttnn/unit_tests/operations/test_reduce_scatter_nightly.py
@@ -19,7 +19,6 @@
 @pytest.mark.parametrize(
     "num_devices, num_links",
     [
-        (4, 1),
         (8, 1),
     ],
 )
@@ -92,3 +91,81 @@ def test_reduce_scatter_nightly(
         num_iters=num_iters,
         enable_async=enable_async,
     )
+
+
+@pytest.mark.timeout(120)
+@pytest.mark.parametrize(
+    "num_devices, num_links",
+    [
+        (4, 2),
+    ],
+)
+@pytest.mark.parametrize(
+    "per_chip_output_shape, scatter_dim, layout",
+    [
+        ([1, 8, 1024, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 4, 1024, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 4, 2048, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 32], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 64], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 64, 64], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 128], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 256], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 512], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 2048], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 128, 1024], 3, ttnn.TILE_LAYOUT),
+        # Has worker slice size warning - defaults to 1x1
+        ([1, 1, 128, 8192], 3, ttnn.TILE_LAYOUT),
+        # Always fails with bfp8_b
+        ([1, 1, 2048, 1024], 3, ttnn.TILE_LAYOUT),
+        # Has worker slice size warning - defaults to 1x1
+        ([1, 1, 2048, 8192], 3, ttnn.TILE_LAYOUT),
+    ],
+)
+@pytest.mark.parametrize(
+    "input_dtype",
+    [
+        ttnn.bfloat16,
+        ttnn.bfloat8_b,
+    ],
+)
+@pytest.mark.parametrize(
+    "mem_config",
+    [
+        ttnn.MemoryConfig(buffer_type=ttnn.BufferType.DRAM),
+        ttnn.MemoryConfig(buffer_type=ttnn.BufferType.L1),
+    ],
+)
+@pytest.mark.parametrize("math_op", [ttnn.ReduceType.Sum])
+@pytest.mark.parametrize("enable_async", [True, False])
+def test_reduce_scatter_nightly(
+    pcie_mesh_device,
+    num_devices,
+    per_chip_output_shape,
+    scatter_dim,
+    num_links,
+    math_op,
+    input_dtype,
+    layout,
+    mem_config,
+    use_program_cache,
+    function_level_defaults,
+    enable_async,
+    num_iters=1,
+):
+    run_reduce_scatter_test(
+        pcie_mesh_device,
+        num_devices,
+        per_chip_output_shape,
+        scatter_dim,
+        num_links,
+        math_op,
+        input_dtype,
+        layout,
+        mem_config,
+        use_program_cache,
+        function_level_defaults,
+        num_iters=num_iters,
+        enable_async=enable_async,
+    )