For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

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, Deinitable, ImplicitlyCopyable, 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.

__deinit__

def __deinit__(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 max.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 {imm}:
        _ = 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[*Ts: DeviceGraphInput](ctx: DeviceContext, build: T, *inputs: *Ts.values, *, cache: Pointer[DeviceGraphCache]) -> Self

Builds and instantiates a device graph, reusing a cached one if it can.

Behaves like the uncached overload, except that a graph an earlier call built from equivalent inputs is returned as-is and build is never called.

Example:

from max.gpu.host import (
    DeviceContext, DeviceGraph, DeviceGraphBuilder, DeviceGraphCache
)

def kernel():
    print("replaying")

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

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

    # The second call reuses the graph the first one built.
    var graph = DeviceGraph.create(ctx, build, cache=Pointer(to=cache))
    var same = DeviceGraph.create(ctx, build, cache=Pointer(to=cache))
    ctx.synchronize()

Parameters:

• ​*Ts (DeviceGraphInput): Types of the device graph inputs.

Args:

• ​ctx (DeviceContext): Device context for the target device.
• ​build (T): Callback that adds nodes to the supplied builder, called only on a cache miss.
• ​*inputs (*Ts.values): The device graph inputs the cache key is derived from.
• ​cache (Pointer[DeviceGraphCache]): The cache to consult, and to store a newly built graph in.

Returns:

Self: The instantiated device graph, which may be one a previous call built.

Raises:

If graph builder creation, build, or instantiation fails.

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.

Pass a cache to the overload above to reuse a previously built graph instead of recording one on every call.

Example:

from max.gpu.host import DeviceContext, DeviceGraph, 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 {imm}:
        _ = 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.