Skip to content

Commit 0604b18

Browse files
cuda.core: add GraphBuilder.graph_definition property
Completes step 3 of #1330 by exposing the captured graph as an explicit `GraphDefinition` view that shares ownership of the underlying `CUgraph`. The handle-layer plumbing landed in PR #2008; this commit wires up the user-facing surface and locks in the state-guard rules. State semantics: - PRIMARY builder: only valid after `end_building()`. Before `begin_building()` no graph exists; during capture the driver is the sole writer, so explicit access is unsafe. - CONDITIONAL_BODY builder: valid both before `begin_building()` (the body graph is allocated at conditional-node creation time) and after `end_building()`. This enables a hybrid flow where a conditional body is populated entirely via the explicit API, with no capture at all. - FORKED builder: never valid. Forked builders share the primary's graph; access through the primary instead. Tests cover the happy path, both hybrid flows on conditional bodies (populate-via-explicit-API and capture-then-augment), the three error states (forked, capturing, primary pre-capture), and the shared-ownership guarantee (the `GraphDefinition` survives the builder's `close()`). Co-authored-by: Cursor <[email protected]>
1 parent dad6a42 commit 0604b18

3 files changed

Lines changed: 291 additions & 3 deletions

File tree

‎cuda_core/cuda/core/graph/_graph_builder.pyi‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,48 @@ class GraphBuilder:
129129
def is_join_required(self) -> bool:
130130
"""Returns True if this graph builder must be joined before building is ended."""
131131

132+
@property
133+
def graph_definition(self) -> GraphDefinition:
134+
"""The captured graph as an explicit :class:`~graph.GraphDefinition`.
135+
136+
The returned :class:`~graph.GraphDefinition` is a view of the same
137+
graph this builder is producing: nodes added through it appear in
138+
subsequent :meth:`complete` and :meth:`debug_dot_print` calls, and
139+
the view stays valid even after the builder is closed.
140+
141+
This lets you mix the capture and explicit APIs on a single graph,
142+
for example to inspect what was captured, augment it with extra
143+
nodes, or build a conditional body entirely with the explicit API.
144+
145+
Availability:
146+
147+
- **Primary builders** (created by :meth:`Device.create_graph_builder`
148+
or :meth:`Stream.create_graph_builder`): only after
149+
:meth:`end_building`.
150+
151+
- **Conditional-body builders** (returned by :meth:`if_then`,
152+
:meth:`if_else`, :meth:`while_loop`, :meth:`switch`): both before
153+
:meth:`begin_building` and after :meth:`end_building`. The body
154+
graph already exists when the conditional is created, so you may
155+
populate it through this view without ever calling
156+
:meth:`begin_building` on the body builder.
157+
158+
- **Forked builders** (returned by :meth:`split`): never. Forked
159+
builders share the primary builder's graph; access it through the
160+
primary instead.
161+
162+
Returns
163+
-------
164+
GraphDefinition
165+
A view of the graph being built.
166+
167+
Raises
168+
------
169+
RuntimeError
170+
If the builder is forked, currently building, or (for primary
171+
builders) has not started building yet.
172+
"""
173+
132174
def begin_building(self, mode: str | None='relaxed') -> GraphBuilder:
133175
"""Begins the building process.
134176

‎cuda_core/cuda/core/graph/_graph_builder.pyx‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ from libc.stdint cimport intptr_t
99

1010
from cuda.bindings cimport cydriver
1111

12-
from cuda.core.graph._graph_definition cimport GraphCondition
12+
from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition
1313
from cuda.core.graph._utils cimport _attach_host_callback_to_graph
1414
from cuda.core._resource_handles cimport (
1515
GraphHandle,
@@ -282,6 +282,64 @@ cdef class GraphBuilder:
282282
"""Returns True if this graph builder must be joined before building is ended."""
283283
return self._kind == FORKED
284284

285+
@property
286+
def graph_definition(self) -> GraphDefinition:
287+
"""The captured graph as an explicit :class:`~graph.GraphDefinition`.
288+
289+
The returned :class:`~graph.GraphDefinition` is a view of the same
290+
graph this builder is producing: nodes added through it appear in
291+
subsequent :meth:`complete` and :meth:`debug_dot_print` calls, and
292+
the view stays valid even after the builder is closed.
293+
294+
This lets you mix the capture and explicit APIs on a single graph,
295+
for example to inspect what was captured, augment it with extra
296+
nodes, or build a conditional body entirely with the explicit API.
297+
298+
Availability:
299+
300+
- **Primary builders** (created by :meth:`Device.create_graph_builder`
301+
or :meth:`Stream.create_graph_builder`): only after
302+
:meth:`end_building`.
303+
304+
- **Conditional-body builders** (returned by :meth:`if_then`,
305+
:meth:`if_else`, :meth:`while_loop`, :meth:`switch`): both before
306+
:meth:`begin_building` and after :meth:`end_building`. The body
307+
graph already exists when the conditional is created, so you may
308+
populate it through this view without ever calling
309+
:meth:`begin_building` on the body builder.
310+
311+
- **Forked builders** (returned by :meth:`split`): never. Forked
312+
builders share the primary builder's graph; access it through the
313+
primary instead.
314+
315+
Returns
316+
-------
317+
GraphDefinition
318+
A view of the graph being built.
319+
320+
Raises
321+
------
322+
RuntimeError
323+
If the builder is forked, currently building, or (for primary
324+
builders) has not started building yet.
325+
"""
326+
if self._kind == FORKED:
327+
raise RuntimeError(
328+
"graph_definition is unavailable on forked graph builders; "
329+
"access it through the primary builder instead."
330+
)
331+
if self._state == CAPTURING:
332+
raise RuntimeError(
333+
"graph_definition is unavailable while capture is in "
334+
"progress; call end_building() first."
335+
)
336+
if self._kind == PRIMARY and self._state == CAPTURE_NOT_STARTED:
337+
raise RuntimeError(
338+
"graph_definition is unavailable before begin_building() on "
339+
"a primary builder; no graph has been created yet."
340+
)
341+
return GraphDefinition._from_handle(self._h_graph)
342+
285343
def begin_building(self, mode: str | None = "relaxed") -> GraphBuilder:
286344
"""Begins the building process.
287345

‎cuda_core/tests/graph/test_graph_builder.py‎

Lines changed: 190 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@
55

66
import numpy as np
77
import pytest
8-
from helpers.graph_kernels import compile_common_kernels
8+
from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels
99
from helpers.marks import requires_module
10+
from helpers.misc import try_create_condition
1011

1112
from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch
12-
from cuda.core.graph import GraphBuilder
13+
from cuda.core.graph import GraphBuilder, GraphDefinition
1314

1415

1516
def test_graph_is_building(init_cuda):
@@ -384,3 +385,190 @@ def test_graph_stream_lifetime(init_cuda):
384385

385386
# Destroy the stream
386387
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

Comments
 (0)