Preserve layout-changing clones (#44574)

Signed-off-by: Michael Gschwind <mgschwind@nvidia.com>
Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>

tests/compile/passes/ir/test_clone_cleanup.py

@@ -132,6 +132,25 @@ class TestCloneCleanup:
         assert count_clones(graph_module.graph) == 0
         torch.testing.assert_close(actual, expected)
 
+    def test_keep_clone_that_changes_layout(self, clone_cleanup_pass):
+        """Clone must be kept when it materializes a compact slice layout."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            return x[:, :3].contiguous()
+
+        inp = torch.randn(4, 5)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        assert count_clones(graph_module.graph) == 1
+        assert actual.stride() == expected.stride() == (3, 1)
+        torch.testing.assert_close(actual, expected)
+
     def test_multiple_clones_of_same_input(self, clone_cleanup_pass):
         """Test multiple independent clones of the same input."""
 

vllm/compilation/passes/ir/clone_elimination.py

@@ -16,6 +16,26 @@ from ..vllm_inductor_pass import VllmInductorPass
 logger = init_logger(__name__)
 
 
+def clone_preserves_layout(node: fx.Node, original_node: fx.Node) -> bool:
+    node_val = node.meta.get("val")
+    original_val = original_node.meta.get("val")
+    if node_val is None or original_val is None:
+        return True
+
+    try:
+        node_stride = tuple(node_val.stride())
+        original_stride = tuple(original_val.stride())
+        node_storage_offset = node_val.storage_offset()
+        original_storage_offset = original_val.storage_offset()
+    except (AttributeError, RuntimeError):
+        return True
+
+    return (
+        node_stride == original_stride
+        and node_storage_offset == original_storage_offset
+    )
+
+
 def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool:
     if user.op == "output":
         return False
@@ -79,6 +99,15 @@ class UnsafeCloneEliminationPass(VllmInductorPass):
             original_node = node.args[0]
             assert isinstance(original_node, fx.Node)
 
+            if not clone_preserves_layout(node, original_node):
+                logger.debug(
+                    "Clone removal not possible, clone changes layout: "
+                    "original_node=%s node=%s",
+                    original_node,
+                    node,
+                )
+                continue
+
             # Clone needs to be preserved if node is getting written to and
             # the old value is used again.
             # This could only happen if an inplace implementation was lowered.