Update on "[ET-VK] Implementation of to_dim_order_copy"

Title says it all!

Previously, to_dim_order_copy was handled by removing the op. However, this is not possible if the op is modifying the dtype of the original tensor, so these instances of the op would be skipped by the partitioner.

This diff adds an implementation dtype conversion, which allows to_dim_order_copy to be lowered.

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

[ghstack-poisoned]

Author: ssjia

backends/vulkan/patterns/sdpa.py

@@ -6,6 +6,8 @@
 
 from typing import Any, Optional
 
+import executorch.backends.vulkan.utils as utils
+
 import torch
 
 from executorch.backends.vulkan.patterns.pattern_registry import (
@@ -15,31 +17,18 @@
 )
 
 from executorch.exir import ExportedProgram
-from executorch.exir.dialects._ops import ops as exir_ops
 
 
 def is_update_cache_node(node: Any) -> bool:
-    if not hasattr(node, "target"):
-        return False
+    return utils.node_has_target(node, "llama::update_cache")
 
-    if isinstance(node.target, str):
-        return node.target == "llama::update_cache"
-    elif hasattr(node.target, "name"):
-        return node.target.name() == "llama::update_cache"
-    else:
-        return False
 
+def is_custom_sdpa_node(node: Any) -> bool:
+    return utils.node_has_target(node, "llama::custom_sdpa")
 
-def is_sdpa_with_kv_cache_node(node: Any) -> bool:
-    if not hasattr(node, "target"):
-        return False
 
-    if isinstance(node.target, str):
-        return "sdpa_with_kv_cache" in node.target
-    elif hasattr(node.target, "name"):
-        return "sdpa_with_kv_cache" in node.target.name()
-    else:
-        return False
+def is_sdpa_with_kv_cache_node(node: Any) -> bool:
+    return utils.node_has_target(node, "llama::sdpa_with_kv_cache")
 
 
 class CausalSDPAMatch(PatternMatch):
@@ -97,7 +86,7 @@ def __init__(self, custom_sdpa_node: torch.fx.Node) -> None:
 def find_causal_sdpa_patterns(
     node: torch.fx.Node,
 ) -> Optional[CausalSDPAMatch]:
-    if node.target != exir_ops.edge.llama.custom_sdpa.default:
+    if not is_custom_sdpa_node(node):
         return None
 
     matched_pattern = CausalSDPAMatch(node)

backends/vulkan/utils.py

@@ -373,6 +373,18 @@ def find_quant_user(node: torch.fx.Node) -> Optional[torch.fx.Node]:
     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
 ##