IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /docs/manual/basics.md). For the complete Mojo documentation index, see llms.txt.
Skip to main content
Version: Nightly
For the complete Mojo documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /docs/manual/basics.md).

DeviceGraph

struct DeviceGraph

Represents an instantiated device graph that can be replayed.

A DeviceGraph captures a sequence of GPU operations (such as kernel launches) as a reusable graph. Once instantiated from a DeviceGraphBuilder, the graph can be replayed multiple times at a lower overhead than re-enqueueing each operation individually.

To obtain a DeviceGraph, use DeviceGraph.create().

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

Methods

__init__

def __init__(out self, *, copy: Self)

Creates a copy of an existing device graph by incrementing its reference count.

Args:

  • copy (Self): The device graph to copy.

__del__

def __del__(deinit self)

Releases resources associated with this device graph.

replay

def replay(self)

Replays the captured sequence of GPU operations.

Submits the pre-captured sequence of operations for execution on the device. This is more efficient than re-enqueueing each operation individually because the graph has already been compiled and instantiated by the driver.

Example:

from std.gpu.host import DeviceContext, DeviceGraph, DeviceGraphBuilder

def kernel():
print("replaying")

with DeviceContext() as ctx:
var compiled_fn = ctx.compile_function[kernel]()

def build(mut builder: DeviceGraphBuilder) raises {read}:
_ = builder.add_function(
compiled_fn, grid_dim=1, block_dim=1, dependencies=[]
)

var graph = DeviceGraph.create(ctx, build)
graph.replay()
graph.replay() # replay as many times as needed
ctx.synchronize()

Raises:

If replay fails.

create

static def create(ctx: DeviceContext, build: T) -> Self

Builds and instantiates a device graph within a scoped callback.

Calls build with a fresh DeviceGraphBuilder, then instantiates the result into a replayable DeviceGraph. The builder, and any DeviceGraphNode handles obtained from it, are valid only for the duration of build: their origin is scoped to this call and cannot escape it, so a node handle cannot be stored beyond the callback or used with a different graph.

Example:

from std.gpu.host import DeviceContext, DeviceGraphBuilder

def kernel(x: Int):
print("Value:", x)

with DeviceContext() as ctx:
var compiled_fn = ctx.compile_function[kernel]()

def build(mut builder: DeviceGraphBuilder) raises {read}:
_ = builder.add_function(
compiled_fn, 42, grid_dim=1, block_dim=1, dependencies=[]
)

var graph = DeviceGraph.create(ctx, build)
graph.replay()
ctx.synchronize()

Args:

  • ctx (DeviceContext): Device context for the target device.
  • build (T): Callback that adds nodes to the supplied builder. It receives the builder by mutable reference and therefore cannot instantiate it directly; instantiation happens here once the callback returns.

Returns:

Self: The instantiated device graph.

Raises:

If graph builder creation, build, or instantiation fails.