executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/backends/vulkan/utils.py

@@ -0,0 +1,2056 @@
+# 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.
+
+import operator
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+import torch
+from executorch.backends.vulkan.serialization.vulkan_graph_schema import (
+    VkDataType,
+    VkMemoryLayout,
+    VkStorageType,
+)
+from executorch.exir.backend.canonical_partitioners.config_partitioner import (
+    format_target_name,
+)
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
+from executorch.exir.tensor import TensorSpec
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
+from torch._subclasses.fake_tensor import FakeTensor, FakeTensorConverter
+from torch.export import ExportedProgram
+from torch.export.exported_program import InputKind
+from torch.export.graph_signature import TensorArgument
+
+TorchOpType = Union[EdgeOpOverload, torch._ops.OpOverload, str]
+
+_DQ_OPS = {
+    "dequantize_per_tensor.tensor",
+    "dequantize_per_tensor.default",
+    "dequantize_per_channel.default",
+    "dequantize_per_channel_group.default",
+    "dequantize_per_token.default",
+    "dequantize_affine.default",
+}
+
+_Q_OPS = {
+    "quantize_per_tensor.tensor",
+    "quantize_per_tensor.default",
+    "quantize_per_channel.default",
+    "quantize_per_token.default",
+    "quantize_affine.default",
+}
+
+_VULKAN_DTYPES: Dict[torch.dtype, VkDataType] = {
+    torch.bool: VkDataType.BOOL,
+    torch.uint8: VkDataType.UINT8,
+    torch.int8: VkDataType.INT8,
+    torch.int32: VkDataType.INT32,
+    torch.int64: VkDataType.INT64,
+    torch.float16: VkDataType.FLOAT16,
+    torch.float32: VkDataType.FLOAT32,
+    torch.float64: VkDataType.FLOAT64,
+}
+
+##
+## Dtype sets for per-operator dtype constraints
+##
+
+DtypeSet = Set[torch.dtype]
+
+FP_T: DtypeSet = {torch.float16, torch.float32}
+INT_T: DtypeSet = {torch.int32, torch.int64}
+QINT8_T: DtypeSet = {torch.int8}
+BOOL_T: DtypeSet = {torch.bool}
+ALL_T: DtypeSet = set(_VULKAN_DTYPES.keys())
+NONE_T: DtypeSet = set()  # Marker for non-tensor args (skip validation)
+
+# Composite dtype sets for specific operator requirements
+FP_INT_T: DtypeSet = FP_T | INT_T
+FP_INT_BOOL_T: DtypeSet = FP_T | INT_T | BOOL_T
+
+
+class DtypeSetList:
+    """
+    Wrapper around a list of DtypeSet with broadcasting semantics.
+    If only one DtypeSet is provided, it applies to all positions.
+    """
+
+    def __init__(self, dtype_sets: Union[DtypeSet, List[DtypeSet]]):
+        self.vals: List[DtypeSet] = (
+            dtype_sets if isinstance(dtype_sets, list) else [dtype_sets]
+        )
+
+    def __len__(self) -> int:
+        return len(self.vals)
+
+    def __getitem__(self, idx: int) -> DtypeSet:
+        # Broadcasting: single set applies to all positions
+        if idx > 0 and len(self.vals) == 1:
+            return self.vals[0]
+        if idx >= len(self.vals):
+            return set()
+        return self.vals[idx]
+
+    def is_empty(self) -> bool:
+        return len(self.vals) == 0
+
+    def any_constrained(self) -> bool:
+        """Returns True if any position has dtype constraints."""
+        return any(len(s) > 0 for s in self.vals)
+
+
+##
+## Node type determination
+##
+
+# Convenience type
+MaybeNodeList = Union[torch.fx.Node, List[torch.fx.Node], Tuple[torch.fx.Node]]
+
+
+def is_torch_op_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+
+    if isinstance(node.target, EdgeOpOverload):
+        return True
+    if isinstance(node.target, torch._ops.OpOverload):
+        return True
+
+    return False
+
+
+def is_dequant_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return node_name in _DQ_OPS
+
+
+def is_quant_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return node_name in _Q_OPS
+
+
+def is_choose_qparams_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return "choose_qparams" in node_name
+
+
+def is_dynamic_qscale(node: Any) -> bool:
+    """Check if a scale node is dynamically computed via a choose_qparams op."""
+    return (
+        isinstance(node, torch.fx.Node)
+        and node.target == operator.getitem
+        and is_choose_qparams_node(node.args[0])
+    )
+
+
+def is_dequant_per_channel_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return node_name == "dequantize_per_channel.default"
+
+
+def is_view_copy_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return "view_copy" in node_name
+
+
+def is_linear_node(node: torch.fx.Node) -> bool:
+    if node.op != "call_function":
+        return False
+    node_name = format_target_name(node.target.__name__)  # pyre-ignore
+    return node_name == "linear.default"
+
+
+def is_get_attr_node(node: torch.fx.Node) -> bool:
+    return isinstance(node, torch.fx.Node) and node.op == "get_attr"
+
+
+def is_constant(program: ExportedProgram, node: torch.fx.Node) -> bool:
+    return node.name in program.graph_signature.inputs_to_lifted_tensor_constants
+
+
+def is_param_node(program: ExportedProgram, node: torch.fx.Node) -> bool:
+    """
+    Check if the given node is a parameter within the exported program
+    """
+    return (
+        is_get_attr_node(node)
+        or is_param(program, node)
+        or is_buffer(program, node)
+        or is_lifted_tensor_constant(program, node)
+    )
+
+
+def is_mutable_buffer_node(
+    node: torch.fx.Node, exported_program: ExportedProgram
+) -> bool:
+    if node.target not in exported_program.graph_signature.inputs_to_buffers:
+        return False
+    buf = exported_program.graph_signature.inputs_to_buffers[node.target]
+    return buf in exported_program.graph_signature.buffers_to_mutate.values()
+
+
+def is_symint_node(node: torch.fx.Node) -> bool:
+    """
+    Returns true if the given node produces a SymInt value
+    """
+    if "val" not in node.meta:
+        return False
+
+    if isinstance(node.meta["val"], torch.SymInt):
+        return True
+
+    return False
+
+
+def is_single_tensor_node(node: torch.fx.Node) -> bool:
+    """
+    Returns true if the given node produces a single tensor value
+    """
+    if "val" not in node.meta:
+        return False
+
+    if isinstance(node.meta["val"], FakeTensor):
+        return True
+
+    return False
+
+
+def is_tensor_collection_node(node: Any) -> bool:
+    """
+    Returns true if the given node produces a collection of tensor values
+    """
+    if not isinstance(node, torch.fx.Node):
+        return False
+
+    if "val" not in node.meta:
+        return False
+
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        return all(isinstance(x, FakeTensor) for x in node.meta["val"])
+
+    return False
+
+
+def is_tensor_node(node: Any) -> bool:
+    """
+    Returns true if the given node produces a tensor value, or a collection of tensor values
+    """
+    if not isinstance(node, torch.fx.Node):
+        return False
+
+    if "val" not in node.meta:
+        return False
+
+    if isinstance(node.meta["val"], FakeTensor):
+        return True
+
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        return all(isinstance(x, FakeTensor) for x in node.meta["val"])
+
+    return False
+
+
+def is_tensor_arg_node(node: Any) -> bool:
+    if isinstance(node, torch.fx.Node):
+        return is_tensor_node(node)
+    elif isinstance(node, (list, tuple)):
+        if len(node) == 0:
+            return False
+        return all(is_tensor_node(n) for n in node)
+
+    return False
+
+
+def num_tensor_arg_nodes(node: torch.fx.Node) -> int:
+    """
+    For a given node, return the number of argument nodes that are associated with
+    tensors.
+    """
+    count = 0
+    for arg_node in node.args:
+        if not isinstance(arg_node, torch.fx.Node):
+            continue
+        if is_tensor_node(arg_node):
+            count += 1
+
+    return count
+
+
+def num_tensors_in_node(node: torch.fx.Node) -> int:
+    """
+    Returns the number of tensors associated a given node
+    """
+    if "val" not in node.meta:
+        return 0
+
+    if isinstance(node.meta["val"], FakeTensor):
+        return 1
+
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        if all(isinstance(x, FakeTensor) for x in node.meta["val"]):
+            return len(node.meta["val"])
+
+    return 0
+
+
+def get_vk_datatype(torch_dtype: torch.dtype) -> VkDataType:
+    """
+    Returns Vulkan dtype corresponding to torch dtype
+    """
+    if torch_dtype not in _VULKAN_DTYPES:
+        raise AssertionError(f"Invalid dtype for vulkan_preprocess ({torch_dtype})")
+
+    return _VULKAN_DTYPES[torch_dtype]
+
+
+def output_dtypes_are_supported(node: torch.fx.Node) -> bool:
+    """
+    Returns true if the output of the given tensor node has dtype that
+    is supported by the Vulkan backend.
+    """
+    if not is_tensor_node(node):
+        return True
+
+    # The val metadata must exist after previous check
+    node_val = node.meta.get("val", None)
+    assert node_val is not None
+
+    # Get all the tensor dtypes in the node
+    tensor_dtypes = []
+    if isinstance(node_val, FakeTensor):
+        tensor_dtypes = [node_val.dtype]
+    elif isinstance(node_val, list) or isinstance(node_val, tuple):
+        tensor_dtypes = [x.dtype for x in node_val]
+
+    # Verify that all the tensor_dtypes are in vk_torch_dtypes
+    return all(dtype in _VULKAN_DTYPES for dtype in tensor_dtypes)
+
+
+def input_dtypes_are_supported(node: torch.fx.Node) -> bool:
+    """
+    Returns true if all the inputs to the given tensor node have dtype that
+    is supported by the Vulkan backend.
+    """
+    if not is_tensor_node(node):
+        return True
+
+    # Iterate over all the args of the node
+    for arg_node in node.args:
+        # The arg could be a single node, or a list (e.g., first arg of cat)
+        if isinstance(arg_node, torch.fx.Node):
+            if not output_dtypes_are_supported(arg_node):
+                return False
+        elif isinstance(arg_node, (list, tuple)):
+            if not all(output_dtypes_are_supported(x) for x in arg_node):
+                return False
+
+    return True
+
+
+def io_dtypes_are_supported(node: torch.fx.Node) -> bool:
+    """
+    Returns true if all the inputs and outputs of the given tensor node have
+    dtype that is supported by the Vulkan backend.
+    """
+    if not output_dtypes_are_supported(node):
+        return False
+    if not input_dtypes_are_supported(node):
+        return False
+
+    return True
+
+
+def check_node_dtypes(  # noqa: C901
+    node: torch.fx.Node,
+    inputs_dtypes: DtypeSetList,
+    outputs_dtypes: DtypeSetList,
+) -> Tuple[bool, str]:
+    """
+    Check if all tensor inputs/outputs have dtypes in the allowed sets.
+    Returns (is_valid, reason_string) for better error reporting.
+    """
+    # Check input tensor dtypes
+    for i, arg in enumerate(node.args):
+        allowed_dtypes = inputs_dtypes[i]
+        # Skip non-constrained positions (NO_DTYPE = empty set)
+        if len(allowed_dtypes) == 0:
+            continue
+
+        if is_tensor_node(arg):
+            if isinstance(arg.meta["val"], (list, tuple)):
+                arg_dtype = arg.meta["val"][0].dtype
+            else:
+                arg_dtype = arg.meta["val"].dtype
+            if arg_dtype not in allowed_dtypes:
+                return False, f"input[{i}] dtype {arg_dtype} not in {allowed_dtypes}"
+
+        elif isinstance(arg, (list, tuple)):
+            # Handle tensor list inputs (e.g., cat)
+            for j, sub_arg in enumerate(arg):
+                if is_tensor_node(sub_arg):
+                    sub_dtype = sub_arg.meta["val"].dtype
+                    if sub_dtype not in allowed_dtypes:
+                        return (
+                            False,
+                            f"input[{i}][{j}] dtype {sub_dtype} not in {allowed_dtypes}",
+                        )
+
+    # Check output tensor dtypes
+    out_val = node.meta.get("val")
+    if isinstance(out_val, FakeTensor):
+        allowed_dtypes = outputs_dtypes[0]
+        if len(allowed_dtypes) > 0 and out_val.dtype not in allowed_dtypes:
+            return False, f"output dtype {out_val.dtype} not in {allowed_dtypes}"
+    elif isinstance(out_val, (list, tuple)):
+        for i, t in enumerate(out_val):
+            if isinstance(t, FakeTensor):
+                allowed_dtypes = outputs_dtypes[i]
+                if len(allowed_dtypes) > 0 and t.dtype not in allowed_dtypes:
+                    return False, f"output[{i}] dtype {t.dtype} not in {allowed_dtypes}"
+
+    return True, "dtypes valid"
+
+
+def tensor_node_is_bool(node: torch.fx.Node) -> bool:
+    """
+    Returns true if a given node contains a tensor with bool dtype
+    """
+    if isinstance(node.meta["val"], FakeTensor):
+        return node.meta["val"].dtype == torch.bool
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        for fake_tensor in node.meta["val"]:
+            if isinstance(fake_tensor, FakeTensor):
+                if fake_tensor.dtype == torch.bool:
+                    return True
+    return False
+
+
+def ndim_of(node: Any) -> Optional[int]:
+    """
+    Returns the number of dimensions of the tensor produced by the given node
+    """
+    if not is_single_tensor_node(node):
+        return None
+
+    return node.meta["val"].ndim
+
+
+def is_unsqueezed_vector(node: torch.fx.Node) -> bool:
+    """
+    Returns True if the node's tensor has all dimensions equal to 1 except for the last dimension.
+    """
+    if not is_single_tensor_node(node):
+        return False
+
+    tensor = node.meta["val"]
+    assert isinstance(tensor, FakeTensor)
+
+    if len(tensor.shape) < 1:
+        return False
+    # All dims except last are 1, last can be any size
+    return all(dim == 1 for dim in tensor.shape[:-1])
+
+
+def op_contains_bool_tensor(node: torch.fx.Node) -> bool:
+    """
+    Returns true if the operator used to compute the given node contains a bool tensor
+    """
+    if is_tensor_node(node) and tensor_node_is_bool(node):
+        return True
+
+    for arg_node in node.args:
+        # pyre-ignore[6]
+        if is_tensor_node(arg_node) and tensor_node_is_bool(arg_node):
+            return True
+
+    return False
+
+
+def op_contains_high_dim_tensor(node: torch.fx.Node) -> bool:
+    """
+    Returns true if the operator used to compute the given node contains a tensor
+    with more than 4 dimensions
+    """
+    if is_tensor_node(node) and tensor_node_is_high_dim(node):
+        return True
+
+    for arg_node in node.args:
+        # pyre-ignore[6]
+        if is_tensor_node(arg_node) and tensor_node_is_high_dim(arg_node):
+            return True
+
+    return False
+
+
+def get_primary_arg_idx(self, node: torch.fx.Node) -> Optional[int]:
+    primary_arg_idx: Optional[int] = None
+    for i, arg_node in enumerate(node.args):
+        if self.is_non_constant_tensor_node(arg_node):
+            return i
+
+    return primary_arg_idx
+
+
+def node_comes_from_any_nn_module_in_set(
+    node,
+    nn_module_typenames: Set[str],
+) -> bool:
+    if isinstance(node, (list, tuple)):
+        return all(
+            node_comes_from_any_nn_module_in_set(n, nn_module_typenames) for n in node
+        )
+
+    if not isinstance(node, torch.fx.Node):
+        return False
+
+    nn_module_stack = node.meta.get("nn_module_stack", None)
+    if nn_module_stack is None:
+        return False
+
+    for _, packed in nn_module_stack.items():
+        _, typename = packed
+        for partial_name in nn_module_typenames:
+            if partial_name in typename:
+                return True
+
+    return False
+
+
+def get_tensor_name(exp_prog: ExportedProgram, node: torch.fx.Node) -> str:
+    if node is None:
+        return ""
+    if is_param(exp_prog, node):
+        return exp_prog.graph_signature.inputs_to_parameters[node.name]
+    elif is_buffer(exp_prog, node):
+        return exp_prog.graph_signature.inputs_to_buffers[node.name]
+    elif is_lifted_tensor_constant(exp_prog, node):
+        return exp_prog.graph_signature.inputs_to_lifted_tensor_constants[node.name]
+    else:
+        assert isinstance(node.target, str)
+        return node.target
+
+    return ""
+
+
+def register_param_mutation(ep: ExportedProgram, node: torch.fx.Node, tag: str) -> bool:
+    """Register an in-place mutation of a param/buffer's state-dict storage,
+    keyed on the stable identity of that storage, and report whether the caller
+    should perform the mutation.
+
+    A graph pass that repacks (or otherwise mutates in place) the constant tensor
+    backing a placeholder must run the mutation exactly once per backing storage;
+    repeating it on the already-mutated tensor corrupts it. Keying the guard on
+    node identity is unsafe, because two distinct placeholder nodes can resolve to
+    the same state-dict entry (e.g. aliased weights), and each would miss a
+    per-node flag and re-run the mutation. This guard instead keys on the
+    state-dict FQN that backs `node` -- the same identity the in-place mutation
+    targets (see `update_program_state_dict`, which overwrites
+    `state_dict[input_spec.target]`). The FQN is resolved via `get_tensor_name`
+    (`graph_signature.inputs_to_{parameters,buffers,lifted_tensor_constants}`) and
+    is stable across the mutation, unlike `data_ptr()` which changes once the
+    state-dict tensor object is replaced.
+
+    `tag` identifies the mutation/packing FORMAT, not the op. Two callers that
+    produce the same packed bytes for a weight should pass the same `tag` and
+    share this guard; observing two different tags on one weight means it would be
+    mutated two incompatible ways, which is corruption, so that case raises.
+
+    A per-ExportedProgram registry mapping `param_key -> tag` is lazily
+    initialized on `ep` (`ep._et_vk_param_modification_tags`) without
+    module-scope mutable state. Because it is stored on `ep`, it persists for the
+    lifetime of that ExportedProgram -- across all pass runs on it, not just one
+    -- and there is no reset hook, so a pass that legitimately needs to re-mutate
+    an already-tagged param must clear `ep._et_vk_param_modification_tags`
+    itself.
+
+    Returns True on the first call for a given param (and records the tag),
+    signalling the caller to proceed with the mutation. Returns False on a
+    subsequent call with the SAME tag, signalling the caller to skip (already
+    mutated this way). Raises RuntimeError on a call whose tag differs from the
+    recorded one.
+    """
+    registry: Dict[str, str] = getattr(ep, "_et_vk_param_modification_tags", None)
+    if registry is None:
+        registry = {}
+        ep._et_vk_param_modification_tags = registry
+
+    # The guard keys on the state-dict FQN that get_tensor_name resolves, which
+    # only exists for parameter / buffer / lifted-constant placeholders (via
+    # inputs_to_{parameters,buffers,lifted_tensor_constants}). For any other node
+    # kind get_tensor_name falls through to node.target, which is not a stable
+    # storage key, so the guard would silently key on the wrong identity. Reject
+    # such nodes up front.
+    if not (
+        is_param(ep, node) or is_buffer(ep, node) or is_lifted_tensor_constant(ep, node)
+    ):
+        raise RuntimeError(
+            "register_param_mutation expects a parameter / buffer / "
+            f"lifted-constant placeholder, but got node {node.name!r} "
+            f"(op={node.op!r}, target={node.target!r})."
+        )
+
+    param_key = get_tensor_name(ep, node)
+
+    recorded_tag = registry.get(param_key)
+    if recorded_tag is None:
+        registry[param_key] = tag
+        return True
+
+    if recorded_tag != tag:
+        raise RuntimeError(
+            f"Param tensor {param_key!r} was already modified with tag "
+            f"{recorded_tag!r}, but a conflicting modification with tag {tag!r} "
+            "was attempted; a weight modified two different ways is corrupt."
+        )
+
+    return False
+
+
+def find_dequant_user(node: torch.fx.Node) -> Optional[torch.fx.Node]:
+    """
+    Search the direct users of the given node and return the first one that is a
+    dequantization op. Returns None if no dequantization op is found.
+    """
+    for user in node.users:
+        if is_dequant_node(user):
+            return user
+    return None
+
+
+def find_quant_user(node: torch.fx.Node) -> Optional[torch.fx.Node]:
+    """
+    Search the direct users of the given node and return the first one that is a
+    quantization op. Returns None if no quantization op is found.
+    """
+    for user in node.users:
+        if is_quant_node(user):
+            return user
+
+    return None
+
+
+def node_has_target(node: Any, target: str):
+    if not hasattr(node, "target"):
+        return False
+
+    if isinstance(node.target, str):
+        return node.target == target
+    elif hasattr(node.target, "name"):
+        return node.target.name() == target
+
+    return False
+
+
+##
+## Memory Layout, Storage Type Determination
+##
+
+ImageExtents = Tuple[int, int, int]
+
+DEFAULT_TEXTURE_LIMITS = (16384, 16384, 2048)
+# Conservative 3D texture limit compatible with most desktop/laptop GPUs. The
+# Vulkan spec only guarantees maxImageDimension3D >= 2048, whereas mobile GPUs
+# commonly support 16384. Used when the small_texture_limits option is set.
+SMALL_TEXTURE_LIMITS = (2048, 2048, 2048)
+DEFAULT_BUFFER_LIMIT = 128 * (1024 * 1024)
+
+all_storage_types: Set[VkStorageType] = {
+    VkStorageType.BUFFER,
+    VkStorageType.TEXTURE_3D,
+}
+
+# Memory layouts available to non-quantized tensors
+all_memory_layouts: Set[VkMemoryLayout] = {
+    VkMemoryLayout.TENSOR_WIDTH_PACKED,
+    VkMemoryLayout.TENSOR_HEIGHT_PACKED,
+    VkMemoryLayout.TENSOR_CHANNELS_PACKED,
+}
+
+# Memory layouts available to quantized tensors
+all_quantized_memory_layouts: Set[VkMemoryLayout] = {
+    VkMemoryLayout.PACKED_INT8_4W4C,
+    VkMemoryLayout.PACKED_INT8_4H4W,
+    VkMemoryLayout.PACKED_INT8_4W,
+    VkMemoryLayout.PACKED_INT8_4C1W,
+    VkMemoryLayout.PACKED_INT8_CONV2D,
+}
+
+universal_memory_layout_set: Set[VkMemoryLayout] = (
+    all_memory_layouts | all_quantized_memory_layouts
+)
+
+MemoryLayoutSet = Set[VkMemoryLayout]
+MemoryLayoutSetList = Union[MemoryLayoutSet, List[MemoryLayoutSet]]
+
+_LAYOUT_TO_PACKED_DIM: Dict[VkMemoryLayout, int] = {
+    VkMemoryLayout.TENSOR_WIDTH_PACKED: 0,
+    VkMemoryLayout.TENSOR_HEIGHT_PACKED: 1,
+    VkMemoryLayout.TENSOR_CHANNELS_PACKED: 2,
+    VkMemoryLayout.PACKED_INT8_4W4C: 2,
+    VkMemoryLayout.PACKED_INT8_4H4W: 0,
+    VkMemoryLayout.PACKED_INT8_4C1W: 2,
+    VkMemoryLayout.PACKED_INT8_CONV2D: 2,
+}
+
+
+def packed_dim_of(layout: VkMemoryLayout) -> int:
+    return _LAYOUT_TO_PACKED_DIM[layout]
+
+
+@dataclass(frozen=True)
+class PackedDimInfo:
+    """
+    Describes how tensor data is organized in physical memory, mirroring the
+    C++ PackedDimInfo struct in runtime/api/containers/Tensor.h.
+    """
+
+    packed_dim: int
+    packed_dim_block_size: int
+
+    @classmethod
+    def from_repr(
+        cls,
+        memory_layout: VkMemoryLayout,
+        storage_type: VkStorageType = VkStorageType.BUFFER,
+    ) -> "PackedDimInfo":
+        """
+        Construct a PackedDimInfo based on a memory layout and storage type,
+        mirroring calculate_packed_dim_info in runtime/api/containers/Tensor.cpp.
+        """
+        is_buffer = storage_type == VkStorageType.BUFFER
+
+        if memory_layout == VkMemoryLayout.TENSOR_WIDTH_PACKED:
+            return cls(
+                packed_dim=0,
+                packed_dim_block_size=1 if is_buffer else 4,
+            )
+        elif memory_layout == VkMemoryLayout.TENSOR_HEIGHT_PACKED:
+            return cls(
+                packed_dim=1,
+                packed_dim_block_size=1 if is_buffer else 4,
+            )
+        elif memory_layout == VkMemoryLayout.TENSOR_CHANNELS_PACKED:
+            return cls(
+                packed_dim=2,
+                packed_dim_block_size=1 if is_buffer else 4,
+            )
+        elif memory_layout == VkMemoryLayout.PACKED_INT8_4W:
+            return cls(
+                packed_dim=0,
+                packed_dim_block_size=4,
+            )
+        elif memory_layout == VkMemoryLayout.PACKED_INT8_4W4C:
+            return cls(
+                packed_dim=2,
+                packed_dim_block_size=4,
+            )
+        elif memory_layout == VkMemoryLayout.PACKED_INT8_4H4W:
+            return cls(
+                packed_dim=0,
+                packed_dim_block_size=4,
+            )
+        elif memory_layout == VkMemoryLayout.PACKED_INT8_4C1W:
+            return cls(
+                packed_dim=2,
+                packed_dim_block_size=4 if is_buffer else 16,
+            )
+        elif memory_layout == VkMemoryLayout.PACKED_INT8_CONV2D:
+            return cls(
+                packed_dim=2,
+                packed_dim_block_size=4,
+            )
+        else:
+            raise ValueError(f"Unknown memory layout: {memory_layout}")
+
+
+def within_buffer_limit(node: torch.fx.Node, buffer_limit: int) -> int:
+    """
+    Checks whether the tensors produced by the given node can fit within the device's
+    GPU buffer limit, which represents the maximum number of elements that can be stored
+    in a GPU buffer.
+    """
+    assert is_tensor_node(node)
+
+    if isinstance(node.meta["val"], FakeTensor):
+        return node.meta["val"].numel() < buffer_limit
+    elif isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        return all(x.numel() < buffer_limit for x in node.meta["val"])
+    else:
+        raise RuntimeError(f"Cannot get numel for val of type {type(node.meta['val'])}")
+
+
+def tensor_node_is_high_dim(node: torch.fx.Node) -> bool:
+    """
+    Returns true if a given node contains a tensor with more than 4 dimensions
+    """
+    if isinstance(node.meta["val"], FakeTensor):
+        return len(node.meta["val"].shape) > 4
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        for fake_tensor in node.meta["val"]:
+            if isinstance(fake_tensor, FakeTensor):
+                if len(fake_tensor.shape) > 4:
+                    return True
+    return False
+
+
+def required_image_extents(sizes: torch.Size, layout: VkMemoryLayout) -> ImageExtents:
+    """
+    Calculate the image extents that will be used to represent a tensor with the given sizes
+    and memory layout in the Vulkan Delegate.
+    """
+    width = sizes[-1] if len(sizes) >= 1 else 1
+    height = sizes[-2] if len(sizes) >= 2 else 1
+    channels = sizes[-3] if len(sizes) >= 3 else 1
+    batch = sizes[0] if len(sizes) >= 4 else 1
+
+    if layout == VkMemoryLayout.TENSOR_WIDTH_PACKED:
+        width = (width + 3) // 4
+    elif layout == VkMemoryLayout.TENSOR_HEIGHT_PACKED:
+        height = (height + 3) // 4
+    elif layout == VkMemoryLayout.TENSOR_CHANNELS_PACKED:
+        channels = (channels + 3) // 4
+    elif layout == VkMemoryLayout.PACKED_INT8_4W4C:
+        width = (width + 3) // 4
+        channels = (channels + 3) // 4
+    elif layout == VkMemoryLayout.PACKED_INT8_4H4W:
+        height = (height + 3) // 4
+        width = (width + 3) // 4
+    elif layout == VkMemoryLayout.PACKED_INT8_CONV2D:
+        # Use conservative extents (same as 4W4C) since this is buffer-only
+        width = (width + 3) // 4
+        channels = (channels + 3) // 4
+    else:
+        raise RuntimeError(f"Unsupported memory layout {layout}")
+
+    return width, height, channels * batch
+
+
+def extents_are_valid(extents: ImageExtents, limits: ImageExtents) -> bool:
+    return all(extents[i] <= limits[i] for i in range(len(extents)))
+
+
+def valid_texture_memory_layouts(
+    tensor_sizes: torch.Size, texture_limits: ImageExtents
+) -> Set[VkMemoryLayout]:
+    """
+    Given tensor sizes, determine the set of memory layouts which will prodice a texture
+    that can fit within the specified device limits.
+    """
+    valid_layouts = set()
+    for layout in list(all_memory_layouts):
+        extents = required_image_extents(tensor_sizes, layout)
+        if extents_are_valid(extents, texture_limits):
+            valid_layouts.add(layout)
+
+    return valid_layouts
+
+
+class TensorRepr:
+    """
+    This class is a wrapper around a pair of VkStorageType and VkMemoryLayout which
+    describes how a tensor should be represented in the Vulkan Delegate.
+    """
+
+    def __init__(self, storage_type: VkStorageType, memory_layout: VkMemoryLayout):
+        self.storage_type = storage_type
+        self.memory_layout = memory_layout
+
+    def __str__(self) -> str:
+        return f"TensorRepr({self.storage_type}, {self.memory_layout})"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TensorRepr):
+            return NotImplemented
+        return (
+            self.storage_type == other.storage_type
+            and self.memory_layout == other.memory_layout
+        )
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+
+class TensorReprList:
+    """
+    This class is a wrapper around a list of TensorRepr instances that automatically
+    applies a "broadcasting" mechanism. The broadcasting mechanism allows for a single
+    underlying TensorRepr to be used to represent multiple tensors.
+    """
+
+    def __init__(self, tensor_reprs: Union[TensorRepr, List[TensorRepr]]):
+        self.vals: List[TensorRepr] = (
+            tensor_reprs if isinstance(tensor_reprs, list) else [tensor_reprs]
+        )
+
+    def __len__(self):
+        return len(self.vals)
+
+    def __getitem__(self, idx: int) -> TensorRepr:
+        if idx > 0 and len(self) == 1:
+            return self.vals[0]
+        else:
+            return self.vals[idx]
+
+    def __setitem__(self, idx: int, val: TensorRepr) -> None:
+        if idx > 0 and len(self) == 1:
+            self.vals[0] = val
+        else:
+            self.vals[idx] = val
+
+    def __str__(self) -> str:
+        return f"[{', '.join(str(ts) for ts in self.vals)}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TensorReprList):
+            return NotImplemented
+
+        if len(self) == len(other):
+            for self_val, other_val in zip(self.vals, other.vals):
+                if self_val != other_val:
+                    return False
+
+            return True
+
+        return False
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+    def append(self, val: TensorRepr) -> None:
+        self.vals.append(val)
+
+    def storage_type(self, idx: int = 0) -> VkStorageType:
+        return self.vals[idx].storage_type
+
+    def memory_layout(self, idx: int = 0) -> VkMemoryLayout:
+        return self.vals[idx].memory_layout
+
+
+class TensorRepSet:
+    """
+    This class describes the possible set of representations (i.e. TensorRepr) that may
+    be used to represent a tensor. This set is determined by the implementation of the
+    operator that the tensor participates in as well as the texture extents of the GPU.
+    """
+
+    def __init__(
+        self,
+        buffer_memory_layouts: Set[VkMemoryLayout],
+        texture_memory_layouts: Set[VkMemoryLayout],
+    ):
+        self.valid_buffer_layouts = buffer_memory_layouts
+        self.valid_texture_layouts = texture_memory_layouts
+
+    def __str__(self) -> str:
+        buffer_layouts = ", ".join(layout.name for layout in self.valid_buffer_layouts)
+        texture_layouts = ", ".join(
+            layout.name for layout in self.valid_texture_layouts
+        )
+        return f"TensorRepSet(Buffer Layouts: [{buffer_layouts}], Texture Layouts: [{texture_layouts}])"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TensorRepSet):
+            return NotImplemented
+        return (
+            self.valid_buffer_layouts == other.valid_buffer_layouts
+            and self.valid_texture_layouts == other.valid_texture_layouts
+        )
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+    def copy(self) -> "TensorRepSet":
+        return TensorRepSet(
+            set(self.valid_buffer_layouts), set(self.valid_texture_layouts)
+        )
+
+    def is_empty(self) -> bool:
+        """
+        A TensorRepSet is "empty" if there are no valid representations of the tensor.
+        """
+        return (
+            len(self.valid_buffer_layouts) == 0 and len(self.valid_texture_layouts) == 0
+        )
+
+    def make_intersect(self, other: "TensorRepSet") -> "TensorRepSet":
+        """
+        Merge this TensorRepr with another TensorRepr, returning a new TensorRepr
+        with the intersection of the two.
+        """
+        return TensorRepSet(
+            self.valid_buffer_layouts & other.valid_buffer_layouts,
+            self.valid_texture_layouts & other.valid_texture_layouts,
+        )
+
+    def make_union(self, other: "TensorRepSet") -> "TensorRepSet":
+        """
+        Merge this TensorRepSet with another TensorRepSet, returning a new TensorRepSet
+        with the union of the two.
+        """
+        return TensorRepSet(
+            self.valid_buffer_layouts | other.valid_buffer_layouts,
+            self.valid_texture_layouts | other.valid_texture_layouts,
+        )
+
+    def is_compatible(self, storage: TensorRepr) -> bool:
+        """
+        Check if this TensorRepr is compatible with the given TensorRepSet.
+        """
+        if storage.storage_type == VkStorageType.BUFFER:
+            return storage.memory_layout in self.valid_buffer_layouts
+        elif storage.storage_type == VkStorageType.TEXTURE_3D:
+            return storage.memory_layout in self.valid_texture_layouts
+        else:
+            raise RuntimeError(f"Unsupported storage type {storage.storage_type}")
+
+    def any_in_common(self, other: "TensorRepSet") -> bool:
+        """
+        Check if this TensorRepr has any representations in common with another
+        TensorRepr.
+        """
+        return (
+            len(self.valid_buffer_layouts & other.valid_buffer_layouts) > 0
+            or len(self.valid_texture_layouts & other.valid_texture_layouts) > 0
+        )
+
+    def texture_is_valid(self):
+        return len(self.valid_texture_layouts) > 0
+
+    def buffer_is_valid(self):
+        return len(self.valid_buffer_layouts) > 0
+
+    def first_valid_buffer_layout(self):
+        return list(self.valid_buffer_layouts)[0]
+
+    def first_valid_texture_layout(self):
+        return list(self.valid_texture_layouts)[0]
+
+    def make_tensor_repr(self) -> TensorRepr:
+        """
+        Pick a representation (i.e. TensorRepr) from the set of possible representations.
+        If there are multiple valid representations, then:
+        1. Prefer texture storage over buffer storage
+        2. Pick the first available memory layout.
+        """
+        if self.is_empty():
+            # An empty repset typically means that it is associated with a weight tensor
+            # or non tensor argument. In this case, just return default storage and
+            # layout as placeholder.
+            return TensorRepr(
+                VkStorageType.DEFAULT_STORAGE, VkMemoryLayout.DEFAULT_LAYOUT
+            )
+
+        if self.texture_is_valid():
+            return TensorRepr(
+                VkStorageType.TEXTURE_3D, self.first_valid_texture_layout()
+            )
+
+        else:
+            return TensorRepr(VkStorageType.BUFFER, self.first_valid_buffer_layout())
+
+    def is_constrained(self) -> bool:
+        """
+        A "constrained" RepSet is one that has either:
+        1. A single valid texture memory layout, and no valid buffer memory layouts
+        2. No valid texture memory layouts, and a single valid buffer memory layout
+        3. Is empty
+
+        In this case, it is unambiguous which representation should be used for the
+        tensor.
+        """
+        if self.is_empty():
+            return True
+        elif (
+            len(self.valid_texture_layouts) == 1 and len(self.valid_buffer_layouts) == 0
+        ):
+            return True
+        elif (
+            len(self.valid_texture_layouts) == 0 and len(self.valid_buffer_layouts) == 1
+        ):
+            return True
+        else:
+            return False
+
+    def is_ambiguous(self) -> bool:
+        """
+        An "ambiguous" RepSet is one that is not constrained.
+        """
+        return not self.is_constrained()
+
+    def _possible_pdis(self) -> Set[PackedDimInfo]:
+        buffer_set = set()
+        texture_set = set()
+        for layout in self.valid_buffer_layouts:
+            buffer_set.add(PackedDimInfo.from_repr(layout, VkStorageType.BUFFER))
+        for layout in self.valid_texture_layouts:
+            texture_set.add(PackedDimInfo.from_repr(layout, VkStorageType.TEXTURE_3D))
+        return buffer_set, texture_set
+
+    def has_same_packed_dim_info_set(self, other: "TensorRepSet") -> bool:
+        """
+        Check if self and other produce the exact same sets of PackedDimInfo
+        for both buffer and texture storage types. Completely empty repsets
+        (no layouts for any storage type) are treated as matching any other
+        repset.
+        """
+        other_buf_set, other_tex_set = other._possible_pdis()
+        buf_set, tex_set = self._possible_pdis()
+
+        # A completely empty repset is compatible with anything
+        if not buf_set and not tex_set:
+            return True
+        if not other_buf_set and not other_tex_set:
+            return True
+
+        return other_buf_set == buf_set and other_tex_set == tex_set
+
+    def has_compatible_packed_dim_info_set(self, other: "TensorRepSet") -> bool:
+        """
+        Check if all PackedDimInfos from other are contained within self's
+        PackedDimInfo sets, i.e. self is a superset of other for both buffer
+        and texture PDI sets.
+        """
+        other_buf_set, other_tex_set = other._possible_pdis()
+        buf_set, tex_set = self._possible_pdis()
+
+        for pdi in other_buf_set:
+            if pdi not in buf_set:
+                return False
+
+        for pdi in other_tex_set:
+            if pdi not in tex_set:
+                return False
+
+        return True
+
+    def constrain_to_compatible_packed_dim(
+        self, other: "TensorRepSet"
+    ) -> "TensorRepSet":
+        """
+        Return a new TensorRepSet containing only layouts from self whose
+        PackedDimInfo is present in other's PackedDimInfo sets. If other is
+        completely empty, return a copy of self unchanged. If other has layouts
+        for only one storage type, layouts for the missing storage type are
+        also removed.
+        """
+        other_buf_set, other_tex_set = other._possible_pdis()
+
+        # Completely empty other means no constraint
+        if not other_buf_set and not other_tex_set:
+            return self.copy()
+
+        new_buf = {
+            layout
+            for layout in self.valid_buffer_layouts
+            if other_buf_set
+            and PackedDimInfo.from_repr(layout, VkStorageType.BUFFER) in other_buf_set
+        }
+        new_tex = {
+            layout
+            for layout in self.valid_texture_layouts
+            if other_tex_set
+            and PackedDimInfo.from_repr(layout, VkStorageType.TEXTURE_3D)
+            in other_tex_set
+        }
+        return TensorRepSet(new_buf, new_tex)
+
+
+def make_tensor_repset(tensor_repr: TensorRepr) -> TensorRepSet:
+    """
+    Given a TensorRepr, return a TensorRepSet that contains only that TensorRepr
+    """
+    if tensor_repr.storage_type == VkStorageType.BUFFER:
+        return TensorRepSet({tensor_repr.memory_layout}, set())
+    elif tensor_repr.storage_type == VkStorageType.TEXTURE_3D:
+        return TensorRepSet(set(), {tensor_repr.memory_layout})
+    else:
+        raise RuntimeError(f"Unsupported storage type {tensor_repr.storage_type}")
+
+
+def upper_bound_size(dim: Union[int, torch.SymInt]) -> Optional[int]:
+    """Largest value a (possibly symbolic) tensor dimension can take.
+
+    Returns None if no finite bound is known.
+
+    A symbolic dim compares against a limit using its *hint* -- the size of the
+    example input the model happened to be traced with -- not the maximum the
+    exported range allows. Sizing decisions must use the bound instead, or a
+    model traced with a small example will make a choice that is invalid once it
+    runs at a larger size.
+    """
+    if not isinstance(dim, torch.SymInt):
+        return int(dim)
+    if not dim.node.expr.free_symbols:
+        return int(dim.node.expr)
+    shape_env = dim.node.shape_env
+    if shape_env is None:
+        return None
+    try:
+        upper = shape_env.bound_sympy(dim.node.expr).upper
+    except Exception:
+        return None
+    if upper is None or not upper.is_finite:
+        return None
+    return int(upper)
+
+
+def filter_invalid_reprs(
+    tensor_val: FakeTensor,
+    tensor_repset: TensorRepSet,
+    texture_limits: ImageExtents,
+) -> TensorRepSet:
+    """
+    `tensor_val` represents an actual tensor participating in some operator computation.
+
+    `tensor_repset` represents the set of valid tensor representations that may be used
+    for that tensor that is supported by the op implementation.
+
+    `texture_limits` represents the maximum texture sizes that is supported by the GPU.
+
+    Given the above, return a new TensorRepSet that contains only texture layouts that
+    can be used to produce a valid image texture for the given tensor (i.e. fits within
+    texture limits).
+    """
+    # Size the texture by what the dimension CAN be, not by the example the
+    # model was traced with. An unbounded dim cannot be shown to fit, so it
+    # falls back to buffer storage.
+    bounds = [upper_bound_size(d) for d in tensor_val.shape]
+    valid_texture_layouts = set()
+    if all(b is not None for b in bounds):
+        max_shape = torch.Size(bounds)
+        for memory_layout in tensor_repset.valid_texture_layouts:
+            extents = required_image_extents(max_shape, memory_layout)
+            if extents_are_valid(extents, texture_limits):
+                valid_texture_layouts.add(memory_layout)
+
+    # High dimensional tensors require buffer storage
+    if len(tensor_val.shape) > 4:
+        return TensorRepSet(tensor_repset.valid_buffer_layouts, set())
+
+    return TensorRepSet(tensor_repset.valid_buffer_layouts, valid_texture_layouts)
+
+
+def filter_invalid_reprs_for_node_list(
+    arg_repsets: TensorRepSet,
+    arg_node: List[torch.fx.Node],
+    texture_limits: ImageExtents,
+) -> TensorRepSet:
+    """
+    Wrapper around filter_invalid_reprs for a list of nodes. This will happen
+    for the cat operator, where the first argument is a list of nodes.
+    """
+    # For variable length args, assume that they all need to use the same representation
+    # only one repset should be defined
+    common_tensor_repsets = arg_repsets
+
+    for n in arg_node:
+        assert isinstance(n, torch.fx.Node)
+        common_tensor_repsets = common_tensor_repsets.make_intersect(
+            filter_invalid_reprs(n.meta["val"], common_tensor_repsets, texture_limits)
+        )
+
+    return common_tensor_repsets
+
+
+## Convenience TensorRepSet definitions
+
+# Only includes memory layouts that can be used by non-quantized tensors
+
+CONTIGUOUS_ANY = TensorRepSet(
+    {VkMemoryLayout.TENSOR_WIDTH_PACKED}, {VkMemoryLayout.TENSOR_WIDTH_PACKED}
+)
+CONTIGUOUS_BUFFER = TensorRepSet({VkMemoryLayout.TENSOR_WIDTH_PACKED}, set())
+
+WIDTH_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_WIDTH_PACKED})
+HEIGHT_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_HEIGHT_PACKED})
+CHANNELS_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_CHANNELS_PACKED})
+
+CHANNELS_PACKED_ANY = TensorRepSet(
+    {VkMemoryLayout.TENSOR_CHANNELS_PACKED}, {VkMemoryLayout.TENSOR_CHANNELS_PACKED}
+)
+
+CHANNELS_PACKED_TEXTURE_OR_CONTIGUOUS_BUFFER = TensorRepSet(
+    {VkMemoryLayout.TENSOR_WIDTH_PACKED}, {VkMemoryLayout.TENSOR_CHANNELS_PACKED}
+)
+
+ANY_TEXTURE = TensorRepSet(set(), all_memory_layouts)
+ANY_BUFFER = TensorRepSet(all_memory_layouts, set())
+ANY_STORAGE = TensorRepSet(all_memory_layouts, all_memory_layouts)
+
+# Only includes memory layouts that can be used by quantized tensors
+
+PACKED_INT8_BUFFER = TensorRepSet(all_quantized_memory_layouts, set())
+PACKED_INT8_4W4C_BUFFER = TensorRepSet({VkMemoryLayout.PACKED_INT8_4W4C}, set())
+PACKED_INT8_4H4W_BUFFER = TensorRepSet({VkMemoryLayout.PACKED_INT8_4H4W}, set())
+PACKED_INT8_4W_BUFFER = TensorRepSet({VkMemoryLayout.PACKED_INT8_4W}, set())
+PACKED_INT8_4C1W_BUFFER = TensorRepSet({VkMemoryLayout.PACKED_INT8_4C1W}, set())
+
+PACKED_INT8_CONV2D_BUFFER = TensorRepSet({VkMemoryLayout.PACKED_INT8_CONV2D}, set())
+
+PACKED_INT8_CHANNELS_PACKED_BUFFER = TensorRepSet(
+    {
+        VkMemoryLayout.PACKED_INT8_4W4C,
+        VkMemoryLayout.PACKED_INT8_4C1W,
+        VkMemoryLayout.PACKED_INT8_CONV2D,
+    },
+    set(),
+)
+
+
+# Special use RepSets
+
+NO_STORAGE = TensorRepSet(set(), set())
+# Buffer side admits both float and quantized (PACKED_INT8_*) layouts; texture side
+# is float-only because the Vulkan backend has no quantized texture support
+# (required_image_extents and the texture indexing helpers only know about the
+# float layouts). Used as an intersection identity (e.g. common_arg_repset
+# accumulator) and as a placeholder for non-tensor / not-yet-prepacked args, so
+# narrowing the texture side is non-breaking for those uses while letting it act
+# as a true universal set when intersected against quant-aware repsets.
+ANY_STORAGE_INCL_PACKED_INT8 = TensorRepSet(
+    universal_memory_layout_set, all_memory_layouts
+)
+
+
+class TensorRepSetList:
+    """
+    This class is a wrapper around a list of TensorRepSet instances that automatically
+    applies a "broadcasting" mechanism. The broadcasting mechanism allows for a single
+    underlying TensorRepSet to be used for multiple tensors.
+    """
+
+    def __init__(
+        self,
+        tensor_repsets: Union[TensorRepSet, List[TensorRepSet]],
+    ):
+        self.vals: List[TensorRepSet] = (
+            tensor_repsets if isinstance(tensor_repsets, list) else [tensor_repsets]
+        )
+
+    def __len__(self):
+        return len(self.vals)
+
+    def __getitem__(self, idx: int) -> TensorRepSet:
+        if idx > 0 and len(self) == 1:
+            return self.vals[0]
+        if idx >= len(self.vals):
+            return set()
+        return self.vals[idx]
+
+    def __setitem__(self, idx: int, val: TensorRepSet) -> None:
+        if idx > 0 and len(self.vals) == 1:
+            self.vals[0] = val
+        else:
+            self.vals[idx] = val
+
+    def __str__(self) -> str:
+        return f"[{', '.join(str(ts) for ts in self.vals)}]"
+
+    def append(self, val: TensorRepSet) -> None:
+        return self.vals.append(val)
+
+    def any_is_empty(self) -> bool:
+        if len(self.vals) == 0:
+            return True
+
+        return any(tensor_repr.is_empty() for tensor_repr in self.vals)
+
+
+class OpRepSets:
+    """
+    This class is responsible for representing and managing the set of valid tensor
+    representations that may be used for all input and output tensors of an operator.
+    It is also responsible for maintaining synchronization rules between tensors
+    participating in the computation.
+
+    Currently, three synchronization rules exist:
+    1. All input tensors must use the same representation (e.g. binary ops)
+    2. The "primary" input and output tensors must use the same representation
+       (e.g. group norm; the output is a tuple of out, mean, rstd; out must be the same
+       representation as the first input x, but mean and rstd may use different
+       representations as out)
+    3. All output tensors must use the same representation (e.g. choose qparams)
+
+    Note that "primary" input and output tensor refers to the first non-weight input
+    tensor and the first output tensor. Note that Some operators (such as arange) do not
+    have any tensor inputs.
+
+    Currently, the above three synchronization rules are sufficient to describe the
+    representation requirements of all ET-VK operators.
+
+    This class also provides utilities to constrain the repsets; when applying the
+    constraints, the synchronization rules will be maintained.
+    """
+
+    def __init__(  # noqa: C901
+        self,
+        inputs_repsets: TensorRepSetList,
+        outputs_repsets: TensorRepSetList,
+        op_node: torch.fx.Node,
+        texture_limits: ImageExtents,
+    ):
+        self.op_node = op_node
+
+        # inputs_repset_list is received from the operator registration. If a different
+        # repset is defined for each input tensor, then assume that the input tensor
+        # representations do not need to be synchronized.
+        if len(inputs_repsets) > 1:
+            self.sync_args_repr = False
+        # Otherwise, default to True
+        else:
+            self.sync_args_repr = True
+
+        # outputs_repset_list is received from the operator registration. If a different
+        # repset is defined for each output tensor, then assume that the output tensor
+        # representations do not need to be synchronized.
+        if len(outputs_repsets) > 1:
+            self.sync_outs_repr = False
+        else:
+            self.sync_outs_repr = True
+
+        # Try to determine the index of the "primary" argument, i.e. the first non
+        # constant tensor argument. For the vast majority of operators with tensor
+        # arguments, this will be the first argument.
+        self.primary_arg_idx: Optional[int] = None
+        for i, arg_node in enumerate(self.op_node.args):
+            arg_node_repset = inputs_repsets[i]
+            if not is_tensor_arg_node(arg_node):
+                continue
+            if arg_node_repset is None:
+                continue
+            if arg_node_repset.is_empty():
+                continue
+
+            self.primary_arg_idx = i
+            break
+
+        # If the repset of the primary input and the primary output are the same, then
+        # assume they need to be the same.
+        self.sync_primary_io_repr = self.primary_arg_idx is not None
+        if self.primary_arg_idx is not None:
+            if inputs_repsets[self.primary_arg_idx] != outputs_repsets[0]:
+                self.sync_primary_io_repr = False
+
+        # Now, go through the arguments of the operator and create a filtered repset
+        # for each based on the actual tensor value.
+        args_repset_list = TensorRepSetList([])
+        common_arg_repset = ANY_STORAGE_INCL_PACKED_INT8
+        for i, arg_node in enumerate(op_node.args):
+            arg_repset = inputs_repsets[i]
+
+            # Use ANY_STORAGE_INCL_PACKED_INT8 for non-tensor nodes so they don't cause the op
+            # repsets to appear empty
+            if not is_tensor_arg_node(arg_node):
+                args_repset_list.append(ANY_STORAGE_INCL_PACKED_INT8)
+            # NO_STORAGE is used to denote that an input is either a non tensor arg or
+            # a weight tensor that is not prepacked. Similar to the above, use
+            # ANY_STORAGE_INCL_PACKED_INT8 in this case.
+            elif arg_repset.is_empty():
+                args_repset_list.append(ANY_STORAGE_INCL_PACKED_INT8)
+            else:
+                assert not arg_repset.is_empty()
+
+                arg_repset = self.filter_invalid_reprs_for_arg(
+                    arg_repset, arg_node, texture_limits
+                )
+
+                args_repset_list.append(arg_repset)
+                common_arg_repset = common_arg_repset.make_intersect(arg_repset)
+
+        # Repeat for output tensors.
+        outs_repset_list = TensorRepSetList([])
+        common_out_repset = ANY_STORAGE_INCL_PACKED_INT8
+        if num_tensors_in_node(op_node) == 1:
+            out_val = op_node.meta["val"]
+            # num_tensors_in_node counts tensors, not nesting: an op declared
+            # to return Tensor[] still lands here when it happens to produce
+            # exactly one, and meta["val"] is then a one-element list rather
+            # than a bare FakeTensor.
+            if isinstance(out_val, (list, tuple)):
+                out_val = out_val[0]
+            common_out_repset = filter_invalid_reprs(
+                out_val, outputs_repsets[0], texture_limits
+            )
+            outs_repset_list.append(common_out_repset)
+        # Multiple output tensors
+        else:
+            for i, val in enumerate(op_node.meta["val"]):
+                assert isinstance(val, FakeTensor)
+                out_repset = filter_invalid_reprs(
+                    val, outputs_repsets[i], texture_limits
+                )
+
+                outs_repset_list.append(out_repset)
+                common_out_repset = common_out_repset.make_intersect(out_repset)
+
+        # Apply synchronization rules between the primary input and output
+        primary_repset = NO_STORAGE
+        if self.sync_primary_io_repr:
+            primary_in_repset = (
+                common_arg_repset
+                if self.sync_args_repr
+                else args_repset_list[self.primary_arg_idx]
+            )
+            primary_out_repset = (
+                common_out_repset if self.sync_outs_repr else outs_repset_list[0]
+            )
+            primary_repset = primary_in_repset.make_intersect(primary_out_repset)
+
+            args_repset_list[self.primary_arg_idx] = primary_repset.copy()
+            outs_repset_list[0] = primary_repset.copy()
+
+        # Apply synchronization rules; if either all inputs/outputs must use the same
+        # representation, then only use a single underlying repset.
+        if self.sync_args_repr:
+            common_repset = (
+                primary_repset if self.sync_primary_io_repr else common_arg_repset
+            )
+
+            for i in range(len(args_repset_list)):
+                args_repset_list[i] = common_repset.copy()
+
+        if self.sync_outs_repr:
+            common_repset = (
+                primary_repset if self.sync_primary_io_repr else common_out_repset
+            )
+
+            for i in range(len(outs_repset_list)):
+                outs_repset_list[i] = common_repset.copy()
+
+        # Save the resulting repsets
+        self.args_repset_list = args_repset_list
+        self.outs_repset_list = outs_repset_list
+
+        # Check that synchronization rules are respected.
+        self.assert_sync_contraints()
+
+    def __str__(self) -> str:
+        return f"OpRepSets(ins={self.args_repset_list}, outs={self.outs_repset_list})"
+
+    def filter_invalid_reprs_for_arg(
+        self, arg_repsets: TensorRepSet, arg_node: Any, texture_limits: ImageExtents
+    ) -> TensorRepSet:
+        """
+        Helper function to call filter_invalid_reprs
+        """
+        if isinstance(arg_node, torch.fx.Node) and is_single_tensor_node(arg_node):
+            return filter_invalid_reprs(
+                arg_node.meta["val"], arg_repsets, texture_limits
+            )
+        elif isinstance(arg_node, list) and all(
+            is_single_tensor_node(n) for n in arg_node
+        ):
+            return filter_invalid_reprs_for_node_list(
+                arg_repsets, arg_node, texture_limits
+            )
+        # Special case for getitem; return the repset of the particular val in the
+        # list of tensors that is being extracted.
+        elif (
+            self.op_node.target == operator.getitem and arg_node == self.op_node.args[0]
+        ):
+            idx = self.op_node.args[1]
+            assert isinstance(idx, int)
+            return filter_invalid_reprs(
+                arg_node.meta["val"][idx], arg_repsets, texture_limits
+            )
+
+        raise NotImplementedError(f"Unhandled node type {arg_node}")
+
+    def assert_sync_contraints(self) -> None:
+        if self.sync_args_repr:
+            for i in range(len(self.args_repset_list)):
+                for j in range(i + 1, len(self.args_repset_list)):
+                    ri = self.args_repset_list[i]
+                    rj = self.args_repset_list[j]
+                    if not ri.is_empty() and not rj.is_empty():
+                        assert ri.has_compatible_packed_dim_info_set(
+                            rj
+                        ), f"Synced arg repsets {i} and {j} have incompatible packed dim info: {ri} vs {rj}"
+
+        if self.sync_outs_repr:
+            for i in range(len(self.outs_repset_list)):
+                for j in range(i + 1, len(self.outs_repset_list)):
+                    ri = self.outs_repset_list[i]
+                    rj = self.outs_repset_list[j]
+                    if not ri.is_empty() and not rj.is_empty():
+                        assert ri.has_compatible_packed_dim_info_set(
+                            rj
+                        ), f"Synced out repsets {i} and {j} have incompatible packed dim info: {ri} vs {rj}"
+
+        if self.sync_primary_io_repr:
+            primary_arg = self.args_repset_list[self.primary_arg_idx]
+            primary_out = self.outs_repset_list[0]
+            if not primary_arg.is_empty() and not primary_out.is_empty():
+                assert primary_arg.has_compatible_packed_dim_info_set(
+                    primary_out
+                ), f"Primary arg and out repsets have incompatible packed dim info: {primary_arg} vs {primary_out}"
+
+    def any_is_empty(self) -> bool:
+        return (
+            self.args_repset_list.any_is_empty() or self.outs_repset_list.any_is_empty()
+        )
+
+    def get_arg_repset(self, i: int):
+        return self.args_repset_list[i]
+
+    def get_out_repset(self, i: int):
+        return self.outs_repset_list[i]
+
+    def try_constrain_with_arg_repset(
+        self, arg_i: int, source_repset: TensorRepSet
+    ) -> bool:
+        """
+        Attempt to constrain the repsets of the tensors participating in this operator
+        based on an "existing" repset of an argument. The existing repset can have two
+        sources:
+        * A representation may have been determined for the argument already from a
+          prior operator
+        * The output repset of the operator which produces the argument
+
+        If the existing repset of the argument is compatible with the current operator,
+        then constrain the repsets of this operator and apply synchronization rules.
+
+        This process tries to minimize the number of transition nodes that will need to
+        be inserted by tag_memory_meta_pass.py by maintaining existing representations
+        for as long as possible.
+        """
+        arg_current_repset = self.args_repset_list[arg_i]
+
+        if arg_current_repset == source_repset:
+            return False
+
+        if not arg_current_repset.any_in_common(source_repset):
+            return False
+
+        # Compute the narrowed repset (intersection of current arg and source).
+        narrowed = arg_current_repset.make_intersect(source_repset)
+
+        if self.sync_primary_io_repr:
+            # Check that the narrowed result is compatible with the output.
+            # Using the intersection rather than the raw source_repset avoids
+            # rejecting valid constraints where the source has extra layouts
+            # (e.g. ANY_TEXTURE includes HP/CP) that don't exist in the output
+            # but also don't appear in the intersection.
+            if not self.get_out_repset(0).has_compatible_packed_dim_info_set(narrowed):
+                return False
+
+        # If this point is reached, then it is possible to constrain
+        self.args_repset_list[arg_i] = narrowed
+
+        # Propagate to other synced args via packed-dim compatibility
+        if self.sync_args_repr:
+            for i in range(len(self.args_repset_list)):
+                if i != arg_i:
+                    self.args_repset_list[i] = self.args_repset_list[
+                        i
+                    ].constrain_to_compatible_packed_dim(narrowed)
+
+        # Propagate to output via packed-dim compatibility
+        if self.sync_primary_io_repr and (
+            arg_i == self.primary_arg_idx or self.sync_args_repr
+        ):
+            self.outs_repset_list[0] = self.outs_repset_list[
+                0
+            ].constrain_to_compatible_packed_dim(narrowed)
+
+            # Propagate to other synced outputs via packed-dim compatibility
+            if self.sync_outs_repr:
+                for i in range(len(self.outs_repset_list)):
+                    if i != 0:
+                        self.outs_repset_list[i] = self.outs_repset_list[
+                            i
+                        ].constrain_to_compatible_packed_dim(self.outs_repset_list[0])
+
+        self.assert_sync_contraints()
+        return True
+
+    def try_constrain_with_out_repset(self, required_repset: TensorRepSet) -> bool:
+        """
+        Attempt to constrain the output repsets of the tensors participating in this
+        operator based the repset required by a downstream operator.
+        """
+        out_current_repset = self.outs_repset_list[0]
+
+        if out_current_repset == required_repset:
+            return False
+
+        if not out_current_repset.any_in_common(required_repset):
+            return False
+
+        narrowed = out_current_repset.make_intersect(required_repset)
+        self.outs_repset_list[0] = narrowed
+
+        # Propagate to other synced outputs via packed-dim compatibility
+        if self.sync_outs_repr:
+            for i in range(len(self.outs_repset_list)):
+                if i != 0:
+                    self.outs_repset_list[i] = self.outs_repset_list[
+                        i
+                    ].constrain_to_compatible_packed_dim(narrowed)
+
+        # Propagate to primary arg via packed-dim compatibility
+        if self.sync_primary_io_repr:
+            self.args_repset_list[self.primary_arg_idx] = self.args_repset_list[
+                self.primary_arg_idx
+            ].constrain_to_compatible_packed_dim(narrowed)
+
+            # Propagate to other synced args via packed-dim compatibility
+            if self.sync_args_repr:
+                for i in range(len(self.args_repset_list)):
+                    if i != self.primary_arg_idx:
+                        self.args_repset_list[i] = self.args_repset_list[
+                            i
+                        ].constrain_to_compatible_packed_dim(
+                            self.args_repset_list[self.primary_arg_idx]
+                        )
+
+        self.assert_sync_contraints()
+        return True
+
+    def pick_representations(self) -> Tuple[TensorReprList, TensorReprList]:
+        """
+        For each tensor participating in the op, pick a representation for it among the
+        possible represetntation sets.
+        """
+        args_repr_list = TensorReprList([])
+        outs_repr_list = TensorReprList([])
+
+        for i in range(len(self.op_node.args)):
+            arg_repset = self.args_repset_list[i]
+            args_repr_list.append(arg_repset.make_tensor_repr())
+
+        for i in range(num_tensors_in_node(self.op_node)):
+            out_repset = self.outs_repset_list[i]
+            outs_repr_list.append(out_repset.make_tensor_repr())
+
+        return args_repr_list, outs_repr_list
+
+
+##
+## TensorSpec Utils
+##
+
+
+def has_node_spec_attr(node: torch.fx.Node, attr: str) -> bool:
+    return "spec" in node.meta and hasattr(node.meta["spec"], attr)
+
+
+def set_node_spec_attr(node: torch.fx.Node, attr: str, value):
+    assert "spec" in node.meta
+    spec = node.meta["spec"]
+    if isinstance(spec, TensorSpec):
+        setattr(spec, attr, value)
+    elif isinstance(spec, (list, tuple)):
+        # Special case if value is a list/tuple of the same length as the
+        # collection of tensors in the node. In this case, treat the value list
+        # as a list of values to set indivudually for each tensor in the node
+        if isinstance(value, (list, tuple)) and len(spec) == len(value):
+            assert len(spec) == len(value)
+            for s, v in zip(spec, value):
+                assert isinstance(s, TensorSpec)
+                setattr(s, attr, v)
+        # Otherwise, set the attribute to value for all tensors in the list
+        else:
+            for s in spec:
+                assert isinstance(s, TensorSpec)
+                setattr(s, attr, value)
+    else:
+        raise RuntimeError(f"Cannot set attr for spec of type {type(spec)}")
+
+
+def get_node_spec_attr(node: torch.fx.Node, attr: str, return_first: bool = True):
+    assert "spec" in node.meta
+    spec = node.meta["spec"]
+    if isinstance(spec, TensorSpec):
+        return getattr(spec, attr) if hasattr(spec, attr) else None
+    elif isinstance(spec, (list, tuple)):
+        if return_first:
+            return getattr(spec[0], attr) if hasattr(spec[0], attr) else None
+        else:
+            return [getattr(s, attr) if hasattr(s, attr) else None for s in spec]
+    else:
+        raise RuntimeError(f"Cannot get attr for spec of type {type(spec)}")
+
+
+def get_node_storage_type(node: torch.fx.Node) -> Optional[VkStorageType]:
+    return get_node_spec_attr(node, "vk_storage_type")
+
+
+def get_node_memory_layout(node: torch.fx.Node) -> Optional[VkMemoryLayout]:
+    return get_node_spec_attr(node, "vk_memory_layout")
+
+
+def has_node_repr(node) -> bool:
+    if isinstance(node, (list, tuple)):
+        return all(has_node_spec_attr(n, "etvk_node_repr") for n in node)
+    else:
+        return has_node_spec_attr(node, "etvk_node_repr")
+
+
+def set_node_repr(node: torch.fx.Node, node_repr: Union[TensorRepr, TensorReprList]):
+    if isinstance(node_repr, TensorReprList):
+        # Convert to a regular list so taht `set_node_spec_attr` can attach each entry
+        # to a separate TensorSpec
+        node_repr_list = [node_repr[i] for i in range(num_tensors_in_node(node))]
+        set_node_spec_attr(node, "etvk_node_repr", node_repr_list)
+    else:
+        set_node_spec_attr(node, "etvk_node_repr", node_repr)
+
+
+def get_node_repr(node) -> Union[TensorRepr, TensorReprList]:
+    if isinstance(node, (list, tuple)):
+        raise NotImplementedError("get_node_repr not implemented for list of nodes")
+    else:
+        return get_node_spec_attr(node, "etvk_node_repr", False)
+
+
+##
+## Graph Pattern Matching
+##
+
+
+def maybe_skip_q_dq_arg_chain(
+    arg: torch.fx.node.Argument,
+) -> Tuple[Optional[torch.fx.Node], Optional[torch.fx.Node], Optional[torch.fx.Node]]:
+    """
+    Check if the given node argument is part of a Quantize/Dequantize chain produced by
+    the quant workflow. If so, return the source tensor that is the input to the Q/DQ
+    chain and the quantize/dequantize nodes in the chain. Otherwise, return the argument
+    as is and None, None
+    """
+    if not isinstance(arg, torch.fx.Node):
+        return None, None, None
+
+    # If the arg is a view copy node, check if the original node is a dequant node
+    if is_dequant_node(arg) or (
+        is_view_copy_node(arg) and is_dequant_node(arg.args[0])  # pyre-ignore[6]
+    ):
+        dequant_node = arg
+        if is_view_copy_node(arg):
+            dequant_node = arg.args[0]
+
+        quant_node = dequant_node.args[0]  # pyre-ignore[16]
+        assert isinstance(quant_node, torch.fx.Node)
+        source_arg = quant_node.args[0]
+        assert isinstance(source_arg, torch.fx.Node)
+        assert isinstance(dequant_node, torch.fx.Node)
+        return source_arg, quant_node, dequant_node
+    else:
+        return arg, None, None
+
+
+def trace_args_until_placeholder(
+    node: torch.fx.node.Argument, max_search_depth: int = 4
+) -> Tuple[Optional[torch.fx.Node], List[torch.fx.Node]]:
+    """
+    Trace through node.args[0] of a given initial node until a placeholder node is found
+    then return it and the list of nodes traversed. If no placeholder node is found,
+    returns None and an empty list.
+    """
+    cur_node = node
+    search_depth = 0
+
+    if not isinstance(cur_node, torch.fx.Node):
+        return None, []
+
+    traversed = [cur_node]
+    while cur_node.op != "placeholder" and search_depth < max_search_depth:
+        # Break if cur_node has no args
+        if len(cur_node.args) == 0:
+            break
+
+        cur_node = cur_node.args[0]
+        if not isinstance(cur_node, torch.fx.Node):
+            break
+        traversed.append(cur_node)
+        search_depth += 1
+
+    if not isinstance(cur_node, torch.fx.Node):
+        return None, []
+    if cur_node.op != "placeholder":
+        return None, []
+
+    assert isinstance(cur_node, torch.fx.Node)
+    return cur_node, traversed
+
+
+def is_in_4bit_range(tensor: torch.Tensor) -> bool:
+    """
+    Check if the given tensor is in the range of 4-bit quantization and is of integer type.
+    """
+    if tensor.dtype not in (torch.int8, torch.uint8):
+        return False
+
+    return tensor.min().item() >= -8 and tensor.max().item() <= 7
+
+
+def is_in_8bit_range(tensor: torch.Tensor) -> bool:
+    """
+    Check if the given tensor is in the range of 4-bit quantization and is of integer type.
+    """
+    if tensor.dtype not in (torch.int8, torch.uint8):
+        return False
+
+    return tensor.min().item() >= -128 and tensor.max().item() <= 127
+
+
+##
+## Misc
+##
+
+
+def normalize_dims(dims: Union[int, List[int]], ndim: int) -> Union[int, List[int]]:
+    """
+    Normalize dimension indices to be non-negative and within [0, ndim).
+    Accepts a single int or a list of ints.
+    """
+    if isinstance(dims, int):
+        if dims < 0:
+            dims += ndim
+
+        return dims
+
+    normalized = []
+    for d in dims:
+        if d < 0:
+            d += ndim
+        normalized.append(d)
+
+    return normalized
+
+
+def nchw_dim_to_whcn_dim(nchw_dim: int, ndim: int) -> int:
+    # Handle negative indices for nchw_dim
+    if nchw_dim < 0:
+        nchw_dim += ndim
+
+    assert nchw_dim >= 0 and nchw_dim < ndim
+    whcn_dim = (ndim - 1) - nchw_dim
+    return whcn_dim
+
+
+def get_tensor_val_str(tensor_val: FakeTensor) -> str:
+    return f"{tensor_val.dtype}: {tensor_val.shape}"
+
+
+def get_node_val_str(node: torch.fx.Node) -> str:
+    if is_single_tensor_node(node):
+        assert isinstance(node.meta["val"], FakeTensor)
+        return get_tensor_val_str(node.meta["val"])
+    elif is_tensor_collection_node(node):
+        assert isinstance(node.meta["val"], (list, tuple))
+        return f"[{', '.join(get_tensor_val_str(t) for t in node.meta['val'])}]"
+    else:
+        if "val" not in node.meta:
+            return str(node)
+        return str(node.meta["val"])
+
+
+def get_arg_node_val_str(arg_node: Any) -> str:
+    if isinstance(arg_node, torch.fx.Node):
+        return get_node_val_str(arg_node)
+    elif isinstance(arg_node, (list, tuple)):
+        return f"[{', '.join(get_arg_node_val_str(n) for n in arg_node)}]"
+    else:
+        return str(arg_node)
+
+
+def node_io_str(node: torch.fx.Node) -> str:
+    target = node.target
+    if isinstance(target, EdgeOpOverload):
+        assert isinstance(target, EdgeOpOverload)
+        target_name = target.__name__
+    elif isinstance(target, torch._ops.OpOverload):
+        assert isinstance(target, torch._ops.OpOverload)
+        target_name = target.name()
+    else:
+        target_name = str(target)
+
+    out_str = f"{get_node_val_str(node)} = {target_name}("
+    for arg in node.args:
+        out_str += get_arg_node_val_str(arg) + ", "
+
+    out_str += " ...)"
+    return out_str
+
+
+def update_program_state_dict(
+    program: ExportedProgram,
+    buffer_name: str,
+    updated_tensor: torch.Tensor,
+) -> None:
+    target_name = None
+    kind = None
+    # Iterate over all the tensors in the graph signature, and find
+    # the one corresponding to the parameter/buffer name
+    for input_ in program.graph_signature.input_specs:
+        if (
+            input_.kind in (InputKind.BUFFER, InputKind.PARAMETER)
+            and isinstance(input_.arg, TensorArgument)
+            and input_.arg.name == buffer_name
+        ):
+            kind = input_.kind
+            target_name = input_.target
+            break
+
+    # Assert that we found the parameter/buffer
+    assert (
+        target_name is not None
+    ), f"could not find {buffer_name} in source program signature"
+    assert target_name in program.state_dict, f"could not find {target_name}"
+
+    if kind == InputKind.PARAMETER:
+        updated_tensor = torch.nn.Parameter(updated_tensor, requires_grad=False)
+
+    # Finally, overwrite the current tensor with updated tensor
+    program.state_dict[target_name] = updated_tensor
+
+
+def align_width_and_update_state_dict(
+    ep: ExportedProgram,
+    node: torch.fx.Node,
+    cur_tensor: torch.Tensor,
+    align_to: int = 4,
+    force_update: bool = False,
+) -> torch.Tensor:
+    """
+    Align the width of the given tensor to the given alignment value and update the
+    state dict of the program with the aligned tensor.
+    """
+    added_padding = False
+    cur_width = cur_tensor.shape[-1]
+    # Only align the width of the tensor if it is not already aligned
+    if cur_width % align_to != 0:
+        num_padding = align_to - (cur_width % align_to)
+        # Align the width of the tensor to the given alignment value
+        aligned_tensor = torch.nn.functional.pad(
+            cur_tensor, (0, num_padding)
+        ).contiguous()
+        added_padding = True
+    else:
+        aligned_tensor = cur_tensor
+
+    if added_padding or force_update:
+        update_program_state_dict(ep, node.name, aligned_tensor)
+        # FakeTensor needs to match updated tensor
+        cur_fake_tensor = node.meta["val"]
+        node.meta["val"] = FakeTensorConverter().from_real_tensor(
+            cur_fake_tensor.fake_mode,
+            aligned_tensor,
+        )
+
+    return aligned_tensor