[ET-VK][AOT][ez] Introduce vulkan export utils lib (pytorch#6605)

Pull Request resolved: pytorch#6600

## Changes

As title. Introduce a common Python utility library for scripts in the Vulkan backend.
ghstack-source-id: 251223077

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

---------

Co-authored-by: Stephen Jia <[email protected]>

Author: pytorchbot

backends/vulkan/_passes/TARGETS

@@ -12,6 +12,7 @@ runtime.python_library(
     deps = [
         "//caffe2:torch",
         "//executorch/exir:pass_base",
+        "//executorch/backends/vulkan:utils_lib",
     ],
 )
 

backends/vulkan/_passes/insert_prepack_nodes.py

@@ -13,10 +13,10 @@
 import torch
 
 from executorch.backends.vulkan.op_registry import handles_own_prepacking
+from executorch.backends.vulkan.utils import is_param_node
 
 from executorch.exir.dialects._ops import ops as exir_ops
 
-from torch._export.utils import is_buffer, is_param
 from torch.export import ExportedProgram
 
 
@@ -31,25 +31,8 @@ def insert_prepack_nodes(program: ExportedProgram) -> ExportedProgram:
     argument into the operator implementation.
     """
 
-    def is_get_attr_node(node: torch.fx.Node) -> bool:
-        return isinstance(node, torch.fx.Node) and node.op == "get_attr"
-
-    def is_constant(node: torch.fx.Node) -> bool:
-        return node.name in program.graph_signature.inputs_to_lifted_tensor_constants
-
-    def is_param_node(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_constant(node)
-        )
-
     def prepack_not_required(node: torch.fx.Node) -> bool:
-        if not is_param_node(node):
+        if not is_param_node(program, node):
             return True
 
         for user in node.users:

backends/vulkan/serialization/vulkan_graph_builder.py

@@ -12,6 +12,11 @@
 import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema
 
 import torch
+from executorch.backends.vulkan.utils import (
+    is_constant,
+    is_get_attr_node,
+    is_param_node,
+)
 from executorch.exir.backend.utils import DelegateMappingBuilder
 
 from executorch.exir.tensor import TensorSpec
@@ -68,34 +73,12 @@ def get_vk_datatype(torch_dtype: torch.dtype) -> vk_graph_schema.VkDataType:
         else:
             raise AssertionError(f"Invalid dtype for vulkan_preprocess ({torch_dtype})")
 
-    def is_constant(self, node: Node):
-        return (
-            node.name in self.program.graph_signature.inputs_to_lifted_tensor_constants
-        )
-
-    def is_get_attr_node(self, node: Node) -> bool:
-        """
-        Returns true if the given node is a get attr node for a tensor of the model
-        """
-        return isinstance(node, Node) and node.op == "get_attr"
-
-    def is_param_node(self, node: Node) -> bool:
-        """
-        Check if the given node is a parameter within the exported program
-        """
-        return (
-            self.is_get_attr_node(node)
-            or is_param(self.program, node)
-            or is_buffer(self.program, node)
-            or self.is_constant(node)
-        )
-
     def get_constant(self, node: Node) -> Optional[torch.Tensor]:
         """
         Returns the constant associated with the given node in the exported program.
         Returns None if the node is not a constant within the exported program
         """
-        if self.is_constant(node):
+        if is_constant(self.program, node):
             constant_name = (
                 self.program.graph_signature.inputs_to_lifted_tensor_constants[
                     node.name
@@ -116,9 +99,9 @@ def get_param_tensor(self, node: Node) -> torch.Tensor:
             tensor = get_param(self.program, node)
         elif is_buffer(self.program, node):
             tensor = get_buffer(self.program, node)
-        elif self.is_constant(node):
+        elif is_constant(self.program, node):
             tensor = self.get_constant(node)
-        elif self.is_get_attr_node(node):
+        elif is_get_attr_node(node):
             # This is a hack to support both lifted and unlifted graph
             try:
                 tensor = getattr(node.graph.owning_module, node.target)
@@ -132,7 +115,7 @@ def get_param_tensor(self, node: Node) -> torch.Tensor:
 
     def maybe_add_constant_tensor(self, node: Node) -> int:
         constant_id = -1
-        if self.is_param_node(node):
+        if is_param_node(self.program, node):
             constant_id = len(self.const_tensors)
             self.const_tensors.append(self.get_param_tensor(node))
 
@@ -280,7 +263,7 @@ def process_placeholder_node(self, node: Node) -> None:
         if len(node.users) == 0:
             return None
         ids = self.create_node_value(node)
-        if not self.is_param_node(node):
+        if not is_param_node(self.program, node):
             if isinstance(ids, int):
                 self.input_ids.append(ids)
             else:

backends/vulkan/targets.bzl

@@ -213,6 +213,19 @@ def define_common_targets(is_fbcode = False):
     ## AOT targets
     ##
     if is_fbcode:
+        runtime.python_library(
+            name = "utils_lib",
+            srcs = [
+                "utils.py",
+            ],
+            visibility = [
+                "//executorch/backends/vulkan/...",
+            ],
+            deps = [
+                "//caffe2:torch",
+            ]
+        )
+
         runtime.python_library(
             name = "custom_ops_lib",
             srcs = [

backends/vulkan/utils.py

@@ -0,0 +1,30 @@
+# 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 torch
+from torch._export.utils import is_buffer, is_param
+
+from torch.export import ExportedProgram
+
+
+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_constant(program, node)
+    )