# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-unsafe

import operator

from typing import Callable, Dict, Optional, Set, Union

import executorch.backends.vulkan.custom_ops_lib  # noqa

import torch

from executorch.backends.vulkan.serialization.vulkan_graph_schema import (
    VkMemoryLayout,
    VkStorageType,
)

from executorch.backends.vulkan.utils import (
    all_memory_layouts,
    all_packed_dims,
    PackedDim,
)
from executorch.exir.dialects._ops import ops as exir_ops

from executorch.exir.dialects.edge._ops import EdgeOpOverload
from torch._subclasses.fake_tensor import FakeTensor

######################
## OpFeatures class ##
######################


def allow_node(node: torch.fx.Node) -> bool:
    return True


class TextureImplFeatures:
    __slots__ = [
        "valid_packed_dims",
        "uses_axis_map",
    ]

    def __init__(
        self,
        uses_axis_map: bool = False,
        valid_packed_dims: Optional[Set[PackedDim]] = None,
    ):
        self.uses_axis_map: bool = uses_axis_map
        self.valid_packed_dims = set()
        if valid_packed_dims is not None:
            self.valid_packed_dims = valid_packed_dims

    def valid_memory_layouts(self) -> Set[VkMemoryLayout]:
        """
        Derive the set of memory layouts supported by the texture implementation based
        on the valid packed dimensions.
        """
        layouts = set()

        if PackedDim.WIDTH in self.valid_packed_dims:
            layouts.add(VkMemoryLayout.TENSOR_WIDTH_PACKED)

        if PackedDim.HEIGHT in self.valid_packed_dims:
            layouts.add(VkMemoryLayout.TENSOR_HEIGHT_PACKED)

        if PackedDim.CHANNELS in self.valid_packed_dims:
            layouts.add(VkMemoryLayout.TENSOR_CHANNELS_PACKED)

        return layouts


class OpFeatures:
    __slots__ = [
        # None or TextureImplFeatures to specify implementation details of the texture
        # based operator implementation.
        "texture_impl",
        # bool indicating if the operator has a buffer based implementation.
        "buffer_impl",
        # bool indicating if the operator has a resize function, which allows it to
        # support dynamic shape tensors.
        "resize_fn",
        # Optimal
        "optimal_storage",
        "optimal_layout",
        # bool indicating if the operator handles its own prepacking. If this is True,
        # then the insert_prepack_nodes pass will not insert prepack nodes for the args
        # of the op.
        "handles_own_prepacking",
        # Optional dictionary to specify a custom function to calculate the required
        # image extents for a particular argument index.
        "skip_limits_check",
        # Optional check function used during partitioning to determine if a node's
        # inputs are supported by the operator implementation.
        "check_node_fn",
    ]

    def __init__(
        self,
        texture_impl: Optional[TextureImplFeatures] = None,
        buffer_impl: bool = False,
        resize_fn: bool = False,
        optimal_storage: Optional[VkStorageType] = None,
        optimal_layout: Optional[VkMemoryLayout] = None,
        handles_own_prepacking: bool = False,
        skip_limits_check: Optional[Set[int]] = None,
        check_node_fn: Optional[Callable] = None,
    ):
        self.texture_impl: Optional[TextureImplFeatures] = texture_impl
        self.buffer_impl: bool = buffer_impl
        self.resize_fn: bool = resize_fn
        self.optimal_storage: Optional[VkStorageType] = optimal_storage
        self.optimal_layout: Optional[VkMemoryLayout] = optimal_layout
        self.handles_own_prepacking: bool = handles_own_prepacking

        self.skip_limits_check: Set[int] = set()
        if skip_limits_check is not None:
            self.skip_limits_check = skip_limits_check

        self.check_node_fn: Callable = allow_node
        if check_node_fn is not None:
            self.check_node_fn = check_node_fn

    def propose_storage_type(self) -> Optional[VkStorageType]:
        """
        Propose a storage type that should be used for this operator. A proposal can be
        made if one of the following is true:
        1. The operator specifies an optimal storage type
        2. Only one storage type is supported.

        If both storage types are supported and no optimal storage type is specified,
        then None is returned to indicate that there is no preference in storage type.
        """
        if self.optimal_storage is not None:
            return self.optimal_storage

        if self.texture_impl is not None and not self.buffer_impl:
            return VkStorageType.TEXTURE_3D
        elif self.buffer_impl and self.texture_impl is None:
            return VkStorageType.BUFFER

        return None

    def supported_storage_types(self) -> Set[VkStorageType]:
        """
        Return the set of storage types supported by this operator.
        """
        storage_types = set()
        if self.texture_impl is not None:
            storage_types.add(VkStorageType.TEXTURE_3D)
        if self.buffer_impl:
            storage_types.add(VkStorageType.BUFFER)

        return storage_types

    def propose_memory_layout(self, storage: VkStorageType) -> Optional[VkMemoryLayout]:
        """
        Given a storage type as a precondition, propose a memory layout that should be
        used for this operator. A proposal can be made if one of the following is true:
        1. The operator specifies an optimal memory layout
        2. Only one memory layout is supported.

        If multiple memory layouts are supported and no optimal memory layout is
        specified then return None to indicate that the "best" memory layout for the
        operator is ambiguous.
        """
        if self.optimal_layout is not None:
            return self.optimal_layout

        if storage == VkStorageType.TEXTURE_3D:
            assert self.texture_impl is not None
            possible_layouts = self.texture_impl.valid_memory_layouts()
            if len(possible_layouts) == 1:
                return next(iter(possible_layouts))

        return None

    def supported_memory_layouts(self, storage: VkStorageType) -> Set[VkMemoryLayout]:
        """
        Return the set of memory layouts supported by this operator for a given storage
        type.
        """
        if storage == VkStorageType.TEXTURE_3D:
            assert self.texture_impl is not None
            return self.texture_impl.valid_memory_layouts()
        else:
            return all_memory_layouts


#######################
## Operator Registry ##
#######################

OpKey = Union[str, torch._ops.OpOverload, EdgeOpOverload]

vulkan_supported_ops: Dict[OpKey, OpFeatures] = {}


def update_features(aten_op):
    def features_decorator(fn: Callable):
        def update_features_impl(op: OpKey):
            if op in vulkan_supported_ops:
                raise RuntimeError(f"[Vulkan delegate] duplicate registration of {op}!")
            vulkan_supported_ops[op] = OpFeatures()
            vulkan_supported_ops[op] = fn(vulkan_supported_ops[op])

        if isinstance(aten_op, list):
            for op in aten_op:
                update_features_impl(op)
        else:
            update_features_impl(aten_op)

        return fn

    return features_decorator


@update_features(
    [
        operator.getitem,
        # Quantization related ops will be fused via graph passes
        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
        exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor,
        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
        exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
    ]
)
def register_ephemeral_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=True,
        valid_packed_dims=all_packed_dims,
    )
    features.buffer_impl = True
    features.resize_fn = True
    return features


@update_features(
    [
        exir_ops.edge.aten.add.Tensor,
        exir_ops.edge.aten.sub.Tensor,
        exir_ops.edge.aten.minimum.default,
        exir_ops.edge.aten.mul.Tensor,
        exir_ops.edge.aten.div.Tensor,
        exir_ops.edge.aten.div.Tensor_mode,
        exir_ops.edge.aten.pow.Tensor_Tensor,
    ]
)
def register_binary_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=True,
        valid_packed_dims=all_packed_dims,
    )
    features.resize_fn = True
    return features


@update_features(
    [
        exir_ops.edge.aten.abs.default,
        exir_ops.edge.aten.clamp.default,
        exir_ops.edge.aten.cos.default,
        exir_ops.edge.aten.exp.default,
        exir_ops.edge.aten.gelu.default,
        exir_ops.edge.aten.hardshrink.default,
        exir_ops.edge.aten.hardtanh.default,
        exir_ops.edge.aten.neg.default,
        exir_ops.edge.aten.relu.default,
        exir_ops.edge.aten.sigmoid.default,
        exir_ops.edge.aten.sin.default,
        exir_ops.edge.aten.sqrt.default,
        exir_ops.edge.aten.rsqrt.default,
        exir_ops.edge.aten.tanh.default,
    ]
)
def register_unary_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=True,
        valid_packed_dims=all_packed_dims,
    )
    features.buffer_impl = True
    features.resize_fn = True
    return features


@update_features(exir_ops.edge.aten._to_copy.default)
def register_to_copy_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=True,
        valid_packed_dims=all_packed_dims,
    )
    features.resize_fn = True

    def check_to_copy_node(node: torch.fx.Node) -> bool:
        float_dtypes = [torch.float16, torch.float32]

        if len(node.args) != 1:
            return False

        in_arg = node.args[0]
        if not isinstance(in_arg, torch.fx.Node):
            return False

        in_tensor = in_arg.meta.get("val", None)
        out_tensor = node.meta.get("val", None)

        if isinstance(in_tensor, FakeTensor) and isinstance(out_tensor, FakeTensor):
            if out_tensor.dtype in float_dtypes and in_tensor.dtype in float_dtypes:
                return True

        return False

    features.check_node_fn = check_to_copy_node

    return features


@update_features(
    [
        exir_ops.edge.aten.bmm.default,
        exir_ops.edge.aten.mm.default,
        exir_ops.edge.aten.addmm.default,
        exir_ops.edge.aten.linear.default,
    ]
)
def register_mm_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=True,
        valid_packed_dims={
            PackedDim.WIDTH,
            PackedDim.CHANNELS,
        },
    )
    features.buffer_impl = True
    features.resize_fn = True
    features.optimal_storage = VkStorageType.TEXTURE_3D
    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
    features.handles_own_prepacking = True
    return features


@update_features(exir_ops.edge.aten._weight_int8pack_mm.default)
def register_int8_mm_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=False,
        valid_packed_dims={PackedDim.WIDTH},
    )
    features.buffer_impl = True
    features.resize_fn = True
    features.optimal_storage = VkStorageType.TEXTURE_3D
    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
    features.handles_own_prepacking = True
    return features


@update_features(exir_ops.edge.et_vk.linear_weight_int4.default)
def register_int4_mm_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        uses_axis_map=False,
        valid_packed_dims={PackedDim.WIDTH},
    )
    features.resize_fn = True
    features.optimal_storage = VkStorageType.TEXTURE_3D
    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
    features.handles_own_prepacking = True
    return features


@update_features(
    [
        exir_ops.edge.aten._log_softmax.default,
        exir_ops.edge.aten._softmax.default,
    ]
)
def register_softmax_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims=all_packed_dims,
    )
    features.resize_fn = True
    return features


@update_features(
    [
        exir_ops.edge.aten.mean.dim,
        exir_ops.edge.aten.sum.dim_IntList,
        exir_ops.edge.aten.amax.default,
        exir_ops.edge.aten.amin.default,
    ]
)
def register_reduce_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims=all_packed_dims,
    )
    features.resize_fn = True

    def check_reduce_node(node: torch.fx.Node) -> bool:
        dim_list = node.args[1]
        if isinstance(dim_list, list) and len(dim_list) != 1:
            return False

        keepdim = node.args[2]
        if isinstance(keepdim, bool) and not keepdim:
            return False

        return True

    features.check_node_fn = check_reduce_node
    return features


@update_features(
    [
        exir_ops.edge.aten.avg_pool2d.default,
        exir_ops.edge.aten.max_pool2d_with_indices.default,
    ]
)
def register_2d_pool_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.CHANNELS},
    )
    features.resize_fn = True
    return features


@update_features(
    [
        exir_ops.edge.aten.convolution.default,
        exir_ops.edge.et_vk.conv_with_clamp.default,
    ]
)
def register_convolution_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.CHANNELS},
    )
    features.resize_fn = True
    features.optimal_storage = VkStorageType.TEXTURE_3D
    features.optimal_layout = VkMemoryLayout.TENSOR_CHANNELS_PACKED
    features.handles_own_prepacking = True
    features.skip_limits_check = {1, 2}
    return features


@update_features("llama::sdpa_with_kv_cache")
def register_sdpa_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.WIDTH},
    )
    features.resize_fn = True
    features.optimal_storage = VkStorageType.TEXTURE_3D
    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
    features.handles_own_prepacking = True
    return features


@update_features(exir_ops.edge.et_vk.apply_rotary_emb.default)
def register_rotary_emb_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.WIDTH},
    )
    features.resize_fn = True
    return features


@update_features(exir_ops.edge.aten.view_copy.default)
def register_view_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims=all_packed_dims,
    )
    features.resize_fn = True
    return features


# Ops ported from PyTorch Vulkan backend. These ops commonly support channels
# packed tensors only and do not have a resize function.
@update_features(
    [
        # Shape Manipulation
        exir_ops.edge.aten.squeeze_copy.dims,
        exir_ops.edge.aten.unsqueeze_copy.default,
        exir_ops.edge.aten.permute_copy.default,
        exir_ops.edge.aten.t_copy.default,
        # Indexing and lookup
        exir_ops.edge.aten.flip.default,
        exir_ops.edge.aten.index_select.default,
        exir_ops.edge.aten.select_copy.int,
        exir_ops.edge.aten.slice_copy.Tensor,
        # Tensor combination
        exir_ops.edge.aten.cat.default,
        exir_ops.edge.aten.split_with_sizes_copy.default,
        exir_ops.edge.aten.split.Tensor,
        exir_ops.edge.aten.repeat.default,
        # Tensor creation
        exir_ops.edge.aten.arange.start_step,
        exir_ops.edge.aten.clone.default,
        exir_ops.edge.aten.constant_pad_nd.default,
        exir_ops.edge.aten.full.default,
        exir_ops.edge.aten.full_like.default,
        exir_ops.edge.aten.ones.default,
        exir_ops.edge.aten.ones_like.default,
        exir_ops.edge.aten.upsample_nearest2d.vec,
        exir_ops.edge.aten.zeros.default,
        exir_ops.edge.aten.zeros_like.default,
        exir_ops.edge.et_vk.grid_priors.default,
    ]
)
def register_ported_op(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.CHANNELS},
    )
    return features


# Ported ops that support their own prepacking.
@update_features(
    [
        exir_ops.edge.aten.embedding.default,
        exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
        exir_ops.edge.aten.native_layer_norm.default,
    ]
)
def register_ported_ops_with_prepacking(features: OpFeatures):
    features.texture_impl = TextureImplFeatures(
        valid_packed_dims={PackedDim.CHANNELS},
    )
    features.handles_own_prepacking = True
    return features


#######################
## Utility functions ##
#######################


def has_impl(target: OpKey) -> bool:
    if not isinstance(target, str):
        if target not in vulkan_supported_ops:
            return target.name() in vulkan_supported_ops
        return target in vulkan_supported_ops
    else:
        return target in vulkan_supported_ops


def get_op_features(target: OpKey) -> OpFeatures:
    if not isinstance(target, str):
        if target not in vulkan_supported_ops:
            # Try the op's name
            return vulkan_supported_ops[target.name()]

        return vulkan_supported_ops[target]
    else:
        return vulkan_supported_ops[target]


def handles_own_prepacking(target: OpKey) -> bool:
    return get_op_features(target).handles_own_prepacking