Update on "[ET-VK][ez] Clean up organization of supported_ops"

As title. Group supported ops by features instead of op category. This will make it easier to mark that an op has increased its feature set. This also allows the registration code to be simplified a lot.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/partitioner/supported_ops.py

@@ -84,8 +84,6 @@ def __contains__(self, op):
     # Convolution ops
     exir_ops.edge.aten.convolution.default,
     exir_ops.edge.et_vk.conv_with_clamp.default,
-    # Custom ops
-    "llama::sdpa_with_kv_cache",
 ]
 
 NO_DYNAMIC_SHAPE = [

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -119,19 +119,13 @@ def is_node_supported(
     def _is_node_supported(
         self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node
     ) -> bool:
-        target = node.target
-        if node.target == torch.ops.higher_order.auto_functionalized:
-            first_arg = node.args[0]
-            assert isinstance(first_arg, torch._ops.OpOverload)
-            target = first_arg.name()
-
         if self.is_linear_permute(node):
             return True
 
-        if target not in VulkanSupportedOperators._ops:
+        if node.target not in VulkanSupportedOperators._ops:
             return False
 
-        features = VulkanSupportedOperators._ops[target]
+        features = VulkanSupportedOperators._ops[node.target]
 
         if self.require_dynamic_shapes and not features.supports_dynamic_shape:
             return False

backends/vulkan/vulkan_preprocess.py

@@ -36,10 +36,6 @@
 
 from executorch.exir.program._program import _copy_module
 
-from torch.export._remove_auto_functionalized_pass import (
-    unsafe_remove_auto_functionalized_pass,
-)
-
 DEFAULT_DEBUG_HANDLE = 65535
 
 
@@ -52,8 +48,6 @@ def preprocess(  # noqa: C901
         program: ExportedProgram,
         module_compile_spec: List[CompileSpec],
     ) -> PreprocessResult:
-        program = unsafe_remove_auto_functionalized_pass(program)
-
         passes = [
             RemoveCloneOpsTransform(),
             AddmmToLinearTransform(),