[ET-VK] Refine paritioner to account for storage type and memory layout

## Context

There are a variety of ways that tensors can be represented in Vulkan. The two main descriptors for how a tensor is laid out in memory is:

1. Storage Type (buffer or texture)
2. Memory Layout (which dim is packed along a texel, which dim has a stride of 1, etc.)

Due to the differences between buffers and textures, and the differences between different memory layouts, an implementation for an operator may only support a specific set of (storage type, memory layout) combinations.

Furthermore, if an operator implementation supports multiple (storage type, memory layout) combinations, there may be a "preferred" setting which results in optimal performance.

These changes lay the foundation for the implementation of a memory metadata tagging graph transform, which will make sure that all tensors participating in an operator call is has a valid/optimal (storage type, memory layout) setting, and insert transition operators to transfer input tensors to the correct memory settings when necessary.

An additional change that is required arises from the fact that in Vulkan, there is a limit on texture and buffer sizes. Therefore, the partitioner needs to account for the storage types and memory layouts supported by the operator implementation, and check if all tensors participating in a computation can be represented with some storage type, memory layout combination supported by the implementation.


## Changes

Improvements to the operator registry:

* Introduce utility functions to check the optimal and enabled storage types and memory layouts for an operator

Improvements to the Partitioner:

* Account for the storage types and memory layouts supported by an operator when deciding if a node should be partitioned
* Improved logic for fusable ops (i.e. the permute/transpose before a mm which can be fused into linear) to check if the final target op is supported in Vulkan, and only partition those nodes if so. Otherwise, don't partition it so that it can be fused by another backend.

Differential Revision: [D65428843](https://our.internmc.facebook.com/intern/diff/D65428843/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/op_registry.py

@@ -8,44 +8,65 @@
 
 import operator
 
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, Optional, Set, Union
 
 import executorch.backends.vulkan.custom_ops_lib  # noqa
 
 import torch
 
-from executorch.backends.vulkan.serialization.vulkan_graph_schema import VkMemoryLayout
+from executorch.backends.vulkan.serialization.vulkan_graph_schema import (
+    VkMemoryLayout,
+    VkStorageType,
+)
+
+from executorch.backends.vulkan.utils import (
+    all_memory_layouts,
+    all_packed_dims,
+    PackedDim,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from torch._subclasses.fake_tensor import FakeTensor
 
+######################
+## OpFeatures class ##
+######################
+
 
 def allow_node(node: torch.fx.Node) -> bool:
     return True
 
 
 class TextureImplFeatures:
     __slots__ = [
-        # Indicates if the compute shader is agnostic to the packed dimension
-        "uses_packed_dim",
-        # Indicates if the compute shader is agnostic to the texture axis mapping
+        "valid_packed_dims",
         "uses_axis_map",
-        # Specifies a specific set of memory layouts that the shader supports. If it is
-        # and empty list, then the supported memory layouts can be inferred from the
-        # `uses_packed_dim` and `uses_axis_map` flags.
-        "supported_layouts",
     ]
 
     def __init__(
         self,
-        uses_packed_dim: bool = False,
         uses_axis_map: bool = False,
-        supported_layouts: Optional[List[VkMemoryLayout]] = None,
+        valid_packed_dims: Optional[Set[PackedDim]] = None,
     ):
-        self.uses_packed_dim: bool = uses_packed_dim
         self.uses_axis_map: bool = uses_axis_map
-        self.supported_layouts: Optional[List[VkMemoryLayout]] = supported_layouts
+        self.valid_packed_dims = set()
+        if valid_packed_dims is not None:
+            self.valid_packed_dims = valid_packed_dims
+
+    def valid_memory_layouts(self) -> Set[VkMemoryLayout]:
+        layouts = set()
+
+        if PackedDim.WIDTH in self.valid_packed_dims:
+            layouts.add(VkMemoryLayout.TENSOR_WIDTH_PACKED)
+
+        if PackedDim.HEIGHT in self.valid_packed_dims:
+            layouts.add(VkMemoryLayout.TENSOR_HEIGHT_PACKED)
+
+        if PackedDim.CHANNELS in self.valid_packed_dims:
+            layouts.add(VkMemoryLayout.TENSOR_CHANNELS_PACKED)
+
+        return layouts
 
 
 class OpFeatures:
@@ -58,6 +79,9 @@ class OpFeatures:
         # bool indicating if the operator has a resize function, which allows it to
         # support dynamic shape tensors.
         "resize_fn",
+        # Optimal
+        "optimal_storage",
+        "optimal_layout",
         # bool indicating if the operator handles its own prepacking. If this is True,
         # then the insert_prepack_nodes pass will not insert prepack nodes for the args
         # of the op.
@@ -72,17 +96,64 @@ def __init__(
         texture_impl: Optional[TextureImplFeatures] = None,
         buffer_impl: bool = False,
         resize_fn: bool = False,
+        optimal_storage: Optional[VkStorageType] = None,
+        optimal_layout: Optional[VkMemoryLayout] = None,
         handles_own_prepacking: bool = False,
         check_node_fn: Optional[Callable] = None,
     ):
         self.texture_impl: Optional[TextureImplFeatures] = texture_impl
         self.buffer_impl: bool = buffer_impl
         self.resize_fn: bool = resize_fn
+        self.optimal_storage: Optional[VkStorageType] = optimal_storage
+        self.optimal_layout: Optional[VkMemoryLayout] = optimal_layout
         self.handles_own_prepacking: bool = handles_own_prepacking
         self.check_node_fn: Callable = allow_node
         if check_node_fn is not None:
             self.check_node_fn = check_node_fn
 
+    def propose_storage_type(self) -> Optional[VkStorageType]:
+        if self.optimal_storage is not None:
+            return self.optimal_storage
+
+        if self.texture_impl is not None and not self.buffer_impl:
+            return VkStorageType.TEXTURE_3D
+        elif self.buffer_impl and self.texture_impl is None:
+            return VkStorageType.BUFFER
+
+        return None
+
+    def supported_storage_types(self) -> Set[VkStorageType]:
+        storage_types = set()
+        if self.texture_impl is not None:
+            storage_types.add(VkStorageType.TEXTURE_3D)
+        if self.buffer_impl:
+            storage_types.add(VkStorageType.BUFFER)
+
+        return storage_types
+
+    def propose_memory_layout(self, storage: VkStorageType) -> Optional[VkMemoryLayout]:
+        if self.optimal_layout is not None:
+            return self.optimal_layout
+
+        if storage == VkStorageType.TEXTURE_3D:
+            assert self.texture_impl is not None
+            possible_layouts = self.texture_impl.valid_memory_layouts()
+            if len(possible_layouts) == 1:
+                return next(iter(possible_layouts))
+
+        return None
+
+    def supported_memory_layouts(self, storage: VkStorageType) -> Set[VkMemoryLayout]:
+        if storage == VkStorageType.TEXTURE_3D:
+            assert self.texture_impl is not None
+            return self.texture_impl.valid_memory_layouts()
+        else:
+            return all_memory_layouts
+
+
+#######################
+## Operator Registry ##
+#######################
 
 OpKey = Union[str, torch._ops.OpOverload, EdgeOpOverload]
 
@@ -122,8 +193,8 @@ def update_features_impl(op: OpKey):
 )
 def register_ephemeral_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
         uses_axis_map=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.buffer_impl = True
     features.resize_fn = True
@@ -143,8 +214,8 @@ def register_ephemeral_op(features: OpFeatures):
 )
 def register_binary_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
         uses_axis_map=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.resize_fn = True
     return features
@@ -170,8 +241,8 @@ def register_binary_op(features: OpFeatures):
 )
 def register_unary_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
         uses_axis_map=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.buffer_impl = True
     features.resize_fn = True
@@ -181,8 +252,8 @@ def register_unary_op(features: OpFeatures):
 @update_features(exir_ops.edge.aten._to_copy.default)
 def register_to_copy_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
         uses_axis_map=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.resize_fn = True
 
@@ -220,40 +291,43 @@ def check_to_copy_node(node: torch.fx.Node) -> bool:
 )
 def register_mm_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=False,
         uses_axis_map=True,
-        supported_layouts=[
-            VkMemoryLayout.TENSOR_WIDTH_PACKED,
-            VkMemoryLayout.TENSOR_CHANNELS_PACKED,
-        ],
+        valid_packed_dims={
+            PackedDim.WIDTH,
+            PackedDim.CHANNELS,
+        },
     )
     features.buffer_impl = True
     features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
     features.handles_own_prepacking = True
     return features
 
 
 @update_features(exir_ops.edge.aten._weight_int8pack_mm.default)
 def register_int8_mm_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=False,
         uses_axis_map=False,
-        supported_layouts=[VkMemoryLayout.TENSOR_WIDTH_PACKED],
+        valid_packed_dims={PackedDim.WIDTH},
     )
     features.buffer_impl = True
     features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
     features.handles_own_prepacking = True
     return features
 
 
 @update_features(exir_ops.edge.et_vk.linear_weight_int4.default)
 def register_int4_mm_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=False,
         uses_axis_map=False,
-        supported_layouts=[VkMemoryLayout.TENSOR_WIDTH_PACKED],
+        valid_packed_dims={PackedDim.WIDTH},
     )
     features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
     features.handles_own_prepacking = True
     return features
 
@@ -266,7 +340,7 @@ def register_int4_mm_op(features: OpFeatures):
 )
 def register_softmax_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.resize_fn = True
     return features
@@ -282,7 +356,7 @@ def register_softmax_op(features: OpFeatures):
 )
 def register_reduce_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.resize_fn = True
 
@@ -309,7 +383,7 @@ def check_reduce_node(node: torch.fx.Node) -> bool:
 )
 def register_2d_pool_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        valid_packed_dims={PackedDim.CHANNELS},
     )
     features.resize_fn = True
     return features
@@ -323,27 +397,31 @@ def register_2d_pool_op(features: OpFeatures):
 )
 def register_convolution_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        valid_packed_dims={PackedDim.CHANNELS},
     )
     features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_CHANNELS_PACKED
     features.handles_own_prepacking = True
     return features
 
 
 @update_features("llama::sdpa_with_kv_cache")
 def register_sdpa_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_WIDTH_PACKED],
+        valid_packed_dims={PackedDim.WIDTH},
     )
     features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
     features.handles_own_prepacking = True
     return features
 
 
 @update_features(exir_ops.edge.et_vk.apply_rotary_emb.default)
 def register_rotary_emb_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_WIDTH_PACKED],
+        valid_packed_dims={PackedDim.WIDTH},
     )
     features.resize_fn = True
     return features
@@ -352,7 +430,7 @@ def register_rotary_emb_op(features: OpFeatures):
 @update_features(exir_ops.edge.aten.view_copy.default)
 def register_view_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        uses_packed_dim=True,
+        valid_packed_dims=all_packed_dims,
     )
     features.resize_fn = True
     return features
@@ -393,7 +471,7 @@ def register_view_op(features: OpFeatures):
 )
 def register_ported_op(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        valid_packed_dims={PackedDim.CHANNELS},
     )
     return features
 
@@ -408,15 +486,24 @@ def register_ported_op(features: OpFeatures):
 )
 def register_ported_ops_with_prepacking(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
-        supported_layouts=[VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        valid_packed_dims={PackedDim.CHANNELS},
     )
     features.handles_own_prepacking = True
     return features
 
 
-##
-## Utility Functions
-##
+#######################
+## Utility functions ##
+#######################
+
+
+def has_impl(target: OpKey) -> bool:
+    if not isinstance(target, str):
+        if target not in vulkan_supported_ops:
+            return target.name() in vulkan_supported_ops
+        return target in vulkan_supported_ops
+    else:
+        return target in vulkan_supported_ops
 
 
 def get_op_features(target: OpKey) -> OpFeatures:
@@ -430,5 +517,13 @@ def get_op_features(target: OpKey) -> OpFeatures:
         return vulkan_supported_ops[target]
 
 
+def optimal_storage_type(target: OpKey) -> Optional[VkStorageType]:
+    return get_op_features(target).optimal_storage
+
+
+def optimal_memory_layout(target: OpKey) -> Optional[VkMemoryLayout]:
+    return get_op_features(target).optimal_layout
+
+
 def handles_own_prepacking(target: OpKey) -> bool:
     return get_op_features(target).handles_own_prepacking

backends/vulkan/partitioner/TARGETS

@@ -13,6 +13,7 @@ runtime.python_library(
     ],
     deps = [
         "//executorch/backends/vulkan:op_registry",
+        "//executorch/backends/vulkan:utils_lib",
         "//executorch/backends/vulkan:vulkan_preprocess",
         "//executorch/exir:delegate",
         "//executorch/exir:lib",