update SqueezeInt4LinearInputs to process relu/gelu inputs too

Summary: Update/rename SqueezeInt4LinearInputs pass so it wraps gelu/relu with squeeze/unsqueeze view ops too

Differential Revision: D69673068

Author: nathanaelsee

Diff for: backends/transforms/fuse_view_copy.py

@@ -39,8 +39,23 @@ def merge_view_copy_chains(graph: torch.fx.Graph) -> torch.fx.Graph:
     graph.eliminate_dead_code()
     return graph
 
+def remove_noop_view_copy(graph: torch.fx.Graph) -> torch.fx.Graph:
+    """
+    Remove view_copy nodes that are no-ops.
+    """
+    ops = exir_ops.edge
+    view_op = ops.aten.view_copy.default
+    for node in graph.nodes:
+        if node.op == "call_function" and node.target == view_op:
+            input_shape = list(node.args[0].meta["val"].shape)
+            target_shape = node.args[1]
+            if input_shape == target_shape:
+                node.replace_all_uses_with(node.args[0])
+    graph.eliminate_dead_code()
+    return graph
 
 class FuseViewCopyTransform(ExportPass):
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph_module.graph = merge_view_copy_chains(graph_module.graph)
+        graph_module.graph = remove_noop_view_copy(graph_module.graph)
         return PassResult(graph_module, True)

Diff for: backends/vulkan/_passes/TARGETS

@@ -31,9 +31,9 @@ runtime.python_library(
 )
 
 runtime.python_library(
-    name = "squeeze_int4_linear_inputs",
+    name = "squeeze_unsqueeze_inputs",
     srcs = [
-        "squeeze_int4_linear_inputs.py",
+        "squeeze_unsqueeze_inputs.py",
     ],
     visibility = [
         "//executorch/backends/...",
@@ -114,7 +114,7 @@ runtime.python_library(
         ":remove_asserts",
         ":remove_local_scalar_dense",
         ":remove_redundant_ops",
-        ":squeeze_int4_linear_inputs",
+        ":squeeze_unsqueeze_inputs",
         ":tag_memory_meta_pass",
     ]
 )

Diff for: backends/vulkan/_passes/__init__.py

@@ -20,8 +20,8 @@
 from executorch.backends.vulkan._passes.remove_redundant_ops import (
     RemoveRedundantOpsTransform,
 )
-from executorch.backends.vulkan._passes.squeeze_int4_linear_inputs import (
-    SqueezeInt4LinearInputs,
+from executorch.backends.vulkan._passes.squeeze_unsqueeze_inputs import (
+    SqueezeUnsqueezeInputs,
 )
 from executorch.backends.vulkan._passes.tag_memory_meta_pass import TagMemoryMetaPass
 
@@ -32,6 +32,6 @@
     "RemoveAssertsTransform",
     "RemoveLocalScalarDenseOpsTransform",
     "RemoveRedundantOpsTransform",
-    "SqueezeInt4LinearInputs",
+    "SqueezeUnsqueezeInputs",
     "TagMemoryMetaPass",
 ]

Diff for: backends/vulkan/_passes/squeeze_int4_linear_inputs.py renamed to backends/vulkan/_passes/squeeze_unsqueeze_inputs.py

@@ -14,8 +14,13 @@
 
 from torch.fx.node import Argument
 
+class SqueezeUnsqueezeInputs(ExportPass):
+    _squeezable_ops = [
+        exir_ops.edge.et_vk.linear_weight_int4.default,
+        exir_ops.edge.aten.relu.default,
+        exir_ops.edge.aten.gelu.default,
+    ]
 
-class SqueezeInt4LinearInputs(ExportPass):
     def call_operator(
         self,
         op,  # pyre-ignore
@@ -26,7 +31,7 @@ def call_operator(
         def _squeezable(shape: List[int]) -> bool:
             return len(shape) > 2 and 1 in shape
 
-        if op != exir_ops.edge.et_vk.linear_weight_int4.default:
+        if op not in self._squeezable_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         # pyre-ignore[16]: `None` has no attribute `node`

Diff for: backends/vulkan/vulkan_preprocess.py

@@ -26,7 +26,7 @@
     insert_prepack_nodes,
     RemoveLocalScalarDenseOpsTransform,
     RemoveRedundantOpsTransform,
-    SqueezeInt4LinearInputs,
+    SqueezeUnsqueezeInputs,
     TagMemoryMetaPass,
 )
 
@@ -153,7 +153,7 @@ def preprocess(  # noqa: C901
                 RemoveRedundantOpsTransform(),
                 AddmmToLinearTransform(),
                 FuseDequantLinearPass(),
-                SqueezeInt4LinearInputs(),
+                SqueezeUnsqueezeInputs(),
                 FuseViewCopyTransform(),
                 ViewCopyToSqueezeUnsqueezePass(),
                 FuseBatchNormWithConvPass(program),