| # 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 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) |
| 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 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). |
| """ |
| valid_texture_layouts = set() |
| for memory_layout in tensor_repset.valid_texture_layouts: |
| extents = required_image_extents(tensor_val.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()) |
| ALL_STORAGES_REPSET = TensorRepSet( |
| universal_memory_layout_set, universal_memory_layout_set |
| ) |
| |
| |
| 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 = ALL_STORAGES_REPSET |
| for i, arg_node in enumerate(op_node.args): |
| arg_repset = inputs_repsets[i] |
| |
| # Use ALL_STORAGES_REPSET 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(ALL_STORAGES_REPSET) |
| # 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 |
| # ALL_STORAGES_REPSET in this case. |
| elif arg_repset.is_empty(): |
| args_repset_list.append(ALL_STORAGES_REPSET) |
| 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 = ALL_STORAGES_REPSET |
| if num_tensors_in_node(op_node) == 1: |
| common_out_repset = filter_invalid_reprs( |
| op_node.meta["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 |
| |
| if self.sync_primary_io_repr: |
| if not self.get_out_repset(0).has_compatible_packed_dim_info_set( |
| source_repset |
| ): |
| return False |
| |
| # If this point is reached, then it is possible to constrain |
| narrowed = arg_current_repset.make_intersect(source_repset) |
| 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 |
