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

import sys

import torch
import torch.distributed as dist
from torch.distributed._shard import (
    shard_parameter,
)
from torch.testing._internal.common_distributed import (
    requires_nccl,
    skip_if_lt_x_gpu,
)
from torch.testing._internal.common_utils import (
    TEST_WITH_DEV_DBG_ASAN,
    run_tests,
)
from torch.testing._internal.distributed._shard.sharded_tensor import (
    TEST_GPU_NUM,
    ShardedTensorTestBase,
    with_comms,
)
from torch.testing._internal.distributed._shard.sharded_tensor._test_ops_common import (
    clone_module_parameter,
    generate_chunk_sharding_specs_for_test,
    generate_local_weight_sharding_params_for_test,
)

if TEST_WITH_DEV_DBG_ASAN:
    print(
        "Skip dev-asan as torch + multiprocessing spawn have known issues",
        file=sys.stderr,
    )
    sys.exit(0)


class TestShardedEmbeddingBag(ShardedTensorTestBase):
    def _run_sharded_embedding_bag(
        self,
        spec,
        input_size,
        num_embeddings,
        embedding_dim,
        mode,
        include_last_offset=False,
        offset_size=None,
        max_norm=None,
        norm_type=2.0,
        padding_idx=None,
    ):
        # Use same seed.
        torch.manual_seed(0)
        local_embedding_bag = torch.nn.EmbeddingBag(
            num_embeddings,
            embedding_dim,
            mode=mode,
            max_norm=max_norm,
            norm_type=norm_type,
            include_last_offset=include_last_offset,
            padding_idx=padding_idx,
        ).cuda(self.rank)

        sharded_embedding_bag = torch.nn.EmbeddingBag(
            num_embeddings,
            embedding_dim,
            mode=mode,
            max_norm=max_norm,
            norm_type=norm_type,
            include_last_offset=include_last_offset,
            padding_idx=padding_idx,
        )

        # Copy the weights from local embedding bag.
        sharded_embedding_bag.weight = clone_module_parameter(
            local_embedding_bag, "weight"
        )

        # Shard the parameter.
        shard_parameter(sharded_embedding_bag, "weight", spec)

        # Run sharded computation
        torch.manual_seed(self.rank)  # inputs different on each rank
        inp = torch.randint(0, num_embeddings, tuple(input_size)).cuda(self.rank)
        per_sample_weights = None
        if mode == "sum":
            per_sample_weights = torch.rand(*input_size).cuda(self.rank)

        offsets = None
        if len(input_size) == 1:
            # We need to generate certain length offset for each rank.
            # The current implementation and dist API does not support
            # the case when the offset has different lengths.
            # input_size[0] >> offset_size, so the while loop will not
            # for too long.
            while offsets is None or (offsets.size(0) != offset_size):
                offsets = torch.randint(input_size[0], (offset_size,))
                offsets[0] = 0
                if include_last_offset:
                    offsets[-1] = input_size[0]
                offsets = (
                    torch.unique(offsets, sorted=True).contiguous().cuda(self.rank)
                )

        # If max_norm is set, we need to ensure that the renorm has been applied across
        # inputs from all ranks.
        if max_norm is not None:
            gathered_inputs = [torch.zeros_like(inp) for _ in range(TEST_GPU_NUM)]
            dist.all_gather(gathered_inputs, inp)
            unique_inp = torch.unique(torch.cat(gathered_inputs))
            offsets_dummy = torch.tensor([len(unique_inp) // 2]).cuda(self.rank)
            local_embedding_bag(unique_inp, offsets=offsets_dummy)

        sharded_output = sharded_embedding_bag(
            inp,
            offsets=offsets,
            per_sample_weights=per_sample_weights,
        )

        # Run local computation
        local_output = local_embedding_bag(
            inp,
            offsets=offsets,
            per_sample_weights=per_sample_weights,
        )

        # Compare local weight and shared one to ensure the renorm
        # as expected.
        if max_norm is not None:
            sharded_dim = spec.dim
            sharded_weight = sharded_embedding_bag.weight.local_shards()[0].tensor
            (start_pos, chunk_size) = generate_local_weight_sharding_params_for_test(
                local_embedding_bag.weight, sharded_dim, TEST_GPU_NUM, spec, self.rank
            )
            local_weight_narrowed = local_embedding_bag.weight.narrow(
                sharded_dim, start_pos, chunk_size
            )
            self.assertEqual(local_weight_narrowed, sharded_weight)

        # Verify
        self.assertEqual(local_output, sharded_output)

        # Validate for torch.nn.functional.embedding_bag version.
        local_output = torch.nn.functional.embedding_bag(
            inp,
            local_embedding_bag.weight,
            offsets=offsets,
            mode=mode,
            per_sample_weights=per_sample_weights,
            include_last_offset=include_last_offset,
            max_norm=max_norm,
            norm_type=norm_type,
            padding_idx=padding_idx,
        )
        sharded_output = torch.nn.functional.embedding_bag(
            inp,
            sharded_embedding_bag.weight,
            offsets=offsets,
            mode=mode,
            per_sample_weights=per_sample_weights,
            include_last_offset=include_last_offset,
            max_norm=max_norm,
            norm_type=norm_type,
            padding_idx=padding_idx,
        )

        self.assertEqual(local_output, sharded_output)

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(TEST_GPU_NUM)
    @requires_nccl()
    def test_sharded_embedding_bag_colwise(self):
        for spec in generate_chunk_sharding_specs_for_test(1):
            self._test_sharded_embedding_bag_with_test_cases(spec, 1)

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(TEST_GPU_NUM)
    @requires_nccl()
    def test_sharded_embedding_bag_rowwise(self):
        for spec in generate_chunk_sharding_specs_for_test(0):
            self._test_sharded_embedding_bag_with_test_cases(spec, 0)

    def _test_sharded_embedding_bag_with_test_cases(self, spec, sharded_dim):
        self._run_sharded_embedding_bag(spec, [5, 5], 17, 14, "sum")
        self._run_sharded_embedding_bag(spec, [5, 4], 17, 12, "mean")
        self._run_sharded_embedding_bag(spec, [6, 7], 21, 11, "max")
        self._run_sharded_embedding_bag(
            spec, [5, 5], 17, 14, "sum", max_norm=2.5,
        )
        self._run_sharded_embedding_bag(
            spec,
            [5, 4],
            17,
            12,
            "mean",
            max_norm=2.0,
            norm_type=1.0,
        )
        self._run_sharded_embedding_bag(
            spec,
            [6, 7],
            21,
            11,
            "max",
            max_norm=1.5,
            norm_type=1.0,
        )
        self._run_sharded_embedding_bag(spec, [5, 5], 17, 14, "sum", padding_idx=6)
        self._run_sharded_embedding_bag(spec, [8, 6], 24, 13, "sum")
        self._run_sharded_embedding_bag(spec, [4, 3], 16, 14, "max")
        self._run_sharded_embedding_bag(spec, [8], 23, 13, "sum", offset_size=3)
        self._run_sharded_embedding_bag(spec, [5], 17, 12, "mean", offset_size=2)
        self._run_sharded_embedding_bag(spec, [12], 16, 12, "max", offset_size=4)
        self._run_sharded_embedding_bag(
            spec, [8], 23, 13, "sum", offset_size=3, include_last_offset=True
        )
        self._run_sharded_embedding_bag(
            spec, [12], 16, 12, "max", offset_size=4, include_last_offset=True
        )
        self._run_sharded_embedding_bag(
            spec,
            [12],
            17,
            12,
            "sum",
            offset_size=3,
            max_norm=1.25,
        )
        self._run_sharded_embedding_bag(
            spec,
            [5],
            17,
            12,
            "mean",
            offset_size=2,
            max_norm=1.25,
        )
        self._run_sharded_embedding_bag(
            spec,
            [5],
            17,
            12,
            "max",
            offset_size=2,
            max_norm=1.15,
        )
        self._run_sharded_embedding_bag(spec, [4, 3], 16, 14, "sum", padding_idx=12)
        self._run_sharded_embedding_bag(spec, [4, 3], 16, 14, "mean", padding_idx=12)
        self._run_sharded_embedding_bag(spec, [4, 3], 16, 14, "max", padding_idx=12)
        self._run_sharded_embedding_bag(
            spec,
            [12],
            17,
            12,
            "sum",
            offset_size=3,
            max_norm=1.25,
            padding_idx=10,
        )
        self._run_sharded_embedding_bag(
            spec,
            [5],
            17,
            12,
            "mean",
            offset_size=2,
            max_norm=1.25,
            padding_idx=10,
        )
        self._run_sharded_embedding_bag(
            spec,
            [5],
            17,
            12,
            "max",
            offset_size=2,
            max_norm=1.15,
            padding_idx=10,
        )


if __name__ == "__main__":
    run_tests()
