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 math
import operator
from typing import AbstractSet, 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",
        # Like are_node_inputs_supported_fn, for ops whose support depends on which
        # inputs are constants. Called with the node and the set of nodes holding
        # constant data (parameters, unmutated buffers, lifted constants).
        "are_node_inputs_supported_with_constants_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,
        are_node_inputs_supported_with_constants_fn: Optional[Callable] = None,
        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.are_node_inputs_supported_with_constants_fn = (
            are_node_inputs_supported_with_constants_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.floordiv,
        operator.mul,
        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,
    ]
)
def register_ephemeral_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        supports_resize=True,
    )


def _check_copy_is_noop(node: torch.fx.Node) -> bool:
    """Only support copy.default when it's a no-op (same dtype and shape)."""
    src = node.args[1]
    if not isinstance(src, torch.fx.Node):
        return False
    src_val = src.meta.get("val")
    dst_val = node.meta.get("val")
    if src_val is None or dst_val is None:
        return False
    return src_val.dtype == dst_val.dtype and src_val.shape == dst_val.shape


@update_features(
    [
        exir_ops.edge.aten.copy.default,
    ]
)
def register_copy_op():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        supports_resize=True,
        are_node_inputs_supported_fn=_check_copy_is_noop,
    )


# =============================================================================
# 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.hardswish.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,
        exir_ops.edge.aten.log10.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,
    )


# =============================================================================
# binary/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.et_vk.swiglu.default)
def register_swiglu():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_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,
    )


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


@update_features(
    [
        exir_ops.edge.aten.bitwise_and.Tensor,
        exir_ops.edge.aten.bitwise_or.Tensor,
        exir_ops.edge.aten.bitwise_not.default,
        exir_ops.edge.aten.logical_and.default,
        exir_ops.edge.aten.logical_or.default,
    ]
)
def register_bool_binary_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# binary/BinaryScalarOp.cpp
# =============================================================================


def is_scalar_value_supported(value: Any, dtype: torch.dtype) -> bool:
    if type(value) not in (bool, int, float):
        return False
    if isinstance(value, float) and math.isnan(value):
        return False
    if dtype in utils.INT_T:
        return -(2**31) <= value <= 2**31 - 1
    return True


@update_features(
    [
        exir_ops.edge.aten.pow.Tensor_Scalar,
        exir_ops.edge.aten.mul.Scalar,
    ]
)
def register_binary_scalar_ops():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        supports_highdim=True,
        are_node_inputs_supported_fn=lambda node: is_scalar_value_supported(
            node.args[1], node.meta["val"].dtype
        ),
    )


@update_features(exir_ops.edge.aten.eq.Scalar)
def register_eq_scalar():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_T,
        outputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
        are_node_inputs_supported_fn=lambda node: is_scalar_value_supported(
            node.args[1], node.args[0].meta["val"].dtype
        ),
    )


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


@update_features(exir_ops.edge.aten._to_copy.default)
def register_to_copy():
    def pick_to_copy_storage(
        node: torch.fx.Node,
    ) -> Tuple[utils.TensorRepSet, utils.TensorRepSet]:
        in_dtype = node.args[0].meta["val"].dtype  # type: ignore[union-attr]
        out_dtype = node.meta["val"].dtype
        fp_types = {torch.float16, torch.float32}
        if in_dtype in fp_types and out_dtype in fp_types:
            return utils.ANY_STORAGE, utils.ANY_STORAGE
        return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER

    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        outputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        pick_io_storage_fn=pick_to_copy_storage,
    )


# =============================================================================
# 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_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


# =============================================================================
# gemm/matmul/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,
    )


# =============================================================================
# gemm/linear/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,
    )


# =============================================================================
# gemm/linear/qw/{q8csw,legacy}/*.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,
    )


# =============================================================================
# gemm/linear/qw/{q4gsw,dq8ca_q4gsw,legacy}/**/*.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_q4gsw_backward.default)
def register_linear_q4gsw_backward():
    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,
    )


# =============================================================================
# binary/q8ta/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,
    )


# =============================================================================
# Q8taPixelShuffle.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.q8ta_pixel_shuffle.default)
def register_q8ta_pixel_shuffle():
    # The fused kernel is restricted to the channels-packed family
    # (PACKED_INT8_4W4C, PACKED_INT8_4C1W, PACKED_INT8_CONV2D), all of which
    # share packed_dim=C. See add_q8ta_pixel_shuffle_node in Q8taPixelShuffle.cpp
    # for the runtime assertion. The surrounding q8ta_conv2d ops produce
    # PACKED_INT8_4W4C on this model, so the partitioner can route through this
    # op without inserting layout-transition q8ta_clone dispatches.
    return OpFeatures(
        inputs_storage=utils.PACKED_INT8_CHANNELS_PACKED_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 = (
        []
        if node.target
        in (exir_ops.edge.aten.amax.default, exir_ops.edge.aten.amin.default)
        else 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 not 1 <= len(dims_reduced) <= 2:
        return False

    # any.dim can repack the reduced texture after removing the reduction axis.
    if not get_keepdim_setting(node) and node.target != exir_ops.edge.aten.any.dim:
        return False

    if utils.ndim_of(node.args[0]) == 4:
        dims = [dims_reduced] if isinstance(dims_reduced, int) else dims_reduced
        # Textures fold batch into channels; neither axis can be reduced across batches.
        if 0 in dims or (
            1 in dims and utils.upper_bound_size(node.args[0].meta["val"].shape[0]) != 1
        ):
            return False

    return True


def is_reduce_node_supported(node: torch.fx.Node) -> bool:
    if (
        node.target in (exir_ops.edge.aten.sum.dim_IntList, exir_ops.edge.aten.mean.dim)
        and (len(node.args) < 2 or node.args[1] is None)
        and utils.ndim_of(node.args[0]) != 1
    ):
        return False
    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,
    )


@update_features(exir_ops.edge.aten.any.dim)
def register_any_dim():
    return OpFeatures(
        inputs_storage=utils.ANY_TEXTURE,
        inputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
        are_node_inputs_supported_fn=lambda node: (
            utils.ndim_of(node.args[0]) > 0 and is_reduce_node_supported(node)
        ),
        pick_io_storage_fn=pick_storage_for_reduce,
    )


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


def is_argreduce_node_supported(node: torch.fx.Node) -> bool:
    ndim = utils.ndim_of(node.args[0])
    assert ndim is not None
    dim = node.args[1] if len(node.args) > 1 else None
    if dim is None:
        return ndim == 1
    return ndim > 0 and utils.normalize_dims(dim, ndim) == ndim - 1


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


# =============================================================================
# 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/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]
            # There is no transposed conv1d shader. Conv1dAsConv2dPass lowers
            # the ones it can rewrite onto the conv2d transposed path; accept
            # exactly those.
            if transpose and not utils.conv1d_transposed_is_rewritable(node):
                return False

        return True

    def pick_conv_storage(
        node: torch.fx.Node,
    ) -> Tuple[List[utils.TensorRepSet], utils.TensorRepSet]:
        x = node.args[0]
        assert isinstance(x, torch.fx.Node)
        x_shape = x.meta["val"].size()

        # Default: channels-packed texture (conv2d and fallback conv1d)
        input_storage = utils.CHANNELS_PACKED_TEXTURE
        output_storage = utils.CHANNELS_PACKED_TEXTURE

        if len(x_shape) == 3:
            # Conv1d: check if we can use height-packed
            weight = node.args[1]
            assert isinstance(weight, torch.fx.Node)
            w_shape = weight.meta["val"].size()
            groups = node.args[8]

            c_in = x_shape[1]
            c_out = w_shape[0]
            kernel_size = w_shape[2]

            is_pointwise = kernel_size == 1
            is_depthwise = (
                isinstance(groups, int)
                and groups == c_in
                and c_out == c_in
                and w_shape[1] == 1
            )
            if is_pointwise or is_depthwise:
                input_storage = utils.HEIGHT_PACKED_TEXTURE
                output_storage = utils.HEIGHT_PACKED_TEXTURE

        # Build per-input storage list. The convolution op has variable args:
        # aten.convolution.default: input, weight, bias, stride, padding,
        #   dilation, transposed, output_padding, groups
        # et_vk.conv_with_clamp.default: + output_min, output_max
        # All args after input are NO_STORAGE (prepacked or non-tensor)
        inputs = [input_storage] + [utils.NO_STORAGE] * 10
        return inputs, output_storage

    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,
        pick_io_storage_fn=pick_conv_storage,
    )


# =============================================================================
# convolution/conv2d/q8ta/**/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,
    )


# =============================================================================
# gemm/linear/q8ta/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,
    )


# =============================================================================
# SDPA.cpp (fused SDPA entry point)
# =============================================================================


@update_features("et_vk::sdpa")
def register_general_sdpa():
    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,
    )


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


@update_features(exir_ops.edge.et_vk.apply_rotary_emb_interleaved.default)
def register_apply_rotary_emb_interleaved():
    return OpFeatures(
        # freqs_cis is pinned to buffer storage so the shader can compute a
        # flat [N, C] linear address regardless of the tensor's declared rank
        # (callers commonly pass 4D [1, N, C/2, 2] without a preceding view).
        inputs_storage=[
            utils.CONTIGUOUS_ANY,  # x
            utils.CONTIGUOUS_BUFFER,  # freqs_cis
        ],
        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_INCL_PACKED_INT8,
        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_INCL_PACKED_INT8,
        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_INCL_PACKED_INT8,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# alias_copy is a no-op identity operation (same as clone/alias). It is removed
# by RemoveRedundantOpsTransform during preprocess, so the C++ runtime never sees
# it. Registering it here ensures the partitioner delegates it to Vulkan rather
# than creating partition boundaries that break the graph.
@update_features(exir_ops.edge.aten.alias_copy.default)
def register_alias_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE_INCL_PACKED_INT8,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


# =============================================================================
# Unfold.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.unfold_copy.default)
def register_unfold_copy():
    return OpFeatures(
        inputs_storage=utils.ANY_BUFFER,
        inputs_dtypes=utils.FP_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_STORAGE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        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,
    )


# =============================================================================
# SliceScatter.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.slice_scatter.default)
def register_slice_scatter():
    def check_slice_scatter_node(node: torch.fx.Node) -> bool:
        """Only delegate slice_scatter when its window is fixed at build time.

        SliceScatter.cpp normalizes start/end against the scattered dim's size
        and bakes them into push constants when the graph is built, so the
        window must not depend on anything that changes at runtime: a symbolic
        start/end (e.g. a KV-cache position) or a dynamic scattered dim (which
        moves a negative or default start/end) would silently scatter into the
        wrong rows. Same reasoning as _check_pad_is_static. Dynamic sizes on
        the other dims are fine, since the shader reads those from the tensor
        metadata UBOs.
        """
        # Schema: slice_scatter(self, src, dim=0, start=None, end=None, step=1)
        args = list(node.args) + [None] * (6 - len(node.args))
        self_arg, src_arg, dim, start, end, step = args[:6]
        dim = 0 if dim is None else dim
        step = 1 if step is None else step
        if not isinstance(dim, int) or not isinstance(step, int):
            return False
        # A non-positive step has no meaning here and aten rejects it; the
        # shader divides by step, so refuse rather than fault.
        if step <= 0:
            return False
        if not all(i is None or isinstance(i, int) for i in (start, end)):
            return False

        for t in (self_arg, src_arg):
            if not isinstance(t, torch.fx.Node):
                return False
            val = t.meta.get("val", None)
            if not isinstance(val, torch.Tensor):
                return False
            if not -val.dim() <= dim < val.dim():
                return False
            if not isinstance(val.shape[dim], int):
                return False
        # aten casts src to self's dtype, but the shader reads all three
        # tensors with one dtype and SliceScatter.cpp rejects a mismatch, so
        # leave mixed dtypes to the CPU.
        if self_arg.meta["val"].dtype != src_arg.meta["val"].dtype:
            return False
        return True

    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_BOOL_T,
        supports_resize=True,
        are_node_inputs_supported_fn=check_slice_scatter_node,
    )


# =============================================================================
# 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,
        supports_resize=True,
    )


# =============================================================================
# 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
# =============================================================================


def _index_tensor_shapes(node: torch.fx.Node):
    """Return self, index, and axis for the supported form, else None."""
    self_arg = node.args[0]
    indices = node.args[1]

    if not isinstance(self_arg, torch.fx.Node):
        return None
    self_val = self_arg.meta.get("val", None)
    if self_val is None or not isinstance(indices, (list, tuple)):
        return None

    non_none = [(dim, index) for dim, index in enumerate(indices) if index is not None]
    if len(non_none) != 1:
        return None
    index_dim, index_arg = non_none[0]
    if index_dim >= len(self_val.size()) or not isinstance(index_arg, torch.fx.Node):
        return None
    index_val = index_arg.meta.get("val", None)
    if index_val is None:
        return None

    return self_val, index_val, index_dim


def _check_index_tensor_node(node: torch.fx.Node) -> bool:
    shapes = _index_tensor_shapes(node)
    if shapes is None:
        return False
    _, index_val, _ = shapes
    # The gather is expressed as "one index position per output slice", so
    # the index must be 1-D. `self` may be any rank.
    return len(index_val.size()) == 1


def _pick_index_tensor_storage(node: torch.fx.Node):
    shapes = _index_tensor_shapes(node)
    # Only the buffer shader handles a higher-rank `self`.
    if shapes is not None and len(shapes[0].size()) > 1:
        return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER
    return utils.ANY_STORAGE, utils.ANY_STORAGE


@update_features(exir_ops.edge.aten.index.Tensor)
def register_index_tensor():
    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,
        pick_io_storage_fn=_pick_index_tensor_storage,
    )


# =============================================================================
# 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,
        supports_resize=True,
    )


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


def _check_pad_is_static(node: torch.fx.Node) -> bool:
    """Only support constant_pad_nd when the pad amounts are static.

    A symbolic pad list is serialized as a VALUELIST rather than an INTLIST, and
    Pad.cpp reads it with get_int_list(), which throws "Expected value to have
    type IntList, got VALUELIST instead".

    Supporting it properly is more than swapping in
    extract_int_or_symint_list(), the way Split.cpp, View.cpp and Expand.cpp
    read their symbolic lists: add_constant_pad_nd_node() folds the amounts into
    a per-dim offset and bakes that into a params buffer at BUILD time, so the
    dispatch would still use stale offsets even if the list were read
    symbolically. The buffer has to be refreshed on resize first. Decline the
    node until then.
    """
    pad = node.args[1]
    if not isinstance(pad, (list, tuple)):
        return False
    return all(isinstance(p, int) for p in pad)


@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,
        are_node_inputs_supported_fn=_check_pad_is_static,
    )


# =============================================================================
# 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,
        supports_resize=True,
        are_node_inputs_supported_fn=lambda node: node.target
        not in (exir_ops.edge.aten.full.default, exir_ops.edge.aten.full_like.default)
        or is_scalar_value_supported(node.args[1], node.meta["val"].dtype),
    )


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


@update_features(
    [
        exir_ops.edge.aten.scalar_tensor.default,
        # EXIR deliberately keeps scalar_tensor in the ATen dialect.
        torch.ops.aten.scalar_tensor.default,
    ]
)
def register_scalar_tensor():
    return OpFeatures(
        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
        inputs_dtypes=utils.FP_INT_T,
        supports_resize=True,
        are_node_inputs_supported_fn=lambda node: is_scalar_value_supported(
            node.args[0], node.meta["val"].dtype
        ),
    )


# =============================================================================
# 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,
    )


# =============================================================================
# PixelShuffle.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.pixel_shuffle.default)
def register_pixel_shuffle():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
    )


# =============================================================================
# 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,
    )


# =============================================================================
# GridSampler2d.cpp
# =============================================================================


@update_features(exir_ops.edge.aten.grid_sampler_2d.default)
def register_grid_sampler_2d():
    # The Vulkan implementation supports bilinear interpolation (0) with
    # either zeros (0) or border (1) padding and either align_corners. That
    # covers RIFE's WarpModule (border, align_corners=True) and the deformable
    # attention in DETR derivatives (zeros, align_corners=False). Reflection
    # padding and nearest/bicubic interpolation are not implemented.
    # The C++ side has VK_CHECK_COND asserts for these,
    # but those abort the whole inference at graph build — for any other model
    # that contains a differently-configured grid_sampler_2d we want graceful
    # CPU fallback, so we gate delegation here.
    #
    # Edge IR can hand us these scalar args as plain Python literals, SymInt /
    # SymBool wrappers, or get_attr-style fx.Node references, so we unwrap
    # each one defensively (mirrors the `isinstance(groups, int)` guard in
    # check_conv_node / pick_conv_storage above). If we can't confidently pull
    # a literal out of any arg, return False so the node stays on CPU instead
    # of hitting a runtime VK_CHECK_COND.
    def _unwrap_literal(arg: object) -> object:
        # Plain Python literal (covers bool, since bool is a subclass of int).
        if isinstance(arg, (bool, int, float)):
            return arg
        # get_attr / constant fx.Node — read the materialized value from meta.
        if isinstance(arg, torch.fx.Node):
            val = arg.meta.get("val", None)
            if isinstance(val, (bool, int, float)):
                return val
            return None
        # Symbolic int/bool (or anything else int-convertible) — try once.
        try:
            return int(arg)  # pyre-ignore[6]
        except (TypeError, ValueError):
            return None

    def check_grid_sampler_2d_node(node: torch.fx.Node) -> bool:
        # Schema: aten::grid_sampler_2d(input, grid, interpolation_mode,
        #                               padding_mode, align_corners)
        if len(node.args) < 5:
            return False

        interp = _unwrap_literal(node.args[2])
        padding = _unwrap_literal(node.args[3])
        align_corners = _unwrap_literal(node.args[4])

        if interp is None or padding is None or align_corners is None:
            return False

        # mode: 0 = bilinear. padding: 0 = zeros, 1 = border (2 = reflection
        # needs a coordinate fold the shader does not implement).
        # align_corners is free: both settings are specialization constants.
        return interp == 0 and padding in (0, 1)

    return OpFeatures(
        inputs_storage=[
            utils.CHANNELS_PACKED_TEXTURE,  # input  : [N, C, Hin, Win]
            utils.CONTIGUOUS_BUFFER,  # grid   : [N, Hout, Wout, 2]
        ],
        inputs_dtypes=utils.FP_T,
        supports_resize=True,
        are_node_inputs_supported_fn=check_grid_sampler_2d_node,
    )


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


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


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


@update_features(exir_ops.edge.aten.embedding.default)
def register_embedding():
    def check_embedding_weight_size(node: torch.fx.Node) -> bool:
        weight = node.args[0]
        if isinstance(weight, torch.fx.Node) and utils.is_tensor_node(weight):
            numel = weight.meta["val"].numel()
            if numel > utils.DEFAULT_BUFFER_LIMIT:
                return False
        return True

    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=[utils.FP_T, utils.INT_T],
        supports_prepacking=True,
        supports_resize=True,
        are_node_inputs_supported_fn=check_embedding_weight_size,
    )


# =============================================================================
# EmbeddingQ4gsw (Quantized Embedding)
# =============================================================================


@update_features(exir_ops.edge.quantized_decomposed.embedding_4bit.dtype)
def register_quantized_decomposed_embedding_4bit():
    def check_embedding_4bit_weight_size(node: torch.fx.Node) -> bool:
        weight = node.args[0]
        if isinstance(weight, torch.fx.Node) and utils.is_tensor_node(weight):
            numel = weight.meta["val"].numel()
            if numel > utils.DEFAULT_BUFFER_LIMIT:
                return False
        return True

    return OpFeatures(
        inputs_storage=utils.ANY_BUFFER,
        supports_prepacking=True,
        supports_resize=True,
        are_node_inputs_supported_fn=check_embedding_4bit_weight_size,
    )


@update_features(exir_ops.edge.et_vk.embedding_q4gsw.default)
def register_embedding_q4gsw():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        supports_prepacking=True,
        supports_resize=True,
    )


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


def _check_batch_norm_is_4d(node: torch.fx.Node) -> bool:
    """Only support batch norm on a 4d input.

    add_native_batch_norm_node() asserts
    VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor") on
    both the input and the output, so partitioning a batch norm whose input is
    not 4d yields a .pte that lowers cleanly and then aborts at execute time.
    Any conv1d model reaches here with rank-3 activations.
    """
    input_node = node.args[0]
    if not isinstance(input_node, torch.fx.Node):
        return False
    val = input_node.meta.get("val")
    return val is not None and val.dim() == 4


@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,
        are_node_inputs_supported_fn=_check_batch_norm_is_4d,
    )


def unwrap_identity_repeat(arg: Any) -> Any:
    """
    Return the source of ``arg`` if it is a ``repeat`` whose every factor is 1
    (its output shape matches its input shape), else ``arg`` itself.
    """
    if (
        isinstance(arg, torch.fx.Node)
        and arg.target == exir_ops.edge.aten.repeat.default
        and isinstance(arg.args[0], torch.fx.Node)
        and arg.args[0].meta["val"].shape == arg.meta["val"].shape
    ):
        return arg.args[0]
    return arg


def can_rewrite_batch_norm_as_group_norm(
    node: torch.fx.Node,
    constant_nodes: AbstractSet[torch.fx.Node],
) -> bool:
    """
    Whether ReplaceInstanceNormPass can rewrite a batch norm node into group norm.

    ``F.instance_norm`` reshapes its input to ``[1, N * C, H, W]`` and lowers to
    ``_native_batch_norm_legit.no_stats``, which normalizes using statistics taken
    over the batch and spatial dims. When the batch dim is 1 that is exactly group
    norm with one group per channel, so the existing group norm kernels cover it.
    A batch dim above 1 (only reachable by calling batch norm directly) is a
    different reduction and is left alone.

    ``constant_nodes`` holds the placeholders whose data is known at export time.
    """
    if node.target != exir_ops.edge.aten._native_batch_norm_legit.no_stats:
        return False

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

    val = input_node.meta.get("val")
    if val is None or val.dim() != 4:
        return False
    # A symbolic batch dim could be above 1 at runtime, and the pass passes the
    # static channel and spatial sizes to group norm.
    if not all(isinstance(d, int) for d in val.shape) or val.shape[0] != 1:
        return False

    # Group norm always applies an affine transform, so both weight and bias must
    # be present. add_native_group_norm_node() prepacks them, so each must be a
    # constant by the time group norm sees it. F.instance_norm passes
    # weight.repeat(N): ReplaceInstanceNormPass looks through that repeat when N
    # is 1, but for N above 1 it is a runtime tensor, so instance norm over a
    # batch is rejected here. So are computed affine args and ones fed in as model
    # inputs. A repeat outside the norm's partition also reaches the delegate as a
    # runtime input, but that is only known once partitions are proposed, so
    # VulkanPartitioner checks it then.
    for affine_arg in (node.args[1], node.args[2]):
        if unwrap_identity_repeat(affine_arg) not in constant_nodes:
            return False

    # Only the normalized output may be consumed. Group norm returns mean and rstd
    # shaped [N, group] where batch norm saves them shaped [C], so the saved
    # statistics are not drop-in replacements.
    for user in node.users:
        if user.op != "call_function" or user.target != operator.getitem:
            return False
        if user.args[1] != 0:
            return False

    return True


@update_features(exir_ops.edge.aten._native_batch_norm_legit.no_stats)
def register_native_batch_norm_legit_no_stats():
    """Instance norm, which ReplaceInstanceNormPass rewrites into group norm.

    ``F.instance_norm`` lowers to this overload. The pass only handles the cases
    can_rewrite_batch_norm_as_group_norm() accepts, so gate partitioning on the
    same predicate.
    """
    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,
        are_node_inputs_supported_with_constants_fn=can_rewrite_batch_norm_as_group_norm,
    )


# =============================================================================
# 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_STORAGE,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
        supports_resize=True,
    )


# =============================================================================
# RmsNorm.cpp
# =============================================================================


@update_features(exir_ops.edge.et_vk.rms_norm.default)
def register_rms_norm():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
        supports_resize=True,
    )


# =============================================================================
# FusedCe.cpp (training)
# =============================================================================


@update_features(exir_ops.edge.et_vk.fused_ce.default)
def register_fused_ce():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=[utils.FP_T, utils.INT_T, utils.NONE_T],
        outputs_dtypes=[utils.FP_T, utils.FP_T],
    )


@update_features(
    [
        exir_ops.edge.aten.ne.Scalar,
        exir_ops.edge.aten.lt.Scalar,
        exir_ops.edge.aten.le.Scalar,
        exir_ops.edge.aten.ge.Scalar,
        exir_ops.edge.aten.gt.Scalar,
    ]
)
def register_compare_scalar_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,
        are_node_inputs_supported_fn=lambda node: is_scalar_value_supported(
            node.args[1], node.args[0].meta["val"].dtype
        ),
    )


@update_features(exir_ops.edge.aten.logical_not.default)
def register_logical_not():
    return OpFeatures(
        inputs_storage=utils.ANY_STORAGE,
        inputs_dtypes=utils.BOOL_T,
        supports_resize=True,
        supports_highdim=True,
    )


@update_features("et_vk::adamw_step")
def register_adamw_step():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
    )


@update_features(exir_ops.edge.et_vk.linear_dW.default)
def register_linear_dW():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
        supports_prepacking=True,
    )


@update_features(exir_ops.edge.et_vk.q4gsw_requant.default)
def register_q4gsw_requant():
    return OpFeatures(
        inputs_storage=utils.CONTIGUOUS_ANY,
        inputs_dtypes=utils.FP_T,
    )


#######################
## 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