partition to_dim_order_copy in XNN delegate

Author: leafs1

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -400,6 +400,11 @@ 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:
+                if node.meta["val"].is_contiguous():
+                    self.mark_as_nchw_node(node)
+                else:
+                    self.mark_as_nhwc_node(node)
             else:
                 # The node can have inputs in any format (but all must be the
                 # same format)

backends/xnnpack/partition/config/__init__.py

@@ -48,6 +48,7 @@
     SoftmaxConfig,
     SquareRootConfig,
     SubConfig,
+    ToDimOrderCopyConfig,
     UpsampleBilinear2dConfig,
 )
 from executorch.backends.xnnpack.partition.config.node_configs import (
@@ -97,6 +98,7 @@
     PreluConfig,
     ReciprocalSquareRootConfig,
     ReLUConfig,
+    ToDimOrderCopyConfig,
     # SDPAConfig, TODO: D60553559: preserving SDPA for fairseq fails
     SigmoidConfig,
     SliceCopyConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -364,6 +364,13 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32]
 
 
+class ToDimOrderCopyConfig(GenericNodePartitionerConfig):
+    target_name = "_to_dim_order_copy.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+
 class MeanDimConfig(GenericNodePartitionerConfig):
     target_name = "mean.dim"
 

backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py

@@ -130,87 +130,22 @@ def test_fp32_channels_last_tagged_reshape_pass(self):
                 .run_method_and_compare_outputs()
             )
 
-    class LinearConv(torch.nn.Module):
+    class LinearConvDimSwap(torch.nn.Module):
         def __init__(self):
             super().__init__()
             self.conv1 = torch.nn.Conv2d(3, 3, 3)
             self.linear1 = torch.nn.Linear(4, 3)
 
         def forward(self, x):
             y = self.linear1(x)
+            y = y.to(memory_format=torch.channels_last)
+            y = y.to(memory_format=torch.contiguous_format)
             return self.conv1(y)
 
-    def test_conv_linear_dim_order_swaps_on_nhwc_input(self):
-        tester = Tester(
-            self.LinearConv().eval(),
-            (torch.randn(1, 3, 6, 4).to(memory_format=torch.channels_last),),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
+    LinearConvDimSwapModule = LinearConvDimSwap()
 
-    def test_conv_linear_dim_order_swaps_on_nchw_input(self):
-        tester = Tester(
-            self.LinearConv().eval(),
-            (torch.randn(1, 3, 6, 4),),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
-
-    class ConvLinearConv(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.conv1 = torch.nn.Conv2d(3, 3, 3)
-            self.linear1 = torch.nn.Linear(4, 4)
-
-        def forward(self, x):
-            y = self.conv1(x)
-            return self.linear1(y)
-
-    def test_linear_conv_dim_order_swaps_on_nhwc_input(self):
-        tester = Tester(
-            self.ConvLinearConv().eval(),
-            (torch.randn(1, 3, 6, 6).to(memory_format=torch.channels_last),),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
-
-    def test_linear_conv_dim_order_swaps_on_nchw_input(self):
-        tester = Tester(
-            self.ConvLinearConv().eval(),
-            (torch.randn(1, 3, 6, 6),),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
-
-    class Bilinear(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-
-        def forward(self, x):
-            return torch.nn.functional.interpolate(
-                x, scale_factor=2, mode="bilinear", align_corners=True
-            )
-
-    def test_nhwc_input_on_nhwc_op(self):
-        tester = Tester(
-            self.Bilinear().eval(),
-            (
-                torch.arange(8)
-                .reshape(1, 2, 2, 2)
-                .to(torch.float32)
-                .to(memory_format=torch.channels_last),
-            ),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
-
-    def test_nchw_input_on_nhwc_op(self):
-        tester = Tester(
-            self.Bilinear().eval(),
-            (torch.arange(8).reshape(1, 2, 2, 2).to(torch.float32),),
-        )
-
-        tester.export().to_edge_transform_and_lower().to_executorch().serialize().run_method_and_compare_outputs()
+    def test_conv_linear_dim_order_swap_partitioner(self):
+        self.run_tester(self.LinearConvDimSwapModule, (torch.randn(1, 3, 6, 4),))
 
     def test_qs8_channels_last_tagged_reshape_pass(self):
         for module, num_reshape in self.modules.items():