Add view_copy/static_reshape support to XNNPACK delegate

Author: GregoryComer

backends/xnnpack/operators/__init__.py

@@ -49,4 +49,5 @@
     op_static_resize_bilinear_2d,
     op_sub,
     op_to_copy,
+    op_view_copy,
 )

backends/xnnpack/operators/op_skip_ops.py

@@ -77,16 +77,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,92 @@
+# 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.
+
+from typing import cast, Dict
+
+import torch
+from executorch.backends.transforms import get_shape
+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,
+    PERM_NHWC_TO_NCHW,
+)
+from torch.fx.experimental.symbolic_shapes import free_symbols
+
+
+@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",
+        )
+        input_shape = input_node.meta["val"].shape
+
+        # 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),
+            "Symbol 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)
+
+        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

@@ -47,6 +47,7 @@
     SquareRootConfig,
     SubConfig,
     UpsampleBilinear2dConfig,
+    ViewCopyConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
     BatchNormConfig,
@@ -100,6 +101,7 @@
     SquareRootConfig,
     SubConfig,
     UpsampleBilinear2dConfig,
+    ViewCopyConfig,
     # Quant/Dequant Op Configs
     QuantizedPerTensorConfig,
     DeQuantizedPerTensorConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -336,6 +336,32 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         return torch.ops.aten.upsample_bilinear2d.vec
 
 
+class ViewCopyConfig(GenericNodePartitionerConfig):
+    target_name = "view_copy.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
+
+    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"
 

backends/xnnpack/test/ops/test_view_copy.py

@@ -0,0 +1,180 @@
+# 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 unittest
+
+import torch
+from executorch.backends.xnnpack.test.tester import Export, Tester
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import Dim
+
+
+class TestViewCopy(unittest.TestCase):
+    class View(torch.nn.Module):
+        def __init__(self, new_shape):
+            super().__init__()
+            self.new_shape = new_shape
+
+        def forward(self, x):
+            z = x.view(self.new_shape)
+            return z
+
+    def test_fp16_view_copy(self):
+        inputs = (torch.randn(4, 4).to(torch.float16),)
+        (
+            Tester(self.View((2, 8)), inputs)
+            .export()
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_view_copy(self):
+        inputs = (torch.randn(4, 4),)
+        (
+            Tester(self.View((2, 8)), inputs)
+            .export()
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_view_copy_inferred_dim(self):
+        inputs = (torch.randn(4, 4),)
+        (
+            Tester(self.View((-1, 8)), inputs)
+            .export()
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_view_copy_dynamic_shape_first_dim(self):
+        inputs = (torch.randn(4, 4),)
+        test_inputs = (
+            (torch.randn(2, 4),),
+            (torch.randn(10, 4),),
+        )
+        dynamic_shapes = {
+            "x": {0: Dim("x", min=1, max=10)},
+        }
+        tester = (
+            Tester(self.View((-1, 2)), inputs)
+            .export(Export(dynamic_shapes=dynamic_shapes))
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+        for test_input in test_inputs:
+            tester.run_method_and_compare_outputs(inputs=test_input)
+
+    def test_fp32_view_copy_dynamic_shape_last_dim(self):
+        inputs = (torch.randn(2, 4, 4),)
+        test_inputs = (
+            (torch.randn(2, 4, 2),),
+            (torch.randn(2, 4, 10),),
+        )
+        dynamic_shapes = {
+            "x": {2: Dim("x", min=1, max=10) * 2},
+        }
+        tester = (
+            Tester(self.View((-1, 4, 2)), inputs)
+            .export(Export(dynamic_shapes=dynamic_shapes))
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+        for test_input in test_inputs:
+            tester.run_method_and_compare_outputs(inputs=test_input)
+
+    def test_fp32_view_copy_unsupported_dynamism(self):
+        class SymbolicView(torch.nn.Module):
+            def forward(self, x):
+                return x.view(x.shape[0] // 2, x.shape[1] * 2)
+
+        inputs = (torch.randn(4, 4),)
+        dynamic_shapes = {
+            "x": {1: Dim("x", min=1, max=10) * 2},
+        }
+        (
+            Tester(SymbolicView(), inputs)
+            .export(Export(dynamic_shapes=dynamic_shapes))
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {  # Expect no delegation.
+                    torch.ops.higher_order.executorch_call_delegate: 0,
+                    exir_ops.edge.aten.view_copy.default: 1,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_fp32_view_copy_static_symbolic_arg(self):
+        class SymbolicView(torch.nn.Module):
+            def forward(self, x):
+                return x.view(x.shape[0] // 2, x.shape[1] * 2)
+
+        inputs = (torch.randn(4, 4),)
+        (
+            Tester(SymbolicView(), inputs)
+            .export()
+            .check_node_count({torch.ops.aten.view.default: 1})
+            .to_edge_transform_and_lower()
+            .check_node_count(
+                {
+                    torch.ops.higher_order.executorch_call_delegate: 1,
+                    exir_ops.edge.aten.view_copy.default: 0,
+                }
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )