Register clone_dim_order op; add test for op replacement

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,12 +65,23 @@ 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
 """
 DimOrderOpsMap = {
     exir_ops.edge.aten._to_copy.default: exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
     exir_ops.edge.aten.empty.memory_format: exir_ops.edge.dim_order_ops._empty_dim_order.default,
+    exir_ops.edge.aten.clone.default: exir_ops.edge.dim_order_ops._clone_dim_order.default,
 }
 
 """

exir/tests/test_memory_format_ops_pass.py

@@ -28,6 +28,7 @@
     MemoryFormatOpsPassTestUtils,
     MemoryFormatTestSet,
     PropagateToCopyChannalsLastModule,
+    SimpleCloneChannelsLastModule,
     SimpleEmptyChannelLastModule,
     SimpleEmptyContiguoustModule,
     SimpleToCopyChannelsLastModule,
@@ -389,3 +390,17 @@ def test_mobilenet_v3_xnnpack(self) -> None:
                 rtol=1e-3,
             ),
         )
+
+    def test_op_clone_dim_order_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))
+
+        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

@@ -38,6 +38,10 @@
         "torch.ops.aten.empty.memory_format",
         "executorch_exir_dialects_edge__ops_dim_order_ops__empty_dim_order_default",
     ),
+    torch.ops.aten.clone.default: (
+        "torch.ops.aten.clone.default",
+        "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default",
+    ),
 }
 
 
@@ -70,6 +74,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x.to(dtype=torch.double, memory_format=torch.channels_last)
 
 
+class SimpleCloneContiguousModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.clone(memory_format=torch.contiguous_format)
+
+
+class SimpleCloneChannelsLastModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.clone(memory_format=torch.channels_last)
+
+
 class SimpleEmptyContiguoustModule(torch.nn.Module):
     def __init__(self):
         super().__init__()