Register _clone_dim_order and add operator tests

Author: keyprocedure

exir/passes/dim_order_ops_registry.py

@@ -28,6 +28,14 @@
     "_empty_dim_order.out(int[] size, *, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)"
 )
 
+lib.define(
+    "_clone_dim_order(Tensor self, *, bool non_blocking=False, int[]? dim_order=None) -> Tensor"
+)
+
+lib.define(
+    "_clone_dim_order.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)"
+)
+
 
 def _op_impl(target, *args, **kwargs):
     kwargs["memory_format"] = get_memory_format(kwargs.get("dim_order", None))
@@ -57,6 +65,16 @@ def _empty_dim_order_out_impl(*args, **kwargs):
     return _op_impl(torch.ops.aten.empty.out, *args, **kwargs)
 
 
+@impl(lib, "_clone_dim_order", "CompositeImplicitAutograd")
+def _clone_dim_order_impl(*args, **kwargs):
+    return _op_impl(torch.ops.aten.clone.default, *args, **kwargs)
+
+
+@impl(lib, "_clone_dim_order.out", "CompositeImplicitAutograd")
+def _clone_dim_order_out_impl(*args, **kwargs):
+    return _op_impl(torch.ops.aten.clone.out, *args, **kwargs)
+
+
 """
 Defines a map of edge ops to the corresponding dim_order ops for quick lookup
 """

exir/tests/test_memory_format_ops_pass.py

@@ -28,8 +28,6 @@
     MemoryFormatOpsPassTestUtils,
     MemoryFormatTestSet,
     PropagateToCopyChannalsLastModule,
-    SimpleCloneChannelsLastModule,
-    SimpleCloneContiguousModule,
     SimpleEmptyChannelLastModule,
     SimpleEmptyContiguoustModule,
     SimpleToCopyChannelsLastModule,
@@ -93,36 +91,40 @@ 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_dim_order_preserves_channels_last(self):
+        x = torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last)
+        y = torch.ops.dim_order_ops._clone_dim_order.default(x)
+
+        assert y.is_contiguous(
+            memory_format=torch.channels_last
+        ), "_clone_dim_order output is not in channels_last memory format."
+        assert torch.allclose(x, y)
+
+    def test_op_clone_dim_order_to_contiguous(self):
+        x = torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last)
+        contiguous_dim_order = get_dim_order(torch.contiguous_format, x.dim())
+        y = torch.ops.dim_order_ops._clone_dim_order.default(
+            x, dim_order=contiguous_dim_order
         )
 
-    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,
-            ),
+        assert (
+            y.is_contiguous()
+        ), "_clone_dim_order output is not in contiguous memory format"
+        assert torch.allclose(x, y)
+
+    def test_op_clone_dim_order_out_to_channels_last(self):
+        x = torch.randn(2, 3, 4, 5).contiguous()
+        y = torch.empty_like(x, memory_format=torch.channels_last)
+        channels_last_dim_order = get_dim_order(torch.channels_last, y.dim())
+        torch.ops.dim_order_ops._clone_dim_order.out(
+            x, dim_order=channels_last_dim_order, out=y
         )
 
+        assert y.is_contiguous(
+            memory_format=torch.channels_last
+        ), "_clone_dim_order output is not in channels_last memory format"
+        assert torch.allclose(x, y)
+
     def test_op_dim_order_update(self) -> None:
         MemoryFormatOpsPassTestUtils.memory_format_test_runner(
             self,