|
5 | 5 |
|
6 | 6 | import numpy as np |
7 | 7 | import pytest |
8 | | -from helpers.graph_kernels import compile_common_kernels |
| 8 | +from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels |
9 | 9 | from helpers.marks import requires_module |
| 10 | +from helpers.misc import try_create_condition |
10 | 11 |
|
11 | 12 | from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch |
12 | | -from cuda.core.graph import GraphBuilder |
| 13 | +from cuda.core.graph import GraphBuilder, GraphDefinition |
13 | 14 |
|
14 | 15 |
|
15 | 16 | def test_graph_is_building(init_cuda): |
@@ -384,3 +385,190 @@ def test_graph_stream_lifetime(init_cuda): |
384 | 385 |
|
385 | 386 | # Destroy the stream |
386 | 387 | stream.close() |
| 388 | + |
| 389 | + |
| 390 | +# --------------------------------------------------------------------------- |
| 391 | +# GraphBuilder.graph_definition |
| 392 | +# --------------------------------------------------------------------------- |
| 393 | + |
| 394 | + |
| 395 | +def test_graph_definition_returns_graph_definition_after_end_building(init_cuda): |
| 396 | + """Primary builder exposes its captured graph as a GraphDefinition after end_building().""" |
| 397 | + mod = compile_common_kernels() |
| 398 | + empty_kernel = mod.get_kernel("empty_kernel") |
| 399 | + |
| 400 | + gb = Device().create_graph_builder().begin_building() |
| 401 | + launch(gb, LaunchConfig(grid=1, block=1), empty_kernel) |
| 402 | + launch(gb, LaunchConfig(grid=1, block=1), empty_kernel) |
| 403 | + gb.end_building() |
| 404 | + |
| 405 | + gd = gb.graph_definition |
| 406 | + assert isinstance(gd, GraphDefinition) |
| 407 | + # The captured graph must contain the launched kernels. |
| 408 | + assert len(gd.nodes()) == 2 |
| 409 | + |
| 410 | + |
| 411 | +def test_graph_definition_raises_before_begin_building(init_cuda): |
| 412 | + """Primary builder has no graph allocated before begin_building().""" |
| 413 | + gb = Device().create_graph_builder() |
| 414 | + with pytest.raises(RuntimeError, match="before begin_building"): |
| 415 | + _ = gb.graph_definition |
| 416 | + |
| 417 | + |
| 418 | +def test_graph_definition_raises_during_capture(init_cuda): |
| 419 | + """graph_definition is unsafe while the driver is actively capturing.""" |
| 420 | + gb = Device().create_graph_builder().begin_building() |
| 421 | + try: |
| 422 | + with pytest.raises(RuntimeError, match="capture is in"): |
| 423 | + _ = gb.graph_definition |
| 424 | + finally: |
| 425 | + gb.end_building() |
| 426 | + |
| 427 | + |
| 428 | +def test_graph_definition_raises_for_forked(init_cuda): |
| 429 | + """Forked builders share the primary's graph; their property must raise.""" |
| 430 | + mod = compile_common_kernels() |
| 431 | + empty_kernel = mod.get_kernel("empty_kernel") |
| 432 | + |
| 433 | + gb = Device().create_graph_builder().begin_building() |
| 434 | + launch(gb, LaunchConfig(grid=1, block=1), empty_kernel) |
| 435 | + primary, sibling = gb.split(2) |
| 436 | + try: |
| 437 | + with pytest.raises(RuntimeError, match="forked"): |
| 438 | + _ = sibling.graph_definition |
| 439 | + finally: |
| 440 | + sibling = GraphBuilder.join(primary, sibling) |
| 441 | + sibling.end_building() |
| 442 | + |
| 443 | + |
| 444 | +def test_graph_definition_shares_ownership(init_cuda): |
| 445 | + """Closing the builder must not invalidate a held GraphDefinition.""" |
| 446 | + mod = compile_common_kernels() |
| 447 | + empty_kernel = mod.get_kernel("empty_kernel") |
| 448 | + |
| 449 | + gb = Device().create_graph_builder().begin_building() |
| 450 | + launch(gb, LaunchConfig(grid=1, block=1), empty_kernel) |
| 451 | + gb.end_building() |
| 452 | + |
| 453 | + gd = gb.graph_definition |
| 454 | + gb.close() |
| 455 | + # The shared CUgraph keeps the graph alive. |
| 456 | + assert len(gd.nodes()) == 1 |
| 457 | + |
| 458 | + |
| 459 | +def test_graph_definition_round_trips_through_explicit_api(init_cuda): |
| 460 | + """Mutating via the explicit API survives complete() and runs correctly.""" |
| 461 | + mod = compile_common_kernels() |
| 462 | + add_one = mod.get_kernel("add_one") |
| 463 | + |
| 464 | + launch_stream = Device().create_stream() |
| 465 | + mr = LegacyPinnedMemoryResource() |
| 466 | + b = mr.allocate(4) |
| 467 | + arr = np.from_dlpack(b).view(np.int32) |
| 468 | + arr[0] = 0 |
| 469 | + |
| 470 | + gb = launch_stream.create_graph_builder().begin_building() |
| 471 | + launch(gb, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data) |
| 472 | + gb.end_building() |
| 473 | + |
| 474 | + # Add a second add_one through the explicit GraphDefinition view. |
| 475 | + gd = gb.graph_definition |
| 476 | + captured_node = next(iter(gd.nodes())) |
| 477 | + captured_node.launch(LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data) |
| 478 | + assert len(gd.nodes()) == 2 |
| 479 | + |
| 480 | + graph = gb.complete() |
| 481 | + graph.launch(launch_stream) |
| 482 | + launch_stream.sync() |
| 483 | + assert arr[0] == 2 |
| 484 | + |
| 485 | + b.close() |
| 486 | + |
| 487 | + |
| 488 | +@requires_module(np, "2.1") |
| 489 | +def test_graph_definition_hybrid_conditional_body(init_cuda): |
| 490 | + """Populate a conditional body entirely through the explicit API. |
| 491 | +
|
| 492 | + This is the headline hybrid flow enabled by the new property: |
| 493 | + ``if_then`` returns a ``GraphBuilder`` for the body, but instead of |
| 494 | + calling ``begin_building`` and capturing into it, we reach for |
| 495 | + ``graph_definition`` and add nodes through the explicit API. |
| 496 | + """ |
| 497 | + mod = compile_conditional_kernels(int) |
| 498 | + add_one = mod.get_kernel("add_one") |
| 499 | + set_handle = mod.get_kernel("set_handle") |
| 500 | + |
| 501 | + launch_stream = Device().create_stream() |
| 502 | + mr = LegacyPinnedMemoryResource() |
| 503 | + b = mr.allocate(4) |
| 504 | + arr = np.from_dlpack(b).view(np.int32) |
| 505 | + arr[0] = 0 |
| 506 | + |
| 507 | + gb = Device().create_graph_builder().begin_building() |
| 508 | + condition = try_create_condition(gb) |
| 509 | + launch(gb, LaunchConfig(grid=1, block=1), set_handle, condition, 1) |
| 510 | + body_gb = gb.if_then(condition) |
| 511 | + |
| 512 | + # Skip body_gb.begin_building() entirely -- the body graph already |
| 513 | + # exists at conditional-node creation time and is exposed here. |
| 514 | + body_def = body_gb.graph_definition |
| 515 | + assert isinstance(body_def, GraphDefinition) |
| 516 | + assert len(body_def.nodes()) == 0 |
| 517 | + body_def.launch(LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data) |
| 518 | + |
| 519 | + graph = gb.end_building().complete() |
| 520 | + graph.launch(launch_stream) |
| 521 | + launch_stream.sync() |
| 522 | + assert arr[0] == 1 |
| 523 | + |
| 524 | + b.close() |
| 525 | + |
| 526 | + |
| 527 | +@requires_module(np, "2.1") |
| 528 | +def test_graph_definition_conditional_body_after_capture(init_cuda): |
| 529 | + """Capture into a conditional body, then augment it via the explicit API.""" |
| 530 | + mod = compile_conditional_kernels(int) |
| 531 | + add_one = mod.get_kernel("add_one") |
| 532 | + set_handle = mod.get_kernel("set_handle") |
| 533 | + |
| 534 | + launch_stream = Device().create_stream() |
| 535 | + mr = LegacyPinnedMemoryResource() |
| 536 | + b = mr.allocate(4) |
| 537 | + arr = np.from_dlpack(b).view(np.int32) |
| 538 | + arr[0] = 0 |
| 539 | + |
| 540 | + gb = Device().create_graph_builder().begin_building() |
| 541 | + condition = try_create_condition(gb) |
| 542 | + launch(gb, LaunchConfig(grid=1, block=1), set_handle, condition, 1) |
| 543 | + body_gb = gb.if_then(condition).begin_building() |
| 544 | + |
| 545 | + # Capture one increment into the body. |
| 546 | + launch(body_gb, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data) |
| 547 | + body_gb.end_building() |
| 548 | + |
| 549 | + # Add a second increment via the explicit API on the same body graph. |
| 550 | + body_def = body_gb.graph_definition |
| 551 | + captured_node = next(iter(body_def.nodes())) |
| 552 | + captured_node.launch(LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data) |
| 553 | + assert len(body_def.nodes()) == 2 |
| 554 | + |
| 555 | + graph = gb.end_building().complete() |
| 556 | + graph.launch(launch_stream) |
| 557 | + launch_stream.sync() |
| 558 | + assert arr[0] == 2 |
| 559 | + |
| 560 | + b.close() |
| 561 | + |
| 562 | + |
| 563 | +@requires_module(np, "2.1") |
| 564 | +def test_graph_definition_conditional_body_during_capture_raises(init_cuda): |
| 565 | + """The CAPTURING-state guard fires for conditional bodies too.""" |
| 566 | + gb = Device().create_graph_builder().begin_building() |
| 567 | + condition = try_create_condition(gb) |
| 568 | + body_gb = gb.if_then(condition).begin_building() |
| 569 | + try: |
| 570 | + with pytest.raises(RuntimeError, match="capture is in"): |
| 571 | + _ = body_gb.graph_definition |
| 572 | + finally: |
| 573 | + body_gb.end_building() |
| 574 | + gb.end_building() |
0 commit comments