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

@@ -173,6 +173,23 @@ def test_fp32_channels_last_tagged_reshape_pass(self):
                 .run_method_and_compare_outputs()
             )
 
+    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)
+
+    LinearConvDimSwapModule = LinearConvDimSwap()
+
+    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():
             (