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.