# Copyright (c) Meta Platforms, Inc. and affiliates
# implement matrix related ops for distributed tensor
import itertools
from typing import List, Optional

import torch
from torch.distributed._tensor._op_schema import OpSchema, OpStrategy, PlacementStrategy
from torch.distributed._tensor.ops.basic_strategy import gen_einsum_strategies
from torch.distributed._tensor.ops.utils import (
    generate_redistribute_costs,
    infer_broadcast_dims_map,
    is_tensor_shardable,
    map_placements_after_broadcast,
    register_op_strategy,
)
from torch.distributed._tensor.placement_types import (
    DTensorSpec,
    Placement,
    Replicate,
    Shard,
)

from torch.distributed.device_mesh import DeviceMesh

aten = torch.ops.aten


@register_op_strategy(aten.t.default)
def transpose_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    self_strategy = op_schema.args_schema[0]
    assert isinstance(self_strategy, OpStrategy)

    transpose_strategies = []
    for input_strategy in self_strategy.strategies:
        input_spec = input_strategy.output_spec
        # follow the input spec but transpose the Shard placements
        output_placements = [
            Shard(1 - p.dim) if isinstance(p, Shard) else p
            for p in input_spec.placements
        ]
        transpose_strategy = PlacementStrategy(
            output_specs=DTensorSpec(
                mesh=input_strategy.output_spec.mesh,
                placements=tuple(output_placements),
            ),
            input_specs=(input_strategy.output_spec,),
        )
        transpose_strategies.append(transpose_strategy)

    return OpStrategy(strategies=transpose_strategies)


def _mm_like_strategy(
    mm_equation: str, mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    self_strategy, mat2_strategy = op_schema.args_schema
    assert isinstance(self_strategy, OpStrategy)
    assert isinstance(mat2_strategy, OpStrategy)
    # generate all possible strategies for mm
    mm_strategy = gen_einsum_strategies(mm_equation, mesh)
    # filter out invalid strategies and associate costs
    strategies = mm_strategy.strategies
    filtered_strategies = []
    for strtg in strategies:
        assert strtg.input_specs is not None
        self_spec = strtg.input_specs[0]
        mat2_spec = strtg.input_specs[1]
        if is_tensor_shardable(self_strategy.shape, self_spec) and is_tensor_shardable(
            mat2_strategy.shape, mat2_spec
        ):
            redistribute_cost = [
                generate_redistribute_costs(self_strategy, self_spec),
                generate_redistribute_costs(mat2_strategy, mat2_spec),
            ]
            strtg.redistribute_cost = redistribute_cost
            filtered_strategies.append(strtg)

    mm_strategy.strategies = filtered_strategies

    return mm_strategy


def _addmm_like_strategy(
    mm_equation: str, mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    self_strategy, mat1_strategy, mat2_strategy = op_schema.args_schema
    assert isinstance(self_strategy, OpStrategy)
    assert isinstance(mat1_strategy, OpStrategy)
    assert isinstance(mat2_strategy, OpStrategy)
    self_shape = self_strategy.shape
    mm_out_shape = torch.Size(
        [
            mat2_strategy.shape[-1] if i == len(mat1_strategy.shape) - 1 else dim_size
            for i, dim_size in enumerate(mat1_strategy.shape)
        ]
    )
    # generate all possible strategies for mm
    mm_strategy = gen_einsum_strategies(mm_equation, mesh)
    # filter out invalid strategies and associate costs
    strategies = mm_strategy.strategies
    filtered_strategies = []
    for strtg in strategies:
        # construct new strategy by consider the self arg
        assert strtg.input_specs is not None
        mat1_spec = strtg.input_specs[0]
        mat2_spec = strtg.input_specs[1]
        out_spec = strtg.output_spec

        # self arg's spec should follow the output of mm, but need
        # to consider broadcast for the self arg
        broadcast_dims_map = infer_broadcast_dims_map(mm_out_shape, self_shape)
        self_placements = map_placements_after_broadcast(
            out_spec.placements, mm_out_shape, broadcast_dims_map
        )
        self_spec = DTensorSpec(mesh=mesh, placements=self_placements)

        if is_tensor_shardable(mat1_strategy.shape, mat1_spec) and is_tensor_shardable(
            mat2_strategy.shape, mat2_spec
        ):
            # update input specs with new self spec
            strtg.input_specs = (self_spec, mat1_spec, mat2_spec)

            # associate costs
            redistribute_cost = [
                generate_redistribute_costs(self_strategy, self_spec),
                generate_redistribute_costs(mat1_strategy, mat1_spec),
                generate_redistribute_costs(mat2_strategy, mat2_spec),
            ]
            strtg.redistribute_cost = redistribute_cost
            filtered_strategies.append(strtg)

    mm_strategy.strategies = filtered_strategies

    return mm_strategy


@register_op_strategy(aten.mm.default)
def mm_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    return _mm_like_strategy("mk,kn->mn", mesh, op_schema)


@register_op_strategy(aten.addmm.default)
def addmm_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    return _addmm_like_strategy("mk,kn->mn", mesh, op_schema)


@register_op_strategy(aten.bmm.default)
def bmm_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    return _mm_like_strategy("bmk,bkn->bmn", mesh, op_schema)


@register_op_strategy(aten.baddbmm.default)
def baddmm_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    return _addmm_like_strategy("bmk,bkn->bmn", mesh, op_schema)


@register_op_strategy(aten._scaled_dot_product_flash_attention.default)
def scaled_dot_product_flash_attention_strategy(
    mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    # NOTE: currently we only support some simple strategies to support tensor parallelism
    # TODO: sdpa might be a good candidate for us to explore decomposed sharding propagation
    # as it involves: matmul, pointwise, reduction ops together.
    return_debug_mask = len(op_schema.args_schema) >= 6 and op_schema.args_schema[5]
    q_input_strategy = op_schema.args_schema[0]
    assert isinstance(q_input_strategy, OpStrategy)
    # assuming q/k/v have the same shape
    qkv_shape = q_input_strategy.shape

    all_mesh_dim_strategies = []

    for mesh_dim in range(mesh.ndim):
        single_mesh_dim_strategies = []

        # placement list stores placements of [outputs, inputs]
        # in the spda case, we have 3 valid tensor outputs and 3 tensor inputs
        # first we can always accept full replication for both inputs and outputs
        all_replicate: List[Placement] = [Replicate()] * 6
        single_mesh_dim_strategies.append(all_replicate)

        # second we can accept the sharding pattern of tensor parallelism, which
        # shard on the num of head dim
        qkv_sharding = Shard(1)  # num head dim
        output_sharding = Shard(1)  # num head dim
        logsumexp_sharding = Shard(1)  # num head dim
        if return_debug_mask:
            debug_attn_mask_sharding: Placement = Shard(1)  # num head dim
        else:
            # empty debug mask, replicated
            debug_attn_mask_sharding = Replicate()

        num_heads_dim_sharding = [
            output_sharding,
            logsumexp_sharding,
            debug_attn_mask_sharding,
            qkv_sharding,
            qkv_sharding,
            qkv_sharding,
        ]
        single_mesh_dim_strategies.append(num_heads_dim_sharding)

        # Context Parallelism: shards on the sequence dim
        single_mesh_dim_strategies.append(
            [
                Shard(2),  # output
                Shard(2),  # logsumexp
                Shard(2),  # debugattn
                Shard(2),  # q
                Shard(2),  # k
                Shard(2),  # v
            ]
        )

        all_mesh_dim_strategies.append(single_mesh_dim_strategies)

    strategy_combs = itertools.product(*all_mesh_dim_strategies)

    all_strategies = []
    for strategy_comb in strategy_combs:
        spec_list = []
        for specs in zip(*strategy_comb):
            spec_list.append(DTensorSpec(mesh, tuple(specs)))

        assert len(spec_list) == 6
        input_expected_specs = spec_list[3:]
        output_specs: List[Optional[DTensorSpec]] = list(spec_list[:3])
        # fix up output_specs and fill in None for the int and empty tensor return values
        for i in range(2, 8):
            output_specs.insert(i, None)
        if all(is_tensor_shardable(qkv_shape, spec) for spec in input_expected_specs):
            # only add to the strategy list when all inputs are shardable
            redistribute_cost = []
            for input_idx, spec in enumerate(input_expected_specs):
                qkv_strategy = op_schema.args_schema[input_idx]
                assert isinstance(qkv_strategy, OpStrategy)
                qkv_tensor_meta = qkv_strategy.strategies[0].output_spec.tensor_meta
                spec.tensor_meta = qkv_tensor_meta
                redistribute_cost.append(
                    generate_redistribute_costs(qkv_strategy, spec)
                )

            strat = PlacementStrategy(
                output_specs=tuple(output_specs),
                input_specs=tuple(input_expected_specs),
                redistribute_cost=redistribute_cost,
            )
            all_strategies.append(strat)

    return OpStrategy(all_strategies)


@register_op_strategy(aten._scaled_dot_product_flash_attention_backward.default)
def scaled_dot_product_flash_attention_backward_strategy(
    mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    q_input_strategy = op_schema.args_schema[1]
    assert isinstance(q_input_strategy, OpStrategy)
    # assuming q/k/v have the same shape
    qkv_shape = q_input_strategy.shape

    tensor_input_indices = [
        i
        for i, arg_spec in enumerate(op_schema.args_schema)
        if isinstance(arg_spec, OpStrategy)
    ]
    num_tensor_inputs = len(tensor_input_indices)

    all_mesh_dim_strategies = []

    for mesh_dim in range(mesh.ndim):
        single_mesh_dim_strategies = []

        # placement list stores placements of [outputs, inputs]
        # in the spda backward case, we have 3 tensor outputs and 6 to 10 tensor inputs
        # first we can always accept full replication for both inputs and outputs
        all_replicate: List[Placement] = [Replicate()] * (3 + num_tensor_inputs)

        single_mesh_dim_strategies.append(all_replicate)

        # second we can accept the sharding pattern of tensor parallelism, which
        # shard on the num of head dim
        grad_output_sharding = Shard(1)  # num head dim
        qkv_sharding = Shard(1)  # num head dim
        output_sharding = Shard(1)  # num head dim
        logsumexp_sharding = Shard(1)  # num head dim
        grad_qkv_sharding = Shard(1)  # num head dim

        num_heads_dim_sharding: List[Placement] = [
            grad_qkv_sharding,
            grad_qkv_sharding,
            grad_qkv_sharding,
            grad_output_sharding,
            qkv_sharding,
            qkv_sharding,
            qkv_sharding,
            output_sharding,
            logsumexp_sharding,
        ]
        # accept replicate on the rest tensor inputs, potentially
        # cum_seq_q, cum_seq_k, philox_seed, philox_offset
        # at indices 6, 7, 12, 13, respectively
        num_heads_dim_sharding.extend([Replicate()] * (num_tensor_inputs - 6))
        single_mesh_dim_strategies.append(num_heads_dim_sharding)

        # Context Parallelism: shards on the sequence dim
        seq_dim_sharding: List[Placement] = [
            Shard(2),  # grad_q
            Shard(2),  # grad_k
            Shard(2),  # grad_v
            Shard(2),  # grad_output
            Shard(2),  # q
            Shard(2),  # k
            Shard(2),  # v
            Shard(2),  # output
            Shard(2),  # logsumexp
        ]
        # accept replicate on the rest tensor inputs, potentially
        # cum_seq_q, cum_seq_k, philox_seed, philox_offset
        # at indices 6, 7, 12, 13, respectively
        seq_dim_sharding.extend([Replicate()] * (num_tensor_inputs - 6))
        single_mesh_dim_strategies.append(seq_dim_sharding)

        all_mesh_dim_strategies.append(single_mesh_dim_strategies)

    strategy_combs = itertools.product(*all_mesh_dim_strategies)

    all_strategies = []
    for strategy_comb in strategy_combs:
        spec_list = []
        for specs in zip(*strategy_comb):
            spec_list.append(DTensorSpec(mesh, tuple(specs)))

        assert len(spec_list) == 3 + num_tensor_inputs
        input_expected_specs = spec_list[3:]
        output_specs: List[Optional[DTensorSpec]] = list(spec_list[:3])
        if all(
            is_tensor_shardable(qkv_shape, spec) for spec in input_expected_specs[:6]
        ):
            # only add to the strategy list when all inputs are shardable
            redistribute_cost = []
            for input_idx, spec in enumerate(input_expected_specs):
                qkv_strategy = op_schema.args_schema[tensor_input_indices[input_idx]]
                assert isinstance(qkv_strategy, OpStrategy)
                qkv_tensor_meta = qkv_strategy.strategies[0].output_spec.tensor_meta
                spec.tensor_meta = qkv_tensor_meta
                redistribute_cost.append(
                    generate_redistribute_costs(qkv_strategy, spec)
                )

            strat = PlacementStrategy(
                output_specs=tuple(output_specs),
                input_specs=tuple(input_expected_specs),
                redistribute_cost=redistribute_cost,
            )
            all_strategies.append(strat)

    return OpStrategy(all_strategies)


@register_op_strategy(aten.constant_pad_nd.default)
def constant_pad_nd_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> OpStrategy:
    # TODO(d4l3k); implement a more correct strategy for constant_pad_nd
    return OpStrategy(
        [
            PlacementStrategy(
                output_specs=DTensorSpec(mesh, (Replicate(),)),
                input_specs=(
                    DTensorSpec(mesh, (Replicate(),)),
                    DTensorSpec(mesh, (Replicate(),)),
                ),
                redistribute_cost=[[1]],
            )
        ]
    )


@register_op_strategy(aten._scaled_dot_product_efficient_attention.default)
def scaled_dot_product_efficient_attention_strategy(
    mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    # NOTE: currently we only support some simple strategies to support tensor parallelism
    q_input_strategy = op_schema.args_schema[0]
    assert isinstance(q_input_strategy, OpStrategy)
    # assuming q/k/v have the same shape
    qkv_shape = q_input_strategy.shape
    has_attn_bias = op_schema.args_schema[3] is not None
    compute_log_sumexp = op_schema.args_schema[4]

    all_mesh_dim_strategies = []

    for mesh_dim in range(mesh.ndim):
        single_mesh_dim_strategies = []

        # placement list stores placements of [outputs, inputs]
        # in the spda case, we have 2 valid tensor outputs and 3 or 4 tensor inputs
        # first we can always accept full replication for both inputs and outputs
        all_replicate: List[Placement] = [Replicate()] * (5 + has_attn_bias)
        single_mesh_dim_strategies.append(all_replicate)

        # second we can accept the sharding pattern of tensor parallelism, which
        # shard on the heads dimension
        qkv_sharding = Shard(1)
        output_sharding = Shard(1)
        if compute_log_sumexp:
            logsumexp_sharding: Placement = Shard(1)
        else:
            # empty logsumexp, replicated
            logsumexp_sharding = Replicate()

        num_heads_dim_sharding = [
            output_sharding,
            logsumexp_sharding,
            qkv_sharding,
            qkv_sharding,
            qkv_sharding,
        ]
        if has_attn_bias:
            num_heads_dim_sharding.append(Shard(1))
        single_mesh_dim_strategies.append(num_heads_dim_sharding)

        all_mesh_dim_strategies.append(single_mesh_dim_strategies)

    strategy_combs = itertools.product(*all_mesh_dim_strategies)

    all_strategies = []
    for strategy_comb in strategy_combs:
        spec_list = []
        for specs in zip(*strategy_comb):
            spec_list.append(DTensorSpec(mesh, tuple(specs)))

        assert len(spec_list) == (5 + has_attn_bias)
        input_expected_specs = spec_list[2:]
        output_specs: List[Optional[DTensorSpec]] = list(spec_list[:2])
        # fill in None for the scalar tensor return values
        # namely philox_seed and philox_offset
        output_specs.extend([None, None])
        if all(is_tensor_shardable(qkv_shape, spec) for spec in input_expected_specs):
            # only add to the strategy list when all inputs are shardable
            redistribute_cost = []
            for input_idx, spec in enumerate(input_expected_specs):
                qkv_strategy = op_schema.args_schema[input_idx]
                assert isinstance(qkv_strategy, OpStrategy)
                qkv_tensor_meta = qkv_strategy.strategies[0].output_spec.tensor_meta
                spec.tensor_meta = qkv_tensor_meta
                redistribute_cost.append(
                    generate_redistribute_costs(qkv_strategy, spec)
                )

            strat = PlacementStrategy(
                output_specs=tuple(output_specs),
                input_specs=tuple(input_expected_specs),
                redistribute_cost=redistribute_cost,
            )
            all_strategies.append(strat)

    return OpStrategy(all_strategies)


@register_op_strategy(aten._scaled_dot_product_efficient_attention_backward.default)
def scaled_dot_product_efficient_attention_backward_strategy(
    mesh: DeviceMesh, op_schema: OpSchema
) -> OpStrategy:
    q_input_strategy = op_schema.args_schema[1]
    assert isinstance(q_input_strategy, OpStrategy)
    # assuming q/k/v have the same shape
    qkv_shape = q_input_strategy.shape
    has_attn_bias = op_schema.args_schema[4] is not None

    tensor_input_indices = [
        i
        for i, arg_spec in enumerate(op_schema.args_schema)
        if isinstance(arg_spec, OpStrategy)
    ]

    all_mesh_dim_strategies = []

    for mesh_dim in range(mesh.ndim):
        single_mesh_dim_strategies = []

        # placement list stores placements of [outputs, inputs]
        # in the spda backward case, we have 4 tensor outputs and 8 or 9 tensor inputs
        # NOTE: Output sharding of grad_bias on heads dim if attn_bias is present;
        #       otherwise grad_bias will be empty and its DTensorSpec will be removed.
        # first we can always accept full replication for both inputs and outputs
        all_replicate: List[Placement] = [Replicate()] * (12 + has_attn_bias)

        single_mesh_dim_strategies.append(all_replicate)

        # second we can accept the sharding pattern of tensor parallelism, which
        # shard on the heads dimension
        grad_output_sharding = Shard(1)
        qkv_sharding = Shard(1)
        output_sharding = Shard(1)
        logsumexp_sharding = Shard(1)
        grad_qkv_sharding = Shard(1)
        grad_bias_sharding = Shard(1)

        num_heads_dim_sharding: List[Placement] = [
            grad_qkv_sharding,
            grad_qkv_sharding,
            grad_qkv_sharding,
            grad_bias_sharding,
            grad_output_sharding,
            qkv_sharding,
            qkv_sharding,
            qkv_sharding,
            # the place for optional input attn_bias,
            output_sharding,
            logsumexp_sharding,
        ]
        # input sharding of attn_bias on heads dim if present
        if has_attn_bias:
            num_heads_dim_sharding.insert(8, Shard(1))
        # accept replicate on the rest scalar tensor inputs
        # namely philox_seed and philox_offset
        num_heads_dim_sharding.extend([Replicate(), Replicate()])
        single_mesh_dim_strategies.append(num_heads_dim_sharding)

        all_mesh_dim_strategies.append(single_mesh_dim_strategies)

    strategy_combs = itertools.product(*all_mesh_dim_strategies)

    all_strategies = []
    for strategy_comb in strategy_combs:
        spec_list = []
        for specs in zip(*strategy_comb):
            spec_list.append(DTensorSpec(mesh, tuple(specs)))

        assert len(spec_list) == (12 + has_attn_bias)
        input_expected_specs = spec_list[4:]
        output_specs: List[Optional[DTensorSpec]] = list(spec_list[:4])
        # remove the DTensorSpec of output grad_bias if it's empty
        if not has_attn_bias:
            output_specs[-1] = None
        if all(is_tensor_shardable(qkv_shape, spec) for spec in input_expected_specs):
            # only add to the strategy list when all inputs are shardable
            redistribute_cost = []
            for input_idx, spec in enumerate(input_expected_specs):
                qkv_strategy = op_schema.args_schema[tensor_input_indices[input_idx]]
                assert isinstance(qkv_strategy, OpStrategy)
                qkv_tensor_meta = qkv_strategy.strategies[0].output_spec.tensor_meta
                spec.tensor_meta = qkv_tensor_meta
                redistribute_cost.append(
                    generate_redistribute_costs(qkv_strategy, spec)
                )

            strat = PlacementStrategy(
                output_specs=tuple(output_specs),
                input_specs=tuple(input_expected_specs),
                redistribute_cost=redistribute_cost,
            )
            all_strategies.append(strat)

    return OpStrategy(all_strategies)