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/runtime/XNNExecutor.cpp

@@ -114,13 +114,8 @@ ET_NODISCARD Error XNNExecutor::prepare_args(EValue** args) {
           XNN_MAX_TENSOR_DIMS,
           num_dims);
 
-<<<<<<< HEAD
       for (int j = 0; j < num_dims; ++j) {
         dims[j] = tensor->size(static_cast<int>(dim_order[j]));
-=======
-      for (int i = 0; i < num_dims; ++i) {
-        dims[i] = tensor->size(static_cast<int>(dim_order[i]));
->>>>>>> b5785b2fda (support channels last dim order in xnnpack)
       }
       status =
           xnn_reshape_external_value(runtime_.get(), ext_id, num_dims, dims);
@@ -229,15 +224,9 @@ ET_NODISCARD Error XNNExecutor::resize_outputs(EValue** args) const {
         Internal,
         "Failed to retrieve dim order from tensor!");
 
-<<<<<<< HEAD
     for (int j = 0; j < num_dim; ++j) {
       expected_output_size[static_cast<int>(dim_order[j])] =
           static_cast<SizesType>(dims[j]);
-=======
-    for (int i = 0; i < num_dim; ++i) {
-      expected_output_size[static_cast<int>(dim_order[i])] =
-          static_cast<SizesType>(dims[i]);
->>>>>>> b5785b2fda (support channels last dim order in xnnpack)
     }
 
     executorch::aten::ArrayRef<SizesType> output_size{

backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py

@@ -130,87 +130,23 @@ 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()
-
-    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),
-            ),
-        )
+    LinearConvDimSwapModule = LinearConvDimSwap()
 
-        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_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_qs8_channels_last_tagged_reshape_pass(self):
         for module, num_reshape in self.modules.items():