Add view_copy/static_reshape support to XNNPACK delegate (#7959)

Summary:
Add support for delegating view_copy in the XNNPACK delegate via the XNN static_reshape operator. This includes support for up to one dynamic dimension. It also includes conditional support for NHWC, so long as batch and channel are fixed (if batch or channel is modified, the reshape has to be done in the native dim order).


Test Plan: I've added test coverage for view_copy, including dynamic shape support and practitioner constraints, to backends/xnnpack/test/ops/test_view_copy.py.If also added additional tests to cover the new view logic in channels_last_tagged_reshape_pass.

Differential Revision: D68691788

Pulled By: GregoryComer

Author: GregoryComer

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -15,7 +15,7 @@
     is_tagged_as_implicit_q_dq,
     tag_as_implicit_q_dq,
 )
-from executorch.backends.xnnpack.utils.utils import is_param_node
+from executorch.backends.xnnpack.utils.utils import get_input_node, is_param_node
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import PassResult
 
@@ -89,6 +89,22 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass):
     # is done
     PARTNER_NODE = "XNN_CHANNELS_LAST_TAGGED_RESHAPE_PARTNER_NODE"
 
+    @staticmethod
+    def is_view_dim_order_invariant(node: torch.fx.Node) -> bool:
+        # View must be done in NCHW dim order if channel or batch is changed,
+        # or if rank is not 4.
+        in_shape = get_input_node(node, 0).meta["val"].shape
+        out_shape = node.meta["val"].shape
+
+        if len(in_shape) != 4 or len(out_shape) != 4:
+            return False
+
+        # Are batch and channel modified? If so, return false.
+        if in_shape[0] != out_shape[0] or in_shape[1] != out_shape[1]:
+            return False
+
+        return True
+
     @staticmethod
     def mark_as_nhwc_node(node: torch.fx.Node) -> None:
         node.meta[ChannelsLastTaggedReshapePass.XNN_NHWC_NODE] = True
@@ -161,6 +177,13 @@ def requires_nhwc_input(self, node: torch.fx.Node) -> bool:
         return node.target in self.memory_sensitive_ops_nhwc
 
     def requires_nchw_inputs(self, node: torch.fx.Node) -> bool:
+        # Views depend on whether batch or channel are modified.
+        if (
+            node.target == exir_ops.edge.aten.view_copy.default
+            and not self.is_view_dim_order_invariant(node)
+        ):
+            return True
+
         return node.target in self.memory_sensitive_ops_nchw
 
     def can_be_converted_to_nhwc(self, node: torch.fx.Node) -> bool:

backends/xnnpack/operators/__init__.py

@@ -53,4 +53,5 @@
     op_sub,
     op_tanh,
     op_to_copy,
+    op_view_copy,
 )

backends/xnnpack/operators/op_skip_ops.py

@@ -59,16 +59,6 @@ class OpTCopyDefault(OpSkipOps):
     target = "aten.t_copy.default"
 
 
-@register_node_visitor
-class OpViewCopyDefault(OpSkipOps):
-    """
-    currently, do nothing if node is view_copy.default
-    need to handle this later on, currently view it as one of skip ops
-    """
-
-    target = "aten.view_copy.default"
-
-
 @register_node_visitor
 class OpSymSizeInt(OpSkipOps):
     """

backends/xnnpack/operators/op_view_copy.py

@@ -0,0 +1,96 @@
+# 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.
+
+# pyre-unsafe
+
+from typing import Dict
+
+import torch
+from executorch.backends.xnnpack.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNGraph,
+    XNNStaticReshape,
+    XNode,
+)
+from executorch.backends.xnnpack.utils.utils import (
+    check_or_raise,
+    get_input_node,
+    PERM_NCHW_TO_NHWC,
+)
+
+
+@register_node_visitor
+class ViewCopyVisitor(NodeVisitor):
+    target = "aten.view_copy.default"
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        xnn_graph: XNNGraph,
+        vals_to_ids: Dict[torch.fx.Node, int],
+        debug_handle: int,
+    ) -> None:
+        self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids)
+
+        input_node = get_input_node(node, 0)
+
+        # input
+        input_id = vals_to_ids[input_node]
+
+        # output
+        output_id = vals_to_ids[node]
+
+        # input shape
+        check_or_raise(
+            "val" in input_node.meta,
+            "Missing val in tensor metadata for input when serializing XNNStaticReshape",
+        )
+
+        # output shape
+        check_or_raise(
+            "val" in node.meta,
+            "Missing val in tensor metadata for input when serializing XNNStaticReshape",
+        )
+
+        new_shape = node.args[1]
+        check_or_raise(
+            all(isinstance(d, int) for d in new_shape),
+            "Symbolic reshape parameter is not supported in XNNStaticReshape",
+        )
+
+        # PyTorch uses -1 for inferred dims, whereas XNNPACK expects 0.
+        new_shape = tuple(d if d != -1 else 0 for d in new_shape)
+
+        # Handle NCHW dim order - if this op is in NCHW order, we need to permute the
+        # view shape correspondingly.
+        if "XNN_NHWC_NODE" in node.meta:
+            check_or_raise(len(new_shape) == 4, "Invalid NCHW shape")
+            new_shape = [new_shape[PERM_NCHW_TO_NHWC[n]] for n in range(4)]
+
+        num_dynamic_dims = sum(1 for d in new_shape if d == 0)
+
+        check_or_raise(
+            num_dynamic_dims <= 1,
+            "XNNPACK reshape only supports 1 dynamic dimension.",
+        )
+
+        ser_node = XNode(
+            xnode_union=XNNStaticReshape(
+                num_dims=len(new_shape),
+                new_shape=new_shape,
+                input_id=input_id,
+                output_id=output_id,
+                flags=0,
+            ),
+            debug_handle=debug_handle,
+        )
+        xnn_graph.xnodes.append(ser_node)

backends/xnnpack/partition/config/__init__.py

@@ -53,6 +53,7 @@
     TanhConfig,
     ToDimOrderCopyConfig,
     UpsampleBilinear2dConfig,
+    ViewCopyConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
     BatchNormConfig,
@@ -112,6 +113,7 @@
     SquareRootConfig,
     SubConfig,
     UpsampleBilinear2dConfig,
+    ViewCopyConfig,
     # Quant/Dequant Op Configs
     QuantizedPerTensorConfig,
     DeQuantizedPerTensorConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -379,6 +379,31 @@ class ExpConfig(GenericNodePartitionerConfig):
     def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32]
 
+class ViewCopyConfig(GenericNodePartitionerConfig):
+    target_name = "view_copy.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        """
+        XNNPACK's static_reshape only supports 1 dynamic dimension
+        """
+        if not self.check_common_constraints(node, ep):
+            return False
+
+        new_shape = node.args[1]
+        if not all(isinstance(n, int) for n in new_shape):
+            why(node, reason="symbolic reshape is not supported")
+            return False
+
+        dynamic_dim_count = sum(1 for d in new_shape if d == -1)
+        if dynamic_dim_count > 1:
+            why(node, reason="only a single dynamic dimension is supported")
+            return False
+
+        return True
+
 
 class FloorConfig(GenericNodePartitionerConfig):
     target_name = "floor.default"