# Owner(s): ["oncall: distributed"]

import torch
import torch.distributed as dist
import torch.distributed._functional_collectives as funcol
import torch.nn as nn
from torch.distributed.tensor import DeviceMesh, DTensor, Shard
from torch.distributed.tensor.debug import CommDebugMode
from torch.testing._internal.common_distributed import requires_accelerator_dist_backend
from torch.testing._internal.common_utils import run_tests, TestCase
from torch.testing._internal.distributed._tensor.common_dtensor import MLPModule
from torch.testing._internal.distributed.fake_pg import FakeStore


c10d_functional = torch.ops.c10d_functional
c10d_ops = torch.ops.c10d
device_type = (
    acc.type if (acc := torch.accelerator.current_accelerator(True)) else "cpu"
)


class TestCommMode(TestCase):
    def tearDown(self):
        super().tearDown()
        dist.destroy_process_group()

    def setUp(self):
        super().setUp()
        self.world_size = 2
        store = FakeStore()
        dist.init_process_group(
            backend="fake", rank=1, world_size=self.world_size, store=store
        )
        self.device_type = device_type
        self.world_pg = dist.distributed_c10d._get_default_group()

    def checksAssert(self, comm_mode, key, expected_value, expected_total_value):
        comm_counts = comm_mode.get_comm_counts()
        self.assertEqual(comm_mode.get_total_counts(), expected_total_value)
        self.assertEqual(comm_counts[key], expected_value)

        return

    def test_comm_mode(self):
        world_pg = self.world_pg

        class WrapperModel(nn.Module):
            def __init__(self, device):
                super().__init__()
                self.model = MLPModule(device=device)

            def forward(self, x):
                x = funcol.all_gather_tensor(x, 0, world_pg)
                x = funcol.reduce_scatter_tensor(x, "sum", 0, world_pg)
                out = self.model(x)
                return funcol.all_reduce(out, "sum", world_pg)

        model = WrapperModel(self.device_type)

        comm_mode = CommDebugMode()
        with comm_mode:
            model(torch.randn(20, 10, device=self.device_type))

        comm_counts = comm_mode.get_comm_counts()
        self.assertEqual(comm_mode.get_total_counts(), 3)
        self.assertEqual(comm_counts[c10d_functional.all_reduce], 1)
        self.assertEqual(comm_counts[c10d_functional.all_gather_into_tensor], 1)
        self.assertEqual(comm_counts[c10d_functional.reduce_scatter_tensor], 1)

    def test_comm_mode_coalesced(self):
        world_pg = self.world_pg

        class WrapperModelCoalesced(nn.Module):
            def __init__(self, device):
                super().__init__()
                self.model = MLPModule(device=device)

            def forward(self, x):
                x = funcol.all_gather_tensor(x, 0, world_pg)
                x = funcol.reduce_scatter_tensor(x, "sum", 0, world_pg)
                out = self.model(x)
                return funcol.all_reduce_coalesced([out], "sum", world_pg)

        model = WrapperModelCoalesced(self.device_type)

        comm_mode = CommDebugMode()
        with comm_mode:
            model(torch.randn(20, 10, device=self.device_type))

        comm_counts = comm_mode.get_comm_counts()
        self.assertEqual(comm_mode.get_total_counts(), 3)
        self.assertEqual(comm_counts[c10d_functional.all_reduce_coalesced], 1)
        self.assertEqual(comm_counts[c10d_functional.all_gather_into_tensor], 1)
        self.assertEqual(comm_counts[c10d_functional.reduce_scatter_tensor], 1)

    def test_comm_mode_with_dtensor(self):
        mesh = DeviceMesh(self.device_type, list(range(self.world_size)))

        def f(x, y):
            return torch.mm(x, y)

        comm_mode = CommDebugMode()
        x = torch.randn(4, 8, requires_grad=True)
        y = torch.randn(4, 32, requires_grad=True)
        x_dtensor = DTensor.from_local(x, mesh, [Shard(0)], run_check=False)
        y_dtensor = DTensor.from_local(y, mesh, [Shard(0)], run_check=False)

        with comm_mode:
            f(x_dtensor, y_dtensor)

        comm_counts = comm_mode.get_comm_counts()
        self.assertEqual(comm_mode.get_total_counts(), 1)
        self.assertEqual(comm_counts[c10d_functional.all_reduce], 0)
        self.assertEqual(comm_counts[c10d_functional.all_gather_into_tensor], 1)
        self.assertEqual(comm_counts[c10d_functional.reduce_scatter_tensor], 0)

    @requires_accelerator_dist_backend(["nccl", "xccl"])
    def test_comm_mode_with_c10d(self):
        if not torch.accelerator.is_available():
            return

        inp = torch.rand(2, 8, 16).to(device_type)
        all_gather_out = inp.new_empty(self.world_size * 2, 8, 16)

        comm_mode = CommDebugMode()

        # tests c10d all_reduce tracing
        with comm_mode:
            dist.all_reduce(inp)

        self.checksAssert(comm_mode, c10d_ops.allreduce_, 1, 1)

        # tests c10d all_gather_into_tensor tracing
        with comm_mode:
            dist.all_gather_into_tensor(all_gather_out, inp)

        self.checksAssert(comm_mode, c10d_ops._allgather_base_, 1, 1)

        # tests c10d reduce_scatter tracing
        with comm_mode:
            dist.reduce_scatter_tensor(inp, all_gather_out)

        self.checksAssert(comm_mode, c10d_ops._reduce_scatter_base_, 1, 1)

        # tests c10d broadcast tracing
        with comm_mode:
            dist.broadcast(inp, 0)

        self.checksAssert(comm_mode, c10d_ops.broadcast_, 1, 1)

        # tests c10d gather tracing
        with comm_mode:
            dist.gather(inp, None, 0)

        self.checksAssert(comm_mode, c10d_ops.gather_, 1, 1)

        # tests c10d reduce tracing
        with comm_mode:
            dist.reduce(inp, 0)

        self.checksAssert(comm_mode, c10d_ops.reduce_, 1, 1)

        # tests c10d scatter tracing
        with comm_mode:
            dist.scatter(inp, None, 0)

        self.checksAssert(comm_mode, c10d_ops.scatter_, 1, 1)

        # tests c10d all_gather tracing
        output_list = []

        with comm_mode:
            dist.all_gather(output_list, inp, None)

        self.checksAssert(comm_mode, c10d_ops.allgather_, 1, 1)

        # tests c10d allgather_coalesced_ tracing
        output_list = []

        with comm_mode:
            dist.all_gather_coalesced(output_list, [inp], None)

        self.checksAssert(comm_mode, c10d_ops.allgather_coalesced_, 1, 1)

        # tests c10d allgather_into_tensor_coalesced_ tracing
        with comm_mode, dist._coalescing_manager():
            dist.all_gather_into_tensor(all_gather_out, inp)

        self.checksAssert(comm_mode, c10d_ops.allgather_into_tensor_coalesced_, 1, 1)

        # tests c10d allreduce_coalesced
        with comm_mode:
            dist.all_reduce_coalesced(inp)

        self.checksAssert(comm_mode, c10d_ops.allreduce_coalesced_, 1, 1)

        # tests c10d reduce_scatter_
        with comm_mode:
            dist.reduce_scatter(all_gather_out, [inp])

        self.checksAssert(comm_mode, c10d_ops.reduce_scatter_, 1, 1)

        # tests c10d reduce_scatter_tensor_coalesced
        with comm_mode, dist._coalescing_manager():
            dist.reduce_scatter_tensor(all_gather_out, inp)

        self.checksAssert(comm_mode, c10d_ops.reduce_scatter_tensor_coalesced_, 1, 1)

        # tests c10d alltoall_
        with comm_mode:
            dist.all_to_all([inp], [inp])

        self.checksAssert(comm_mode, c10d_ops.alltoall_, 1, 1)

        # tests c10d alltoall_base_
        with comm_mode:
            dist.all_to_all_single(inp, inp)

        self.checksAssert(comm_mode, c10d_ops.alltoall_base_, 1, 1)


if __name__ == "__main__":
    run_tests()
