Update on "[ET-VK] Add pass to remove local_scalar_dense"

## Context

Scalar tensors (i.e. tensors with only 1 element) are often passed in to functions as scalars via

```
scalar_tensor[0].item()
```

This translates to the following chain in the graph

```
index_select = index_select(scalar_tensor, ...)
scalar = local_scalar_dense(index_select)
```

This diff introduces a pass to remove the `local_scalar_dense` "chain" in favor of passing in the input tensor directly.

Note that this replacement only occurs if the original tensor is a scalar tensor.

In the Vulkan backend, these scalar tensors will be represented as symbolic integers instead of actual tensors, which is why this replacement is valid. However, it may not a valid replacement for other backends.

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

[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

@@ -143,22 +143,16 @@ 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 self.is_in_local_scalar_dense_chain(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

@@ -40,10 +40,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
 
 
@@ -56,8 +52,6 @@ def preprocess(  # noqa: C901
         program: ExportedProgram,
         module_compile_spec: List[CompileSpec],
     ) -> PreprocessResult:
-        program = unsafe_remove_auto_functionalized_pass(program)
-
         passes = [
             RemoveCloneOpsTransform(),
             AddmmToLinearTransform(),