Update on "[ET-VK][AOT] Define pass application order"

## Changes

The goal of this diff is to enforce a specific structure in how graph transform passes are applied during `vulkan_preprocess`. This will help make sure that certain passes are applied at the correct time, and that pre-requisite conditions for passes are fulfilled before they are applied.

See the comments in `vulkan_preprocess.py` for more details.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/_passes/insert_prepack_nodes.py

@@ -71,10 +71,10 @@ def prepack_not_required(node: torch.fx.Node) -> bool:
                 exir_ops.edge.et_vk.prepack.default,
                 (node,),
             )
-            # This pass assumes that the SpecPropPass() has already been applied, and
-            # validate that the original node is marked as a constant. Constant tensors
-            # do not participate in memory planning.
+            # This pass assumes that the SpecPropPass() has already been applied
             assert "spec" in node.meta
+            # Validate that the original node is marked as a constant. Constant tensors
+            # do not participate in memory planning.
             assert node.meta["spec"].const
             prepack_node.meta["spec"] = deepcopy(node.meta["spec"])
             # Set the mem_obj_id to -1 to indicate that this node requires a dedicated

backends/vulkan/vulkan_preprocess.py

@@ -90,7 +90,8 @@ def preprocess(  # noqa: C901
         program = unsafe_remove_auto_functionalized_pass(program)
 
         # First, apply passes that fuse/remove operators to consolidate the graph
-        # structure but still preserve an "ATen-compliant" graph structure.
+        # structure but still preserve an "ATen-compliant" graph structure (i.e. all
+        # arguments to aten operators must match the ATen function schema).
         program = apply_passes(
             program,
             [
@@ -104,8 +105,9 @@ def preprocess(  # noqa: C901
         )
 
         # Next annotate tensor nodes with TensorSpec structs which is needed for dynamic
-        # shapes and memory planning. Until this point, the graph must be "ATen compliant"
-        # (e.g. all arguments to aten operators must match the ATen function schema).
+        # shapes and memory planning. Until this point, the graph must be ATen compliant
+        # because SpecPropPass will be calling the underlying ATen operators during its
+        # execution.
         program = apply_passes(program, [SpecPropPass()])
 
         # Apply graph transforms which either require `TensorSpec`s to have been created