Update base for Update on "[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance."

Adding push constant and ubo verison of select and slice ops to improve memory and performance.
*   Updated `transfer_buffer.yaml` and `transfer_texture.yaml` to include `UBO_PARAMS` parameter and generate variants for `select` and `slice` ops with UBO parameters.
*   Updated `transfer.glsl` to generate ubo and push constant versions of `select` and `slice` ops with UBO parameters.

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

[ghstack-poisoned]

Author: trivedivivek

backends/xnnpack/_passes/__init__.py

@@ -23,6 +23,9 @@
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
+from executorch.backends.xnnpack._passes.remove_redundant_copy_pass import (
+    RemoveRedundantCopyPass,
+)
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 
 from executorch.exir.pass_base import ExportPass
@@ -65,6 +68,7 @@ def __init__(
                 Conv1dUnsqueezePass,
                 PReLUReshapePass,
                 ChannelsLastTaggedReshapePass,
+                RemoveRedundantCopyPass,
             ]
         else:
             self.passes = passes

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from enum import Enum
 from typing import Optional, Tuple
 
 import torch
@@ -19,6 +20,11 @@
 from executorch.exir.pass_base import PassResult
 
 
+class InputDimOrder(Enum):
+    NCHW = 1
+    NHWC = 2
+
+
 # TODO(T151254305) use subgraph_rewriter
 class ChannelsLastTaggedReshapePass(XNNPACKPass):
     """
@@ -83,17 +89,49 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass):
     # is done
     PARTNER_NODE = "XNN_CHANNELS_LAST_TAGGED_RESHAPE_PARTNER_NODE"
 
-    def mark_as_nhwc_node(self, node: torch.fx.Node) -> None:
+    @staticmethod
+    def mark_as_nhwc_node(node: torch.fx.Node) -> None:
         node.meta[ChannelsLastTaggedReshapePass.XNN_NHWC_NODE] = True
 
-    def mark_as_nchw_node(self, node: torch.fx.Node) -> None:
+    @staticmethod
+    def mark_as_nchw_node(node: torch.fx.Node) -> None:
         node.meta[ChannelsLastTaggedReshapePass.XNN_NHWC_NODE] = False
 
-    def is_nhwc_node(self, node: torch.fx.Node) -> bool:
+    def tag_node(self, node: torch.fx.Node) -> None:
+        if node.kwargs["memory_format"] == torch.channels_last:
+            self.mark_as_nhwc_node(node)
+        else:
+            self.mark_as_nchw_node(node)
+
+    @staticmethod
+    def is_nhwc_node(node: torch.fx.Node) -> bool:
+        if is_dequant(node) and len(node.all_input_nodes) > 0:
+            quantize_node = node.args[0]
+            if len(quantize_node.all_input_nodes) > 0:
+                actual_node = quantize_node.args[0]
+                if actual_node.op == "placeholder":
+                    return not actual_node.meta["val"][0].is_contiguous()
+                else:
+                    return actual_node.meta.get(
+                        ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False
+                    )
+
         return node.meta.get(ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False)
 
-    def is_nchw_node(self, node: torch.fx.Node) -> bool:
-        return not self.is_nhwc_node(node)
+    @staticmethod
+    def is_nchw_node(node: torch.fx.Node) -> bool:
+        if is_dequant(node) and len(node.all_input_nodes) > 0:
+            quantize_node = node.args[0]
+            if len(quantize_node.all_input_nodes) > 0:
+                actual_node = quantize_node.args[0]
+                if actual_node.op == "placeholder":
+                    return actual_node.meta["val"][0].is_contiguous()
+                else:
+                    return not actual_node.meta.get(
+                        ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False
+                    )
+
+        return not ChannelsLastTaggedReshapePass.is_nhwc_node(node)
 
     def requires_nhwc_input(self, node: torch.fx.Node) -> bool:
         return node.target in self.memory_sensitive_ops_nhwc
@@ -111,7 +149,7 @@ def can_be_converted_to_nhwc(self, node: torch.fx.Node) -> bool:
         is_nchw_constant = (
             is_param_node(self.exported_program, node)
             and (ChannelsLastTaggedReshapePass.XNN_NHWC_NODE in node.meta)
-            and (self.is_nchw_node(node))
+            and (ChannelsLastTaggedReshapePass.is_nchw_node(node))
         )
         return is_4d and not is_nchw_constant
 
@@ -273,6 +311,22 @@ def insert_copy_and_assign_partner_nodes_quantization_sensitive(
             # in that case
             self.make_partners(original_input, copy_node)
 
+    def input_dim_order(
+        self, input_node: torch.fx.Node, input_order: InputDimOrder
+    ) -> bool:
+        if input_node.op == "placeholder":
+            return (
+                input_node.meta["val"].is_contiguous()
+                if input_order == InputDimOrder.NCHW
+                else not input_node.meta["val"].is_contiguous()
+            )
+        else:
+            return (
+                ChannelsLastTaggedReshapePass.is_nchw_node(input_node)
+                if input_order == InputDimOrder.NCHW
+                else ChannelsLastTaggedReshapePass.is_nhwc_node(input_node)
+            )
+
     def input_to_nhwc(
         self,
         graph_module: torch.fx.GraphModule,
@@ -282,7 +336,7 @@ def input_to_nhwc(
         if is_param_node(self.exported_program, input_node):
             if (
                 ChannelsLastTaggedReshapePass.XNN_NHWC_NODE in input_node.meta
-                and self.is_nchw_node(input_node)
+                and ChannelsLastTaggedReshapePass.is_nchw_node(input_node)
             ):
                 # This constant data tensor has been used somewhere else
                 # in NCHW format so we can't use it here in NHWC format
@@ -296,7 +350,10 @@ def input_to_nhwc(
         if input_node.op == "placeholder":
             if not input_node.meta["val"][0].is_contiguous():
                 return
-        elif self.is_nhwc_node(input_node):
+        elif ChannelsLastTaggedReshapePass.is_nhwc_node(input_node):
+            return
+
+        if self.input_dim_order(input_node, InputDimOrder.NHWC):
             return
 
         if not self.can_be_converted_to_nhwc(input_node):
@@ -326,6 +383,8 @@ def input_to_nhwc(
                     args=(input_node,),
                     memory_format=torch.channels_last,
                 )
+                # Use static method for consistency
+                ChannelsLastTaggedReshapePass.mark_as_nhwc_node(input_node_nhwc)
 
             if is_dynamic_input:
                 # Replace downstream input_nodes with NHWC node
@@ -348,7 +407,7 @@ def input_to_nchw(
         if is_param_node(self.exported_program, input_node):
             if (
                 ChannelsLastTaggedReshapePass.XNN_NHWC_NODE in input_node.meta
-                and self.is_nhwc_node(input_node)
+                and ChannelsLastTaggedReshapePass.is_nhwc_node(input_node)
             ):
                 # This constant data tensor has been used somewhere else
                 # in NHWC format so we can't use it here in NCHW format
@@ -363,7 +422,10 @@ def input_to_nchw(
         if input_node.op == "placeholder":
             if input_node.meta["val"].is_contiguous():
                 return
-        elif self.is_nchw_node(input_node):
+        elif ChannelsLastTaggedReshapePass.is_nchw_node(input_node):
+            return
+
+        if self.input_dim_order(input_node, InputDimOrder.NCHW):
             return
 
         if ChannelsLastTaggedReshapePass.PARTNER_NODE in input_node.meta:
@@ -380,6 +442,7 @@ def input_to_nchw(
                     args=(input_node,),
                     memory_format=torch.contiguous_format,
                 )
+                ChannelsLastTaggedReshapePass.mark_as_nchw_node(input_node_nchw)
 
         self.insert_copy_and_assign_partner_nodes_quantization_sensitive(
             graph_module=graph_module,
@@ -393,7 +456,12 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
         original_nodes = list(graph.nodes)
         for node in original_nodes:
             if len(node.all_input_nodes) == 0:
-                # This node has no inputs so we don't need to change anything
+                # This node has no inputs so we don't need to change anything, but still need to tag input nodes
+                if "val" in node.meta and isinstance(node.meta["val"], torch.Tensor):
+                    if node.meta["val"].is_contiguous():
+                        self.mark_as_nchw_node(node)
+                    else:
+                        self.mark_as_nhwc_node(node)
                 continue
 
             # Need special case for output node because it can have multiple output dim orders as we can output a tuple multiple nodes
@@ -407,10 +475,12 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
             elif self.requires_nhwc_input(node):
                 # Nodes which enter this branch are ones that require their
                 # first input to be nhwc. This makes this node's output nhwc too
-
                 self.input_to_nhwc(graph_module, node.args[0], node)
-                for input_node in node.all_input_nodes:
-                    if input_node.op == "placeholder" and self.is_nhwc_node(input_node):
+                for input_node in node.all_input_nodes[1:]:
+                    if (
+                        input_node.op == "placeholder"
+                        and ChannelsLastTaggedReshapePass.is_nhwc_node(input_node)
+                    ):
                         raise AssertionError(
                             f"Expected {input_node} to be NCHW in channels last reshape pass"
                         )
@@ -419,11 +489,14 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                 # The node requires nchw inputs
                 for input_node in node.all_input_nodes:
                     self.input_to_nchw(graph_module, input_node, node)
+            elif node.target == exir_ops.edge.aten._to_copy.default:
+                self.tag_node(node)
             else:
                 # The node can have inputs in any format (but all must be the
                 # same format)
                 is_or_isnt_nhwc_node = [
-                    self.is_nhwc_node(input_node) for input_node in node.all_input_nodes
+                    ChannelsLastTaggedReshapePass.is_nhwc_node(input_node)
+                    for input_node in node.all_input_nodes
                 ]
                 if all(is_or_isnt_nhwc_node):
                     # All inputs are nhwc so this node's output is nhwc too

backends/xnnpack/_passes/remove_redundant_copy_pass.py

@@ -0,0 +1,158 @@
+# 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 executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import (
+    ChannelsLastTaggedReshapePass,
+)
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+
+
+class RemoveRedundantCopyPass(XNNPACKPass):
+    def _safe_remove_node(self, node, graph):
+        if len(node.users) == 0:
+            graph.erase_node(node)
+
+    def _try_remove_regular_redundant_to_copy(self, node, graph):
+        """
+        Try to remove redundant regular to_copy operations with pattern to_copy1 -> to_copy2 with opposite memory formats
+        """
+        input_node = node.args[0]
+
+        # Check if input is a to_copy with opposite memory format
+        if (
+            input_node.target == exir_ops.edge.aten._to_copy.default
+            and ChannelsLastTaggedReshapePass.is_nchw_node(input_node)
+            != ChannelsLastTaggedReshapePass.is_nchw_node(node)
+            and len(input_node.users) == 1
+        ):  # Ensure the first copy has no other users
+
+            # Get the original input (before the first to_copy)
+            original_input = input_node.args[0]
+
+            # Replace all users of the second to_copy with the original input
+            for user in node.users.copy():
+                user.replace_input_with(node, original_input)
+
+            # Remove both to_copy nodes
+            self._safe_remove_node(node, graph)
+            self._safe_remove_node(input_node, graph)
+
+            return True
+        elif (
+            ChannelsLastTaggedReshapePass.is_nhwc_node(input_node)
+            and ChannelsLastTaggedReshapePass.is_nhwc_node(node)
+        ) or (
+            ChannelsLastTaggedReshapePass.is_nchw_node(input_node)
+            and ChannelsLastTaggedReshapePass.is_nchw_node(node)
+        ):
+            # Replace all users of the second to_copy with the original input
+            for user in node.users.copy():
+                user.replace_input_with(node, input_node)
+            self._safe_remove_node(node, graph)
+            return True
+
+        return False
+
+    def _try_remove_quantized_redundant_to_copy(self, node, graph):
+        """
+        Try to remove redundant to_copy operations in quantized graphs with pattern dq1 -> to_copy1 -> q1 -> dq2 -> to_copy2 -> q2
+        """
+        # Check if this to_copy is followed by a quantize node
+        if len(node.users) != 1:
+            return False
+        q_node = next(iter(node.users))
+        if not is_quant(q_node):
+            return False
+
+        # Check if this to_copy is preceded by a dequantize node
+        dq_node = node.args[0]
+        if not is_dequant(dq_node):
+            return False
+
+        # Get the input to the dequantize node
+        if len(dq_node.all_input_nodes) != 1:
+            return False
+
+        prev_q_node = dq_node.args[0]
+
+        # Check if there's another dequantize -> to_copy -> quantize chain
+        if not is_quant(prev_q_node) or len(prev_q_node.all_input_nodes) != 1:
+            return False
+
+        # Check if there's a to_copy before the previous quantize
+        prev_to_copy = prev_q_node.args[0]
+        if (
+            prev_to_copy.target == exir_ops.edge.aten._to_copy.default
+            and ChannelsLastTaggedReshapePass.is_nchw_node(prev_to_copy)
+            != ChannelsLastTaggedReshapePass.is_nchw_node(node)
+            and len(prev_to_copy.users) == 1
+        ):  # Ensure the first copy has no other users
+            prev_dq_node = prev_to_copy.args[0]
+            if not is_dequant(prev_dq_node) or len(prev_dq_node.all_input_nodes) != 1:
+                return False
+
+            # Get the original input (before the first to_copy)
+            original_input = prev_dq_node.args[0]
+
+            # Replace all users of the second to_copy with the original input
+            for user in q_node.users.copy():
+                user.replace_input_with(q_node, original_input)
+
+            # Remove nodes safely (only if they have no other users)
+            self._safe_remove_node(q_node, graph)
+            self._safe_remove_node(node, graph)
+            self._safe_remove_node(dq_node, graph)
+            self._safe_remove_node(prev_q_node, graph)
+            self._safe_remove_node(prev_to_copy, graph)
+            self._safe_remove_node(prev_dq_node, graph)
+        elif (
+            ChannelsLastTaggedReshapePass.is_nhwc_node(prev_to_copy)
+            and ChannelsLastTaggedReshapePass.is_nhwc_node(node)
+        ) or (
+            ChannelsLastTaggedReshapePass.is_nchw_node(prev_to_copy)
+            and ChannelsLastTaggedReshapePass.is_nchw_node(node)
+        ):
+            # Remove node and the q/dq around it only
+            # Get the original quantized tensor (input to dq_node)
+            original_q_tensor = dq_node.args[0]
+
+            # Replace all users of q_node with the original quantized tensor
+            for user in q_node.users.copy():
+                user.replace_input_with(q_node, original_q_tensor)
+
+            self._safe_remove_node(q_node, graph)
+            self._safe_remove_node(node, graph)
+            self._safe_remove_node(dq_node, graph)
+            return True
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        graph = graph_module.graph
+        original_nodes = list(graph.nodes)
+
+        for node in original_nodes:
+            if len(node.all_input_nodes) == 0:
+                continue
+
+            # Only process to_copy nodes
+            if node.target != exir_ops.edge.aten._to_copy.default:
+                continue
+
+            if is_dequant(node.args[0]):
+                self._try_remove_quantized_redundant_to_copy(node, graph)
+            else:
+                self._try_remove_regular_redundant_to_copy(node, graph)
+
+        graph_module.recompile()
+
+        # Since we are overriding "call", we need to call the parent's "call"
+        # to retrace the graph and regenerate metadata
+        graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, True)