Add aten.clone memory_format check in RemoveCloneOpsTransform

Author: keyprocedure

backends/transforms/remove_clone_ops.py

@@ -15,29 +15,46 @@ def remove_clone_ops(graph: torch.fx.Graph) -> torch.fx.Graph:
     """
     Remove clone op nodes that have the same dim_order as their input, and replace their uses with the input node.
     """
-    clone_op = exir_ops.edge.aten.clone.default
-    clone_dim_order_op = exir_ops.edge.dim_order_ops._clone_dim_order.default
-
     for node in graph.nodes:
         if node.op != "call_function":
             continue
 
-        # Identify clone_dim_order ops with unchanged memory layout.
-        unchanged_layout_clone = (
-            node.target == clone_dim_order_op
-            and "val" in node.meta
-            and "val" in node.args[0].meta
-            and node.meta["val"].dim_order() == node.args[0].meta["val"].dim_order()
-        )
-
-        if node.target == clone_op or unchanged_layout_clone:
+        if is_unchanged_clone(node) or is_unchanged_dim_order_clone(node):
             with graph.inserting_after(node):
                 node.replace_all_uses_with(node.args[0])
 
     graph.eliminate_dead_code()
     return graph
 
 
+def is_unchanged_clone(node: torch.fx.Node) -> bool:
+    """Determine if aten.clone has unchanged memory format."""
+    if node.target != exir_ops.edge.aten.clone.default:
+        return False
+
+    memory_format = node.kwargs.get("memory_format")
+    if memory_format in (None, torch.preserve_format):
+        return True
+
+    input_meta = node.args[0].meta
+    return "val" in input_meta and input_meta["val"].is_contiguous(
+        memory_format=memory_format
+    )
+
+
+def is_unchanged_dim_order_clone(node: torch.fx.Node) -> bool:
+    """Determine if _clone_dim_order has unchanged dim order."""
+    if node.target != exir_ops.edge.dim_order_ops._clone_dim_order.default:
+        return False
+
+    input_meta = node.args[0].meta
+    return (
+        "val" in node.meta
+        and "val" in input_meta
+        and node.meta["val"].dim_order() == input_meta["val"].dim_order()
+    )
+
+
 class RemoveCloneOpsTransform(ExportPass):
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph_module.graph = remove_clone_ops(graph_module.graph)

exir/tests/test_memory_format_ops_pass.py

@@ -326,56 +326,72 @@ def call_operator(self, op, args, kwargs, meta):
         self.assertTrue(is_contiguous_dim_order(expected))
 
     def test_op_clone_replacement_channels_last_survives(self):
-        _clone_dim_order_op_str = (
-            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
-        )
+        clone_op_cases = [
+            # Case testing aten.clone by setting _skip_dim_order to True
+            (True, "executorch_exir_dialects_edge__ops_aten_clone_default"),
+            # Case testing _clone_dim_order by setting _skip_dim_order to False
+            (
+                False,
+                "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default",
+            ),
+        ]
 
-        model = SimpleCloneChannelsLastModule()
-        x = torch.randn(3, 4, 5, 6).to(memory_format=torch.contiguous_format)
+        for skip_dim_order, clone_op_str in clone_op_cases:
+            model = SimpleCloneChannelsLastModule()
+            x = torch.randn(3, 4, 5, 6).to(memory_format=torch.contiguous_format)
 
-        exported = export(model.eval(), (x,), strict=True)
-        before_epm = to_edge(
-            exported, compile_config=EdgeCompileConfig(_skip_dim_order=False)
-        )
+            exported = export(model.eval(), (x,), strict=True)
+            before_epm = to_edge(
+                exported,
+                compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order),
+            )
 
-        updated_epm = before_epm.transform([RemoveCloneOpsTransform()])
+            updated_epm = before_epm.transform([RemoveCloneOpsTransform()])
 
-        FileCheck().check_count(_clone_dim_order_op_str, 1, exactly=True).run(
-            updated_epm.exported_program().graph_module.code
-        )
+            FileCheck().check_count(clone_op_str, 1, exactly=True).run(
+                updated_epm.exported_program().graph_module.code
+            )
 
-        expected = before_epm.exported_program().module()(x)
-        actual = updated_epm.exported_program().module()(x)
-        assert torch.allclose(actual, expected)
-        assert is_channel_last_dim_order(actual)
+            expected = before_epm.exported_program().module()(x)
+            actual = updated_epm.exported_program().module()(x)
+            assert torch.allclose(actual, expected)
+            assert is_channel_last_dim_order(actual)
 
     def test_op_clone_without_transformation_removed(self):
-        _clone_dim_order_op_str = (
-            "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default"
-        )
+        clone_op_cases = [
+            # Case testing aten.clone by setting _skip_dim_order to True
+            (True, "executorch_exir_dialects_edge__ops_aten_clone_default"),
+            # Case testing _clone_dim_order by setting _skip_dim_order to False
+            (
+                False,
+                "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default",
+            ),
+        ]
 
-        model = SimpleCloneChannelsLastModule()
-        x = torch.randn(3, 4, 5, 6).to(memory_format=torch.channels_last)
+        for skip_dim_order, clone_op_str in clone_op_cases:
+            model = SimpleCloneChannelsLastModule()
+            x = torch.randn(3, 4, 5, 6).to(memory_format=torch.channels_last)
 
-        exported = export(model.eval(), (x,), strict=True)
-        before_epm = to_edge(
-            exported, compile_config=EdgeCompileConfig(_skip_dim_order=False)
-        )
+            exported = export(model.eval(), (x,), strict=True)
+            before_epm = to_edge(
+                exported,
+                compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order),
+            )
 
-        FileCheck().check_count(_clone_dim_order_op_str, 1, exactly=True).run(
-            before_epm.exported_program().graph_module.code
-        )
+            FileCheck().check_count(clone_op_str, 1, exactly=True).run(
+                before_epm.exported_program().graph_module.code
+            )
 
-        updated_epm = before_epm.transform([RemoveCloneOpsTransform()])
+            updated_epm = before_epm.transform([RemoveCloneOpsTransform()])
 
-        FileCheck().check_not(_clone_dim_order_op_str).run(
-            updated_epm.exported_program().graph_module.code
-        )
+            FileCheck().check_not(clone_op_str).run(
+                updated_epm.exported_program().graph_module.code
+            )
 
-        expected = before_epm.exported_program().module()(x)
-        actual = updated_epm.exported_program().module()(x)
-        assert torch.allclose(actual, expected)
-        assert is_channel_last_dim_order(actual)
+            expected = before_epm.exported_program().module()(x)
+            actual = updated_epm.exported_program().module()(x)
+            assert torch.allclose(actual, expected)
+            assert is_channel_last_dim_order(actual)
 
     def test_resnet18(self) -> None:
         model = torchvision.models.resnet18()