op_registry.py

# 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 Any, Callable, Dict, List, Optional, Tuple, Union

import executorch.backends.vulkan.custom_ops_lib  # noqa
import executorch.backends.vulkan.utils as utils
import torch
from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.dialects.edge._ops import EdgeOpOverload

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


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


class OpFeatures:
    __slots__ = [
        # Dtype-related slots:
        "inputs_dtypes",  # DtypeSetList for input tensor dtype constraints
        "outputs_dtypes",  # DtypeSetList for output tensor dtype constraints
        # Storage-related slots:
        # Sets of possible (storage types, memory layouts) to use for the input tensor(s)
        "inputs_storage",
        # Sets of possible (storage types, memory layouts) to use for the output tensor(s)
        "outputs_storage",
        # bool indicating if the operator has a resize function, which allows it to
        # support models with dynamic shape
        "supports_resize",
        # bool indicating if the operator supports tensors with more than 4 dimensions
        "supports_highdim",
        # 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.
        "supports_prepacking",
        # Optional check function used during partitioning to determine if a node's
        # inputs are supported by the operator implementation.
        "are_node_inputs_supported_fn",
        # Optional function to determine valid representation sets for input and outputs
        # once a node's actual inputs are known.
        "pick_io_storage_fn",
    ]

    def __init__(
        self,
        inputs_dtypes: Optional[Union[utils.DtypeSet, List[utils.DtypeSet]]] = None,
        outputs_dtypes: Optional[Union[utils.DtypeSet, List[utils.DtypeSet]]] = None,
        inputs_storage: Optional[
            Union[utils.TensorRepSet, List[utils.TensorRepSet]]
        ] = None,
        outputs_storage: Optional[
            Union[utils.TensorRepSet, List[utils.TensorRepSet]]
        ] = None,
        supports_resize: bool = False,
        supports_highdim: bool = False,
        supports_prepacking: bool = False,
        are_node_inputs_supported_fn: Optional[Callable] = allow_node,
        pick_io_storage_fn: Optional[Callable] = None,
    ):
        # Dtype initialization
        self.inputs_dtypes: utils.DtypeSetList = utils.DtypeSetList(
            inputs_dtypes if inputs_dtypes is not None else utils.ALL_T
        )
        self.outputs_dtypes: utils.DtypeSetList = utils.DtypeSetList(
            outputs_dtypes if outputs_dtypes is not None else self.inputs_dtypes[0]
        )

        # Storage initialization
        self.inputs_storage: utils.TensorRepSetList = utils.TensorRepSetList(
            inputs_storage if inputs_storage is not None else []
        )
        self.outputs_storage: utils.TensorRepSetList = utils.TensorRepSetList(
            outputs_storage if outputs_storage is not None else []
        )

        # If output storage is not set, assume that it is derived from the first input
        if self.outputs_storage.any_is_empty():
            self.outputs_storage = utils.TensorRepSetList(self.inputs_storage[0])

        self.supports_resize = supports_resize
        self.supports_highdim = supports_highdim
        self.supports_prepacking = supports_prepacking

        self.are_node_inputs_supported_fn = are_node_inputs_supported_fn
        self.pick_io_storage_fn = pick_io_storage_fn

    def check_dtypes(self, node: torch.fx.Node) -> Tuple[bool, str]:
        """
        Check if all tensor inputs/outputs have dtypes supported by this operator.
        Returns (is_valid, reason_string).
        """
        return utils.check_node_dtypes(
            node,
            self.inputs_dtypes,
            self.outputs_dtypes,
        )

    def make_op_repsets(
        self,
        op_node: torch.fx.Node,
        texture_limits: utils.ImageExtents = utils.DEFAULT_TEXTURE_LIMITS,
    ) -> utils.OpRepSets:
        inputs_storage = self.inputs_storage
        outputs_storage = self.outputs_storage
        if self.pick_io_storage_fn is not None:
            i_storage, o_storage = self.pick_io_storage_fn(op_node)
            inputs_storage = utils.TensorRepSetList(i_storage)
            outputs_storage = utils.TensorRepSetList(o_storage)

        return utils.OpRepSets(inputs_storage, outputs_storage, op_node, texture_limits)


#######################
## 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] = fn()

        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


# =============================================================================
# Ephemeral Operators (no C++ dispatch - handled symbolically)
# =============================================================================


@update_features(
    [
        operator.getitem,
        # Symbolic integer ops (SymIntOps.cpp - symbolic handling)
        torch.ops.aten.sym_size.int,
        operator.add,
        operator.sub,
        operator.lt,
        operator.gt,
        operator.ge,
        operator.le,
        operator.eq,
        # Guard and assert ops
        torch.ops.aten._assert_scalar.default,
        torch.ops.aten.sym_constrain_range_for_size.default,
        # copy.default is a no-op when src dtype matches dst dtype; removed by
        # RemoveRedundantOpsTransform before execution.
        exir_ops.edge.aten.copy.default,
    ]
)
def register_ephemeral_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        supports_resize=True,
    )


# =============================================================================
# UnaryOp.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.abs.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,
        exir_ops.edge.aten.round.default,
        exir_ops.edge.aten.leaky_relu.default,
    ]
)
def register_unaryop_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


@update_features(exir_ops.edge.aten.clamp.default)
def register_clamp():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
    )


# =============================================================================
# BinaryOp.cpp
# =============================================================================


@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_binaryop_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
        supports_highdim=True,
    )


@update_features(
    [
        exir_ops.edge.aten.eq.Tensor,
        exir_ops.edge.aten.lt.Tensor,
        exir_ops.edge.aten.le.Tensor,
        exir_ops.edge.aten.gt.Tensor,
        exir_ops.edge.aten.ge.Tensor,
    ]
)
def register_comparison_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        outputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# BinaryOp.cpp (bitwise)
# =============================================================================


@update_features(exir_ops.edge.aten.bitwise_and.Tensor)
def register_bitwise_and():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# BinaryScalarOp.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.pow.Tensor_Scalar)
def register_pow_tensor_scalar():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# ToCopy.cpp
# =============================================================================


@update_features(exir_ops.edge.aten._to_copy.default)
def register_to_copy():
    def check_to_copy_node(node: torch.fx.Node) -> bool:
        # Only single-arg _to_copy is supported
        return len(node.args) == 1

    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        outputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
        are_node_inputs_supported_fn=check_to_copy_node,
    )


# =============================================================================
# Softmax.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten._log_softmax.default,
        exir_ops.edge.aten._softmax.default,
    ]
)
def register_softmax_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


# =============================================================================
# MatMul.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.bmm.default,
        exir_ops.edge.aten.mm.default,
    ]
)
def register_matmul_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_prepacking=True,
    )


# =============================================================================
# Linear.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.addmm.default,
        exir_ops.edge.aten.linear.default,
    ]
)
def register_linear_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_prepacking=True,
    )


# =============================================================================
# QuantizedLinearQCSNW.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten._weight_int8pack_mm.default,
        exir_ops.edge.et_vk.linear_qcs4w.default,
    ]
)
def register_quantizedlinearqcsnw_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_prepacking=True,
    )


# =============================================================================
# QuantizedLinear.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.et_vk.linear_q8ta_q8csw.default,
        exir_ops.edge.et_vk.linear_q4gsw.default,
    ]
)
def register_quantizedlinear_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
    )


@update_features(exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default)
def register_linear_dq8ca_q4gsw():
    return OpFeatures(
        inputs_storage=[
            utils.CONTIGUOUS_ANY,  # input
            utils.WIDTH_PACKED_TEXTURE,  # input_scale
            utils.WIDTH_PACKED_TEXTURE,  # input_zero_point
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # group_size (scalar)
            utils.NO_STORAGE,  # bias (prepacked)
        ],
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
    )


# =============================================================================
# QuantizeDequantize.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
        exir_ops.edge.quantized_decomposed.quantize_per_token.default,
        exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
        exir_ops.edge.quantized_decomposed.dequantize_per_token.default,
    ]
)
def register_quantizedequantize_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_BUFFER,
        supports_resize=True,
    )


@update_features(
    [
        exir_ops.edge.torchao.quantize_affine.default,
        exir_ops.edge.torchao.dequantize_affine.default,
    ]
)
def register_torchao_quantize_dequantize():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_BUFFER,
        supports_resize=True,
    )


# =============================================================================
# Q8taQuantizeDequantize.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
        exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor,
    ]
)
def register_quantize_per_tensor():
    return OpFeatures(
        inputs_storage=[
            utils.CHANNELS_PACKED_TEXTURE_OR_CONTIGUOUS_BUFFER,
        ],
        outputs_storage=[
            utils.PACKED_INT8_BUFFER,
        ],
        supports_highdim=True,
    )


@update_features(
    [
        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
    ]
)
def register_dequantize_per_tensor():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_BUFFER,
        ],
        outputs_storage=[
            utils.CHANNELS_PACKED_TEXTURE_OR_CONTIGUOUS_BUFFER,
        ],
        supports_highdim=True,
    )


# =============================================================================
# ChooseQParams.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.quantized_decomposed.choose_qparams.tensor,
        exir_ops.edge.quantized_decomposed.choose_qparams_per_token_asymmetric.default,
    ]
)
def register_chooseqparams_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_BUFFER,
        supports_resize=True,
    )


@update_features(exir_ops.edge.torchao.choose_qparams_affine.default)
def register_torchao_choose_qparams_affine():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        outputs_storage=[
            utils.WIDTH_PACKED_TEXTURE,  # scales
            utils.WIDTH_PACKED_TEXTURE,  # zero_points
        ],
        supports_resize=True,
    )


# =============================================================================
# Q8taBinary.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.q8ta_add.default)
def register_q8ta_add():
    return OpFeatures(
        inputs_storage=utils.PACKED_INT8_BUFFER,
        supports_resize=False,
    )


# =============================================================================
# Q8taUnary.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.q8ta_relu.default)
def register_q8ta_relu():
    return OpFeatures(
        inputs_storage=utils.PACKED_INT8_BUFFER,
        supports_resize=True,
    )


# =============================================================================
# =============================================================================


def get_dims_reduced(node: torch.fx.Node) -> Union[int, List[int]]:
    ndim = utils.ndim_of(node.args[0])
    assert ndim is not None
    dims_reduced = None
    if len(node.args) >= 2:
        dims_reduced = node.args[1]

    # If dim_list is None, return a list containing all the dims of the tensor
    if dims_reduced is None:
        dims_reduced = list(range(ndim))

        # Special case for reducing tensors with shape [1, N] - this is equivalent to
        # reducing the last dim.
        if utils.is_unsqueezed_vector(node) and ndim == 2:
            dims_reduced = 1

    if isinstance(dims_reduced, (list, tuple)) and len(dims_reduced) == 1:
        dims_reduced = dims_reduced[0]

    assert isinstance(dims_reduced, (int, list, tuple))
    return utils.normalize_dims(dims_reduced, ndim)


def get_keepdim_setting(node: torch.fx.Node) -> bool:
    for arg in node.args:
        if isinstance(arg, bool):
            return arg

    # Assume false by default
    return False


def is_reduce_node_supported_by_per_row_impl(node: torch.fx.Node) -> bool:
    """
    Checks if a reduction node is supported by the Vulkan backend's reduce per row
    special case implementation.
    """
    input_ndim = utils.ndim_of(node.args[0])
    assert input_ndim is not None
    dims_reduced = get_dims_reduced(node)

    return dims_reduced == input_ndim - 1


def is_reduce_node_supported_by_general_impl(node: torch.fx.Node) -> bool:
    dims_reduced = get_dims_reduced(node)
    # Only 1D and 2D reductions are supported at the moment.
    if isinstance(dims_reduced, (list, tuple)) and len(dims_reduced) > 2:
        return False

    keepdim = get_keepdim_setting(node)
    # keepdim = False is not supported yet for general implementation
    if isinstance(keepdim, bool) and not keepdim:
        return False

    return True


def is_reduce_node_supported(node: torch.fx.Node) -> bool:
    return is_reduce_node_supported_by_per_row_impl(
        node
    ) or is_reduce_node_supported_by_general_impl(node)


def pick_storage_for_reduce(node: torch.fx.Node):
    inputs_storage = utils.NO_STORAGE
    outputs_storage = utils.NO_STORAGE

    ndim = utils.ndim_of(node.args[0])
    dim_list = get_dims_reduced(node)

    if is_reduce_node_supported_by_general_impl(node):
        inputs_storage = inputs_storage.make_union(utils.ANY_TEXTURE)
        outputs_storage = inputs_storage

    # For 1D reductions of the last dim, a special reduce per row case is implemented
    # for buffer backed tensors.
    if is_reduce_node_supported_by_per_row_impl(node):
        inputs_storage = inputs_storage.make_union(utils.CONTIGUOUS_BUFFER)
        outputs_storage = inputs_storage
        return inputs_storage, outputs_storage

    # For 2D reductions, the packed dimension cannot be one of the reduced dims
    if isinstance(dim_list, (list, tuple)) and len(dim_list) == 2:
        # pyre-ignore[6]
        reduce_dim1_whcn = utils.nchw_dim_to_whcn_dim(dim_list[0], ndim)
        # pyre-ignore[6]
        reduce_dim2_whcn = utils.nchw_dim_to_whcn_dim(dim_list[1], ndim)

        possible_packed_dims = {0, 1, 2}
        possible_packed_dims.discard(reduce_dim1_whcn)
        possible_packed_dims.discard(reduce_dim2_whcn)

        packed_dim = possible_packed_dims.pop()
        assert packed_dim in [0, 1, 2]

        if packed_dim == 0:
            inputs_storage = utils.WIDTH_PACKED_TEXTURE
            outputs_storage = utils.WIDTH_PACKED_TEXTURE
        elif packed_dim == 1:
            inputs_storage = utils.HEIGHT_PACKED_TEXTURE
            outputs_storage = utils.HEIGHT_PACKED_TEXTURE
        else:
            inputs_storage = utils.CHANNELS_PACKED_TEXTURE
            outputs_storage = utils.CHANNELS_PACKED_TEXTURE

    return inputs_storage, outputs_storage


@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_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_highdim=True,
        are_node_inputs_supported_fn=is_reduce_node_supported,
        pick_io_storage_fn=pick_storage_for_reduce,
    )


# =============================================================================
# ArgReduce.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.argmax.default,
        exir_ops.edge.aten.argmin.default,
    ]
)
def register_argreduce_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_T,
        outputs_dtypes=utils.INT_T,
        supports_resize=True,
        supports_highdim=True,
        are_node_inputs_supported_fn=is_reduce_node_supported,
        pick_io_storage_fn=pick_storage_for_reduce,
    )


# =============================================================================
# Pool.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.avg_pool2d.default,
        exir_ops.edge.aten.max_pool2d.default,
    ]
)
def register_pool_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


@update_features(exir_ops.edge.aten.max_pool2d_with_indices.default)
def register_max_pool2d_with_indices():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
        outputs_dtypes=[utils.FP_T, utils.INT_T],
        supports_resize=True,
    )


# =============================================================================
# Convolution.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.convolution.default,
        exir_ops.edge.et_vk.conv_with_clamp.default,
    ]
)
def register_convolution_cpp_ops():
    def check_conv_node(node: torch.fx.Node) -> bool:
        x = node.args[0]
        assert isinstance(x, torch.fx.Node)
        x_shape = x.meta["val"].size()
        # 4-D input implies 2D convolution
        if len(x_shape) == 4:
            batches = x.meta["val"].size()[0]
            if batches != 1:
                return False
        # 3-D input implies 1D convolution
        if len(x_shape) == 3:
            transpose = node.args[6]
            # Transposed 1D convolution is not supported yet
            if transpose:
                return False

        return True

    return OpFeatures(
        inputs_storage=[
            utils.CHANNELS_PACKED_TEXTURE,  # input
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # stride (non tensor)
            utils.NO_STORAGE,  # padding (non tensor)
            utils.NO_STORAGE,  # dilation (non tensor)
            utils.NO_STORAGE,  # transposed (non tensor)
            utils.NO_STORAGE,  # output_padding (non tensor)
            utils.NO_STORAGE,  # groups (non tensor)
            utils.NO_STORAGE,  # output_min (non tensor)
            utils.NO_STORAGE,  # output_max (non tensor)
        ],
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_prepacking=True,
        are_node_inputs_supported_fn=check_conv_node,
    )


# =============================================================================
# Q8taConv2d*.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.et_vk.q8ta_conv2d_pw.default,
    ]
)
def register_q8ta_conv_pw_op():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_4W4C_BUFFER,  # input
            utils.NO_STORAGE,  # input_scale (non tensor)
            utils.NO_STORAGE,  # input_zero_point (non tensor)
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # output_scale (non tensor)
            utils.NO_STORAGE,  # output_zero_point (non tensor)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # kernel_size (non tensor)
            utils.NO_STORAGE,  # stride (non tensor)
            utils.NO_STORAGE,  # padding (non tensor)
            utils.NO_STORAGE,  # dilation (non tensor)
            utils.NO_STORAGE,  # groups (non tensor)
            utils.NO_STORAGE,  # original OC count (non tensor)
        ],
        outputs_storage=[
            utils.PACKED_INT8_CHANNELS_PACKED_BUFFER,
        ],
        supports_resize=False,
        supports_prepacking=True,
    )


@update_features(
    [
        exir_ops.edge.et_vk.q8ta_conv2d.default,
        exir_ops.edge.et_vk.q8ta_conv2d_dw.default,
    ]
)
def register_q8ta_conv2d_ops():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_CONV2D_BUFFER,  # input
            utils.NO_STORAGE,  # input_scale (non tensor)
            utils.NO_STORAGE,  # input_zero_point (non tensor)
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # output_scale (non tensor)
            utils.NO_STORAGE,  # output_zero_point (non tensor)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # kernel_size (non tensor)
            utils.NO_STORAGE,  # stride (non tensor)
            utils.NO_STORAGE,  # padding (non tensor)
            utils.NO_STORAGE,  # dilation (non tensor)
            utils.NO_STORAGE,  # groups (non tensor)
            utils.NO_STORAGE,  # original OC count (non tensor)
        ],
        outputs_storage=[
            utils.PACKED_INT8_CHANNELS_PACKED_BUFFER,
        ],
        supports_resize=False,
        supports_prepacking=True,
    )


@update_features(
    [
        exir_ops.edge.et_vk.q8ta_conv2d_transposed.default,
    ]
)
def register_q8ta_conv2d_transposed_op():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_CONV2D_BUFFER,  # input
            utils.NO_STORAGE,  # input_scale (non tensor)
            utils.NO_STORAGE,  # input_zero_point (non tensor)
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # output_scale (non tensor)
            utils.NO_STORAGE,  # output_zero_point (non tensor)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # kernel_size (non tensor)
            utils.NO_STORAGE,  # stride (non tensor)
            utils.NO_STORAGE,  # padding (non tensor)
            utils.NO_STORAGE,  # output_padding (non tensor)
            utils.NO_STORAGE,  # dilation (non tensor)
            utils.NO_STORAGE,  # groups (non tensor)
            utils.NO_STORAGE,  # activation (non tensor)
        ],
        outputs_storage=[
            utils.PACKED_INT8_CHANNELS_PACKED_BUFFER,
        ],
        supports_resize=False,
        supports_prepacking=True,
    )


# =============================================================================
# Q8taLinear.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.q8ta_linear.default)
def register_q8ta_linear():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_4H4W_BUFFER,  # input
            utils.NO_STORAGE,  # input_scale (non tensor)
            utils.NO_STORAGE,  # input_zero_point (non tensor)
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # output_scale (non tensor)
            utils.NO_STORAGE,  # output_zero_point (non tensor)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # activation (non tensor)
        ],
        outputs_storage=[
            utils.PACKED_INT8_4H4W_BUFFER,
        ],
        supports_resize=False,
        supports_prepacking=True,
    )


@update_features(exir_ops.edge.et_vk.q8ta_linear_gemv.default)
def register_q8ta_linear_gemv():
    return OpFeatures(
        inputs_storage=[
            utils.PACKED_INT8_4W_BUFFER,  # input
            utils.NO_STORAGE,  # input_scale (non tensor)
            utils.NO_STORAGE,  # input_zero_point (non tensor)
            utils.NO_STORAGE,  # weight (prepacked)
            utils.NO_STORAGE,  # weight_sums (prepacked)
            utils.NO_STORAGE,  # weight_scales (prepacked)
            utils.NO_STORAGE,  # output_scale (non tensor)
            utils.NO_STORAGE,  # output_zero_point (non tensor)
            utils.NO_STORAGE,  # bias (prepacked)
            utils.NO_STORAGE,  # activation (non tensor)
        ],
        outputs_storage=[
            utils.PACKED_INT8_4W_BUFFER,
        ],
        supports_resize=False,
        supports_prepacking=True,
    )


# =============================================================================
# SDPA.cpp
# =============================================================================


@update_features("llama::sdpa_with_kv_cache")
def register_sdpa_with_kv_cache():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_prepacking=True,
    )


@update_features(
    [
        "llama::update_cache",
        "llama::custom_sdpa",
    ]
)
def register_sdpa_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


# =============================================================================
# RotaryEmbedding.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.apply_rotary_emb.default)
def register_apply_rotary_emb():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Permute.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.permute.default)
def register_permute():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
    )


@update_features(exir_ops.edge.aten.permute_copy.default)
def register_permute_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# View.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.view_copy.default)
def register_view_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


@update_features(exir_ops.edge.dim_order_ops._to_dim_order_copy.default)
def register_to_dim_order_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_BUFFER,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Squeeze.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.squeeze_copy.dims)
def register_squeeze_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Unsqueeze.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.unsqueeze_copy.default)
def register_unsqueeze_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Clone.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.clone.default)
def register_clone():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


@update_features(exir_ops.edge.dim_order_ops._clone_dim_order.default)
def register_clone_dim_order():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Gather.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.gather.default)
def register_gather():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Expand.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.expand_copy.default)
def register_expand_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_BUFFER,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=False,
        supports_highdim=True,
    )


# =============================================================================
# Concat.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.cat.default)
def register_cat():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
    )


# =============================================================================
# Select.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.select_copy.int)
def register_select_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Slice.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.slice_copy.Tensor)
def register_slice_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Split.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.split_with_sizes_copy.default)
def register_split_with_sizes_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Transpose.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.t_copy.default)
def register_t_copy():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
    )


# =============================================================================
# Flip.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.flip.default)
def register_flip():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
    )


# =============================================================================
# IndexSelect.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.index_select.default)
def register_index_select():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
    )


# =============================================================================
# Where.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.where.self)
def register_where():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=[utils.BOOL_T, utils.FP_T, utils.FP_T],
        outputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


# =============================================================================
# IndexTensor.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.index.Tensor)
def register_index_tensor():
    def check_index_tensor_node(node: torch.fx.Node) -> bool:
        self_arg = node.args[0]
        indices = node.args[1]

        # Only support 1D self tensor
        if not isinstance(self_arg, torch.fx.Node):
            return False
        self_val = self_arg.meta.get("val", None)
        if self_val is None:
            return False
        if len(self_val.size()) != 1:
            return False

        # Only support exactly one non-None index tensor
        if not isinstance(indices, (list, tuple)):
            return False
        non_none = [idx for idx in indices if idx is not None]
        if len(non_none) != 1:
            return False

        return True

    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
        are_node_inputs_supported_fn=check_index_tensor_node,
    )


# =============================================================================
# Arange.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.arange.start_step)
def register_arange():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
    )


# =============================================================================
# Pad.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.constant_pad_nd.default)
def register_constant_pad_nd():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
    )


# =============================================================================
# Full.cpp
# =============================================================================


@update_features(
    [
        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.zeros.default,
        exir_ops.edge.aten.zeros_like.default,
    ]
)
def register_full_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
    )


# =============================================================================
# ScalarTensor.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.scalar_tensor.default)
def register_scalar_tensor():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_T,
    )


# =============================================================================
# Upsample.cpp
# =============================================================================


@update_features(
    [
        exir_ops.edge.aten.upsample_nearest2d.vec,
        exir_ops.edge.aten.upsample_bilinear2d.vec,
    ]
)
def register_upsample_cpp_ops():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
    )


# =============================================================================
# GridPriors.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.grid_priors.default)
def register_grid_priors():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
    )


# =============================================================================
# Repeat.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.repeat.default)
def register_repeat():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
    )


# =============================================================================
# Embedding.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.embedding.default)
def register_embedding():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=[utils.FP_T, utils.INT_T],
        supports_prepacking=True,
        supports_resize=True,
    )


# =============================================================================
# BatchNorm.cpp
# =============================================================================


@update_features(exir_ops.edge.aten._native_batch_norm_legit_no_training.default)
def register_native_batch_norm_legit_no_training():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
        supports_resize=True,
    )


# =============================================================================
# GroupNorm.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.native_group_norm.default)
def register_native_group_norm():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_T,
        outputs_storage=[
            utils.CHANNELS_PACKED_TEXTURE,
            utils.CONTIGUOUS_BUFFER,
            utils.CONTIGUOUS_BUFFER,
        ],
        supports_prepacking=True,
    )


# =============================================================================
# NativeLayerNorm.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.native_layer_norm.default)
def register_native_layer_norm():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
        supports_resize=True,
    )


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


def has_impl(target: Any) -> 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: Any) -> 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).supports_prepacking