| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Comprehensive tests for UnsafeCloneEliminationPass. |
| 4 | |
| 5 | This test suite exercises all possible valid FX graph patterns involving clones: |
| 6 | 1. Clone with no users (dead code) |
| 7 | 2. Clone with read-only users |
| 8 | 3. Clone with mutation users |
| 9 | 4. Clone of graph input |
| 10 | 5. Clone with original used after mutation |
| 11 | 6. Clone chains |
| 12 | """ |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | from torch import fx |
| 17 | from torch.fx.experimental.proxy_tensor import make_fx |
| 18 | |
| 19 | from vllm.compilation.passes.fx_utils import find_op_nodes |
| 20 | from vllm.compilation.passes.inductor_pass import get_pass_context, pass_context |
| 21 | from vllm.compilation.passes.ir.clone_elimination import ( |
| 22 | UnsafeCloneEliminationPass, |
| 23 | user_writes_to_node, |
| 24 | ) |
| 25 | from vllm.config import VllmConfig |
| 26 | from vllm.config.utils import Range |
| 27 | |
| 28 | |
| 29 | def count_clones(graph: fx.Graph) -> int: |
| 30 | """Count clone nodes in a graph.""" |
| 31 | return len(list(find_op_nodes(torch.ops.aten.clone.default, graph))) |
| 32 | |
| 33 | |
| 34 | @pytest.fixture(scope="function") |
| 35 | def clone_cleanup_pass(): |
| 36 | return UnsafeCloneEliminationPass(VllmConfig()) |
| 37 | |
| 38 | |
| 39 | @pytest.fixture(autouse=True) |
| 40 | def setup_pass_context(): |
| 41 | """Set up pass context for each test.""" |
| 42 | with pass_context(compile_range=Range(1, 8192)): |
| 43 | yield |
| 44 | |
| 45 | |
| 46 | class TestCloneCleanup: |
| 47 | """Test UnsafeCloneEliminationPass behavior on various graph patterns.""" |
| 48 | |
| 49 | def test_remove_clone_readonly_users(self, clone_cleanup_pass): |
| 50 | """Clone with only read-only users should be removed.""" |
| 51 | |
| 52 | def f(x: torch.Tensor) -> torch.Tensor: |
| 53 | x_clone = x.clone() |
| 54 | return x_clone + 1 |
| 55 | |
| 56 | inp = torch.randn(2, 3) |
| 57 | graph_module = make_fx(f)(inp) |
| 58 | assert count_clones(graph_module.graph) == 1 |
| 59 | |
| 60 | expected = graph_module(inp) |
| 61 | clone_cleanup_pass(graph_module.graph) |
| 62 | graph_module.recompile() |
| 63 | actual = graph_module(inp) |
| 64 | |
| 65 | assert count_clones(graph_module.graph) == 0 |
| 66 | torch.testing.assert_close(actual, expected) |
| 67 | |
| 68 | def test_keep_clone_with_mutation_and_original_used_after(self, clone_cleanup_pass): |
| 69 | """Clone must be kept if it's mutated AND original is used after mutation.""" |
| 70 | |
| 71 | def f(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 72 | x = x.relu() # not a graph param |
| 73 | x_clone = x.clone() |
| 74 | x_clone.add_(1) |
| 75 | return x, x_clone |
| 76 | |
| 77 | inp = torch.randn(2, 3) |
| 78 | graph_module = make_fx(f)(inp) |
| 79 | assert count_clones(graph_module.graph) == 1 |
| 80 | |
| 81 | expected = graph_module(inp) |
| 82 | clone_cleanup_pass(graph_module.graph) |
| 83 | graph_module.recompile() |
| 84 | actual = graph_module(inp) |
| 85 | |
| 86 | # Clone should be KEPT because original is used after mutation |
| 87 | assert count_clones(graph_module.graph) == 1 |
| 88 | torch.testing.assert_close(actual[0], expected[0]) |
| 89 | torch.testing.assert_close(actual[1], expected[1]) |
| 90 | |
| 91 | def test_remove_clone_with_mutation_no_original_use(self, clone_cleanup_pass): |
| 92 | """Clone can be removed if it's mutated but original is not used after.""" |
| 93 | |
| 94 | def f(x: torch.Tensor) -> torch.Tensor: |
| 95 | x = x.relu() # not a graph param |
| 96 | x_clone = x.clone() |
| 97 | x_clone.add_(1) |
| 98 | return x_clone |
| 99 | |
| 100 | inp = torch.randn(2, 3) |
| 101 | graph_module = make_fx(f)(inp) |
| 102 | assert count_clones(graph_module.graph) == 1 |
| 103 | |
| 104 | expected = graph_module(inp) |
| 105 | clone_cleanup_pass(graph_module.graph) |
| 106 | graph_module.recompile() |
| 107 | actual = graph_module(inp) |
| 108 | |
| 109 | assert count_clones(graph_module.graph) == 0 |
| 110 | torch.testing.assert_close(actual, expected) |
| 111 | |
| 112 | def test_clone_chain(self, clone_cleanup_pass): |
| 113 | """Test handling of clone chains: x -> clone1 -> clone2.""" |
| 114 | |
| 115 | def f(x: torch.Tensor) -> torch.Tensor: |
| 116 | x = x.relu() # not a graph param |
| 117 | x1 = x.clone() |
| 118 | x2 = x1.clone() |
| 119 | return x2 + 1 |
| 120 | |
| 121 | inp = torch.randn(2, 3) |
| 122 | graph_module = make_fx(f)(inp) |
| 123 | assert count_clones(graph_module.graph) == 2 |
| 124 | |
| 125 | expected = graph_module(inp) |
| 126 | clone_cleanup_pass(graph_module.graph) |
| 127 | graph_module.recompile() |
| 128 | actual = graph_module(inp) |
| 129 | |
| 130 | # Both clones should be removed |
| 131 | assert count_clones(graph_module.graph) == 0 |
| 132 | torch.testing.assert_close(actual, expected) |
| 133 | |
| 134 | def test_keep_clone_that_changes_layout(self, clone_cleanup_pass): |
| 135 | """Clone must be kept when it materializes a compact slice layout.""" |
| 136 | |
| 137 | def f(x: torch.Tensor) -> torch.Tensor: |
| 138 | return x[:, :3].contiguous() |
| 139 | |
| 140 | inp = torch.randn(4, 5) |
| 141 | graph_module = make_fx(f)(inp) |
| 142 | assert count_clones(graph_module.graph) == 1 |
| 143 | |
| 144 | expected = graph_module(inp) |
| 145 | clone_cleanup_pass(graph_module.graph) |
| 146 | graph_module.recompile() |
| 147 | actual = graph_module(inp) |
| 148 | |
| 149 | assert count_clones(graph_module.graph) == 1 |
| 150 | assert actual.stride() == expected.stride() == (3, 1) |
| 151 | torch.testing.assert_close(actual, expected) |
| 152 | |
| 153 | def test_multiple_clones_of_same_input(self, clone_cleanup_pass): |
| 154 | """Test multiple independent clones of the same input.""" |
| 155 | |
| 156 | def f(x: torch.Tensor) -> torch.Tensor: |
| 157 | x1 = x.clone() |
| 158 | x2 = x.clone() |
| 159 | return x1 + x2 |
| 160 | |
| 161 | inp = torch.randn(2, 3) |
| 162 | graph_module = make_fx(f)(inp) |
| 163 | assert count_clones(graph_module.graph) == 2 |
| 164 | |
| 165 | expected = graph_module(inp) |
| 166 | clone_cleanup_pass(graph_module.graph) |
| 167 | graph_module.recompile() |
| 168 | actual = graph_module(inp) |
| 169 | |
| 170 | # Both clones should be removed (only readonly uses) |
| 171 | assert count_clones(graph_module.graph) == 0 |
| 172 | torch.testing.assert_close(actual, expected) |
| 173 | |
| 174 | def test_no_clones_in_graph(self, clone_cleanup_pass): |
| 175 | """Test pass behavior when graph has no clones.""" |
| 176 | |
| 177 | def f(x: torch.Tensor) -> torch.Tensor: |
| 178 | return x + 1 |
| 179 | |
| 180 | inp = torch.randn(2, 3) |
| 181 | graph_module = make_fx(f)(inp) |
| 182 | assert count_clones(graph_module.graph) == 0 |
| 183 | |
| 184 | expected = graph_module(inp) |
| 185 | clone_cleanup_pass(graph_module.graph) |
| 186 | graph_module.recompile() |
| 187 | actual = graph_module(inp) |
| 188 | |
| 189 | assert count_clones(graph_module.graph) == 0 |
| 190 | torch.testing.assert_close(actual, expected) |
| 191 | |
| 192 | def test_multiple_passes(self, clone_cleanup_pass): |
| 193 | """Test running the pass multiple times (should be idempotent).""" |
| 194 | |
| 195 | def f(x: torch.Tensor) -> torch.Tensor: |
| 196 | x1 = x.clone() |
| 197 | return x1 + 1 |
| 198 | |
| 199 | inp = torch.randn(2, 3) |
| 200 | graph_module = make_fx(f)(inp) |
| 201 | assert count_clones(graph_module.graph) == 1 |
| 202 | |
| 203 | expected = graph_module(inp) |
| 204 | |
| 205 | clone_cleanup_pass(graph_module.graph) |
| 206 | assert count_clones(graph_module.graph) == 0 |
| 207 | graph_module.recompile() |
| 208 | actual = graph_module(inp) |
| 209 | torch.testing.assert_close(actual, expected) |
| 210 | |
| 211 | clone_cleanup_pass(graph_module.graph) |
| 212 | assert count_clones(graph_module.graph) == 0 |
| 213 | graph_module.recompile() |
| 214 | actual = graph_module(inp) |
| 215 | torch.testing.assert_close(actual, expected) |
| 216 | |
| 217 | def test_output_node_no_write(self): |
| 218 | """Output nodes never write to their inputs.""" |
| 219 | |
| 220 | def f(x: torch.Tensor) -> torch.Tensor: |
| 221 | return x |
| 222 | |
| 223 | graph_module = make_fx(f)(torch.randn(2, 3)) |
| 224 | x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] |
| 225 | output_node = [n for n in graph_module.graph.nodes if n.op == "output"][0] |
| 226 | |
| 227 | assert not user_writes_to_node(output_node, x_node) |
| 228 | |
| 229 | def test_readonly_op_no_write(self): |
| 230 | """Readonly operations don't write to inputs.""" |
| 231 | |
| 232 | def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 233 | return x + y |
| 234 | |
| 235 | graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) |
| 236 | placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] |
| 237 | add_node = [ |
| 238 | n |
| 239 | for n in graph_module.graph.nodes |
| 240 | if n.op == "call_function" and n.target == torch.ops.aten.add.Tensor |
| 241 | ][0] |
| 242 | |
| 243 | assert not user_writes_to_node(add_node, placeholders[0]) |
| 244 | assert not user_writes_to_node(add_node, placeholders[1]) |
| 245 | |
| 246 | def test_inplace_op_writes(self): |
| 247 | """Inplace operations write to first argument.""" |
| 248 | |
| 249 | def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 250 | x.add_(y) |
| 251 | return x |
| 252 | |
| 253 | graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) |
| 254 | placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] |
| 255 | add_node = [ |
| 256 | n |
| 257 | for n in graph_module.graph.nodes |
| 258 | if n.op == "call_function" and "add_" in str(n.target) |
| 259 | ][0] |
| 260 | |
| 261 | # add_ writes to first arg but not second |
| 262 | assert user_writes_to_node(add_node, placeholders[0]) |
| 263 | assert not user_writes_to_node(add_node, placeholders[1]) |
| 264 | |
| 265 | def test_copy_writes(self): |
| 266 | """copy_ operation writes to first argument.""" |
| 267 | |
| 268 | def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 269 | x.copy_(y) |
| 270 | return x |
| 271 | |
| 272 | graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) |
| 273 | placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] |
| 274 | copy_node = [ |
| 275 | n |
| 276 | for n in graph_module.graph.nodes |
| 277 | if n.op == "call_function" and "copy_" in str(n.target) |
| 278 | ][0] |
| 279 | |
| 280 | assert user_writes_to_node(copy_node, placeholders[0]) |
| 281 | assert not user_writes_to_node(copy_node, placeholders[1]) |
| 282 | |
| 283 | def test_auto_functionalized_not_a_write(self): |
| 284 | """auto_functionalized ops are follow-up uses, not writes.""" |
| 285 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 286 | |
| 287 | def f(x: torch.Tensor) -> torch.Tensor: |
| 288 | return x |
| 289 | |
| 290 | graph_module = make_fx(f)(torch.randn(2, 3)) |
| 291 | x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] |
| 292 | |
| 293 | # Create an auto_functionalized node in the graph |
| 294 | with graph_module.graph.inserting_before(None): |
| 295 | af_node = graph_module.graph.call_function( |
| 296 | auto_functionalized, kwargs={"input": x_node} |
| 297 | ) |
| 298 | |
| 299 | # auto_functionalized should not be treated as a write |
| 300 | assert not user_writes_to_node(af_node, x_node) |
| 301 | |
| 302 | def test_higher_order_op_conservatively_writes(self): |
| 303 | """Other higher-order operators are conservatively treated as writes.""" |
| 304 | from torch._ops import HigherOrderOperator |
| 305 | |
| 306 | def f(x: torch.Tensor) -> torch.Tensor: |
| 307 | return x |
| 308 | |
| 309 | graph_module = make_fx(f)(torch.randn(2, 3)) |
| 310 | x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] |
| 311 | |
| 312 | # Create a concrete higher-order operator subclass |
| 313 | class MockHigherOrderOp(HigherOrderOperator): |
| 314 | def __call__(self, *args, **kwargs): |
| 315 | return args[0] if args else None |
| 316 | |
| 317 | mock_hoo = MockHigherOrderOp("mock_higher_order_op") |
| 318 | |
| 319 | with graph_module.graph.inserting_before(None): |
| 320 | hoo_node = graph_module.graph.call_function(mock_hoo, args=(x_node,)) |
| 321 | |
| 322 | # Should be conservative and assume it could write |
| 323 | assert user_writes_to_node(hoo_node, x_node) |
| 324 | |
| 325 | |
| 326 | class TestCloneCleanupWithDonatedInputs: |
| 327 | """Test UnsafeCloneEliminationPass with donated input tracking via PassContext.""" |
| 328 | |
| 329 | @pytest.fixture(autouse=True) |
| 330 | def setup_pass_context(self): |
| 331 | """Set up pass context for each test.""" |
| 332 | with pass_context(compile_range=Range(1, 8192)): |
| 333 | yield |
| 334 | |
| 335 | def test_donated_input_clone_removed(self, clone_cleanup_pass): |
| 336 | """Clone of donated input should be removed.""" |
| 337 | |
| 338 | def f(x: torch.Tensor) -> torch.Tensor: |
| 339 | x_clone = x.clone() |
| 340 | x_clone.add_(1) |
| 341 | return x_clone |
| 342 | |
| 343 | inp = torch.randn(2, 3) |
| 344 | graph_module = make_fx(f)(inp) |
| 345 | assert count_clones(graph_module.graph) == 1 |
| 346 | |
| 347 | # Mark first parameter as donated |
| 348 | get_pass_context().donated_input_ids = {0} |
| 349 | |
| 350 | expected = graph_module(inp.clone()) |
| 351 | clone_cleanup_pass(graph_module.graph) |
| 352 | graph_module.recompile() |
| 353 | |
| 354 | # Clone should be removed since input is donated |
| 355 | assert count_clones(graph_module.graph) == 0 |
| 356 | |
| 357 | # Input can be mutated (donated) |
| 358 | inp_copy = inp.clone() |
| 359 | actual = graph_module(inp_copy) |
| 360 | torch.testing.assert_close(actual, expected) |
| 361 | |
| 362 | def test_non_donated_input_clone_kept(self, clone_cleanup_pass): |
| 363 | """Clone of non-donated input with mutation should be kept.""" |
| 364 | |
| 365 | def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 366 | x_clone = x.clone() |
| 367 | x_clone.add_(1) |
| 368 | return x, x_clone |
| 369 | |
| 370 | inp_x = torch.randn(2, 3) |
| 371 | inp_y = torch.randn(2, 3) |
| 372 | graph_module = make_fx(f)(inp_x, inp_y) |
| 373 | assert count_clones(graph_module.graph) == 1 |
| 374 | |
| 375 | # No donated inputs |
| 376 | get_pass_context().donated_input_ids = set() |
| 377 | |
| 378 | expected = graph_module(inp_x.clone(), inp_y.clone()) |
| 379 | clone_cleanup_pass(graph_module.graph) |
| 380 | graph_module.recompile() |
| 381 | |
| 382 | # Clone should be kept since input is not donated and original is used |
| 383 | assert count_clones(graph_module.graph) == 1 |
| 384 | |
| 385 | # Verify inputs are not mutated |
| 386 | inp_x_before = inp_x.clone() |
| 387 | inp_y_before = inp_y.clone() |
| 388 | actual = graph_module(inp_x, inp_y) |
| 389 | torch.testing.assert_close( |
| 390 | inp_x, inp_x_before, msg="Input x should not be mutated" |
| 391 | ) |
| 392 | torch.testing.assert_close( |
| 393 | inp_y, inp_y_before, msg="Input y should not be mutated" |
| 394 | ) |
| 395 | torch.testing.assert_close(actual[0], expected[0]) |
| 396 | torch.testing.assert_close(actual[1], expected[1]) |
| 397 | |
| 398 | def test_mixed_donated_inputs(self, clone_cleanup_pass): |
| 399 | """Test with some inputs donated and some not.""" |
| 400 | |
| 401 | def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 402 | x_clone = x.clone() |
| 403 | x_clone.add_(1) |
| 404 | y_clone = y.clone() |
| 405 | y_clone.add_(2) |
| 406 | return x_clone, y_clone |
| 407 | |
| 408 | inp_x = torch.randn(2, 3) |
| 409 | inp_y = torch.randn(2, 3) |
| 410 | graph_module = make_fx(f)(inp_x, inp_y) |
| 411 | assert count_clones(graph_module.graph) == 2 |
| 412 | |
| 413 | # Only x is donated |
| 414 | get_pass_context().donated_input_ids = {0} |
| 415 | |
| 416 | expected = graph_module(inp_x.clone(), inp_y.clone()) |
| 417 | clone_cleanup_pass(graph_module.graph) |
| 418 | graph_module.recompile() |
| 419 | |
| 420 | # x_clone removed (x is donated), y_clone kept (y is not donated) |
| 421 | assert count_clones(graph_module.graph) == 1 |
| 422 | |
| 423 | # Verify y is not mutated (x can be mutated since it's donated) |
| 424 | inp_y_before = inp_y.clone() |
| 425 | actual = graph_module(inp_x.clone(), inp_y) |
| 426 | torch.testing.assert_close( |
| 427 | inp_y, inp_y_before, msg="Input y should not be mutated" |
| 428 | ) |
| 429 | torch.testing.assert_close(actual[0], expected[0]) |
| 430 | torch.testing.assert_close(actual[1], expected[1]) |