Add graph level and end to end tests for _clone_dim_order op

Author: keyprocedure

exir/tests/test_memory_format_ops_pass.py

@@ -27,8 +27,10 @@
     AmbiguousDimOrderError,
     MemoryFormatOpsPassTestUtils,
     MemoryFormatTestSet,
+    PropagateToCloneChannelsLastModule,
     PropagateToCopyChannalsLastModule,
     SimpleCloneChannelsLastModule,
+    SimpleCloneContiguousModule,
     SimpleEmptyChannelLastModule,
     SimpleEmptyContiguoustModule,
     SimpleToCopyChannelsLastModule,
@@ -92,6 +94,36 @@ def test_op_empty_replacement_contiguous(self) -> None:
             ),
         )
 
+    def test_op_clone_replacement_contiguous(self) -> None:
+        model = SimpleCloneContiguousModule()
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                op=torch.ops.aten.clone.default,
+                sample_input=(
+                    torch.randn((3, 4, 5, 6)).to(memory_format=torch.channels_last),
+                ),
+                target_memory_format=torch.contiguous_format,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+            ),
+        )
+
+    def test_op_clone_replacement_channels_last(self) -> None:
+        model = SimpleCloneChannelsLastModule()
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                op=torch.ops.aten.clone.default,
+                sample_input=(
+                    torch.randn((3, 4, 5, 6)).to(memory_format=torch.contiguous_format),
+                ),
+                target_memory_format=torch.channels_last,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+            ),
+        )
+
     def test_op_dim_order_update(self) -> None:
         MemoryFormatOpsPassTestUtils.memory_format_test_runner(
             self,
@@ -129,6 +161,25 @@ def test_op_dim_order_propagation(self) -> None:
             check_unambiguous_dim_order=True,
         )
 
+    def test_op_clone_dim_order_propagation(self) -> None:
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=PropagateToCloneChannelsLastModule().eval(),
+                op=torch.ops.aten.clone.default,
+                sample_input=(
+                    torch.rand_like(
+                        torch.zeros([2, 2, 2, 2]),
+                        dtype=torch.float32,
+                        memory_format=torch.contiguous_format,
+                    ),
+                ),
+                target_memory_format=torch.channels_last,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+            ),
+            check_unambiguous_dim_order=True,
+        )
+
     def test_op_dim_order_propagation_ambiguous(self) -> None:
         try:
             MemoryFormatOpsPassTestUtils.memory_format_test_runner(
@@ -154,6 +205,29 @@ def test_op_dim_order_propagation_ambiguous(self) -> None:
         except AmbiguousDimOrderError:
             pass  # Expected error
 
+    def test_op_clone_dim_order_graph_replacement(self):
+        model = SimpleCloneChannelsLastModule()
+        x = torch.randn(3, 4, 5, 6).to(memory_format=torch.contiguous_format)
+        _clone_dim_order_op_str = (
+            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
+        )
+
+        exported = export(model.eval(), (x,), strict=True)
+        epm = to_edge(exported, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+
+        # Verify one _clone_dim_order op exists and aten.clone.default nodes have been removed.
+        (
+            FileCheck()
+            .check_not(
+                "aten.clone.default"
+            )  # Check before first _clone_dim_order_op_str match.
+            .check_count(_clone_dim_order_op_str, 1, exactly=True)
+            .check_not(
+                "aten.clone.default"
+            )  # Check after _clone_dim_order_op_str match.
+            .run(epm.exported_program().graph_module.code)
+        )
+
     # Only test dim order replacement result in lean mode test.
     # This test is irrelevant with operator mode.
     def test_dim_order_replacement(self) -> None:
@@ -390,17 +464,3 @@ def test_mobilenet_v3_xnnpack(self) -> None:
                 rtol=1e-3,
             ),
         )
-
-    def test_op_clone_dim_order_registration(self):
-        model = SimpleCloneChannelsLastModule()
-        x = torch.randn(3, 4, 5, 6).to(memory_format=torch.contiguous_format)
-        clone_dim_order_op_str = (
-            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
-        )
-
-        exported = export(model.eval(), (x,), strict=True)
-        epm = to_edge(exported, compile_config=EdgeCompileConfig(_skip_dim_order=False))
-
-        FileCheck().check_count(clone_dim_order_op_str, 1, exactly=True).run(
-            epm.exported_program().graph_module.code
-        )

exir/tests/test_memory_format_ops_pass_utils.py

@@ -122,6 +122,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return t1 * t2
 
 
+class PropagateToCloneChannelsLastModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        t1 = x.clone(memory_format=torch.channels_last)
+        t2 = t1 + t1
+        return t1 * t2
+
+
 class AmbiguousDimOrderError(RuntimeError):
     pass