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).

DeviceGraphBuilder

struct DeviceGraphBuilder[arena_origin: ImmutOrigin]

Builder for explicit device graph construction.

A DeviceGraphBuilder is handed to the callback passed to DeviceGraph.create(). Callers add kernel nodes via add_function() from within that callback, which then instantiates a reusable DeviceGraph.

The builder, and any DeviceGraphNode handles it produces, are valid only for the duration of the callback: their origin is scoped to the DeviceGraph.create call and cannot escape it.

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()

Parameters​

  • ​arena_origin (ImmutOrigin): Origin of the enclosing DeviceGraph.create scope.

Implemented traits​

AnyType, ImplicitlyDeletable, Movable

comptime members​

Node​

comptime Node = DeviceGraphNode[arena_origin]

Node handle type produced by this builder, branded with the builder's DeviceGraph.create scope origin.

Methods​

__del__​

def __del__(deinit self)

Releases resources associated with this graph builder.

context​

def context(self) -> DeviceContext

Returns the device context this builder records against.

Unlike the context() accessors on buffer types, this is a non-raising read of the builder's stored device context (the def declares no raises).

Returns:

DeviceContext: The DeviceContext backing this builder.

add_function​

def add_function[*Ts: DevicePassable](self, f: DeviceFunction[target=f.target, compile_options=f.compile_options, link_options=f.link_options, _ptxas_info_verbose=f._ptxas_info_verbose], *args: *Ts.values, *, grid_dim: Dim, block_dim: Dim, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None)), cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds a type-checked compiled kernel function as a node in this graph.

Parameters:

Args:

Returns:

Self.Node: A handle to the newly added kernel-dispatch node.

Raises:

If adding the node fails.

def add_function[FuncType: def() -> None, //, dump_asm: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, func: FuncType, grid_dim: Dim, block_dim: Dim, *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None)), cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None))) -> Self.Node

Compiles and adds a capturing kernel closure as a node in this graph.

This overload is for kernels that capture variables from their enclosing scope using the {var} capture syntax. Compilation is performed automatically using the DeviceContext that created this builder, so no separate compile step is needed.

Example:

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

with DeviceContext() as ctx:
    var scale: Float32 = 2.0
    var buf = ctx.enqueue_create_buffer[DType.float32](256)
    var ptr = buf.unsafe_ptr()

    def scale_kernel() {var}:
        var i = global_idx.x
        ptr[i] = Float32(i) * scale

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

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

Parameters:

Args:

  • ​func (FuncType): The capturing kernel closure to compile and add as a graph node.
  • ​grid_dim (Dim): Dimensions of the compute grid.
  • ​block_dim (Dim): Dimensions of each thread block.
  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Explicit list of predecessor node handles. An empty list makes the new node a graph root with no predecessors; a non-empty list uses those exact handles as predecessors.
  • ​cluster_dim (OptionalReg[Dim]): Cluster dimensions (optional).
  • ​shared_mem_bytes (OptionalReg[Int]): Amount of dynamic shared memory per block.
  • ​attributes (List[LaunchAttribute]): Launch attributes.
  • ​constant_memory (List[ConstantMemoryMapping]): Constant memory mappings.

Returns:

Self.Node: A handle to the newly added kernel-dispatch node.

Raises:

If adding the node fails.

def add_function[declared_arg_types: TypeList[declared_arg_types.values], //, func: def(*args: *declared_arg_types) thin -> None, *actual_arg_types: DevicePassable, *, link_options: StringSlice[StaticConstantOrigin] = StringSlice(""), dump_asm: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[StaticConstantOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, *args: *actual_arg_types.values, *, grid_dim: Dim, block_dim: Dim, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None)), cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None)), func_attribute: OptionalReg[FuncAttribute] = None) -> Self.Node

Compiles and adds a kernel function as a node in this graph.

This overload takes the kernel as a compile-time parameter and compiles it automatically using the DeviceContext that created this builder, so no separate DeviceContext.compile_function() step is needed. It mirrors the parameter-based DeviceContext.enqueue_function() overload for the non-graph path.

You can pass the function directly to add_function without compiling it first:

from std.gpu.host import DeviceContext, DeviceGraphBuilder

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

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

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

Parameters:

Args:

  • ​*args (*actual_arg_types.values): Variadic arguments which are passed to the func.
  • ​grid_dim (Dim): Dimensions of the compute grid.
  • ​block_dim (Dim): Dimensions of each thread block.
  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Explicit list of predecessor node handles. An empty list makes the new node a graph root with no predecessors; a non-empty list uses those exact handles as predecessors.
  • ​cluster_dim (OptionalReg[Dim]): Cluster dimensions (optional).
  • ​shared_mem_bytes (OptionalReg[Int]): Amount of dynamic shared memory per block.
  • ​attributes (List[LaunchAttribute]): Launch attributes.
  • ​constant_memory (List[ConstantMemoryMapping]): Constant memory mappings.
  • ​func_attribute (OptionalReg[FuncAttribute]): CUfunction_attribute enum.

Returns:

Self.Node: A handle to the newly added kernel-dispatch node.

Raises:

If adding the node fails.

add_copy​

def add_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_buf: HostBuffer[dtype], *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds a host-to-device memcpy node to the graph.

The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Returns:

Self.Node: A handle to the newly added memcpy node.

Raises:

If adding the node fails.

def add_copy[dtype: DType](self, dst_buf: HostBuffer[dtype], src_buf: DeviceBuffer[dtype], *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds a device-to-host memcpy node to the graph.

The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Returns:

Self.Node: A handle to the newly added memcpy node.

Raises:

If adding the node fails.

def add_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_buf: DeviceBuffer[dtype], *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds a device-to-device memcpy node to the graph.

Both buffers must belong to the same context as this builder; cross-context copies are not supported in graphs. The number of bytes copied is determined by the size of the source buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

  • ​dst_buf (DeviceBuffer[dtype]): Device buffer to copy to.
  • ​src_buf (DeviceBuffer[dtype]): Device buffer to copy from. Must be the same size as dst_buf.
  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Explicit list of predecessor node handles. An empty list makes the new node a graph root with no predecessors; a non-empty list uses those exact handles as predecessors.

Returns:

Self.Node: A handle to the newly added memcpy node.

Raises:

If adding the node fails.

add_memset​

def add_memset[dtype: DType](self, dst: DeviceBuffer[dtype], val: Scalar[dtype], *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds a memset node to the graph that sets all elements of dst to val.

Parameters:

  • ​dtype (DType): Type of the data stored in the buffer.

Args:

  • ​dst (DeviceBuffer[dtype]): Destination buffer.
  • ​val (Scalar[dtype]): Value to set all elements of dst to.
  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Explicit list of predecessor node handles. An empty list makes the new node a graph root with no predecessors; a non-empty list uses those exact handles as predecessors.

Returns:

Self.Node: A handle to the newly added memset node.

Raises:

If adding the node fails. The underlying graph APIs cannot express an 8-byte memset whose high and low 32-bit halves differ as a single node, so such patterns will return an error.

add_empty​

def add_empty(self, *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Adds an empty (no-op) node to the graph.

Empty nodes perform no work at execution time. They are used purely for transitive ordering: a single empty node fanned in from m predecessors and out to n successors expresses an m-to-n barrier using m + n edges instead of m * n, and serves as a stable handle for "the completion of this phase" when the producer set is not visible to the consumer.

Args:

  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Explicit list of predecessor node handles. An empty list makes the new node a graph root with no predecessors; a non-empty list uses those exact handles as predecessors.

Returns:

Self.Node: A handle to the newly added empty node.

Raises:

If adding the node fails.

region​

def region(mut self, work: T, *, var dependencies: List[DeviceGraphNode[arena_origin]] = List(__list_literal__=NoneType(None))) -> Self.Node

Runs work and returns a single empty node that joins every node added to this builder during its execution.

The returned handle is suitable for use as a one-element dependencies= entry on a downstream add_* call. The empty node performs no work at execution time; it exists purely as a fan-in barrier so the caller does not need to thread the producer set's individual handles to every consumer.

Every node work adds also depends on the predecessors named in dependencies: while work runs, those handles are injected as ambient predecessors that each add_* call unions into its own dependencies. This makes the region's nodes run after the named predecessors without the closure having to thread the handles through to every add_* call. With the default (empty) dependencies, the region's nodes are unconstrained relative to earlier work.

Example:

from std.gpu.host import DeviceContext, DeviceGraphBuilder

with DeviceContext() as ctx:
    var buf_a = ctx.enqueue_create_buffer[DType.uint8](100)
    var buf_b = ctx.enqueue_create_buffer[DType.uint8](100)
    var buf_c = ctx.enqueue_create_buffer[DType.uint8](100)
    var host_src = ctx.enqueue_create_host_buffer[DType.uint8](100)

    def build(mut builder: DeviceGraphBuilder) raises {read}:
        def add_producers(mut b: DeviceGraphBuilder) raises {read} -> None:
            _ = b.add_memset(buf_a, UInt8(1), dependencies=[])
            _ = b.add_memset(buf_b, UInt8(2), dependencies=[])

        var producers_join = builder.region(add_producers)
        _ = builder.add_copy(
            buf_c, host_src, dependencies=[producers_join]
        )

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

Args:

  • ​work (T): Closure whose effects on this builder are captured. The builder is passed as work's sole argument; the closure must not capture the same builder, since doing so would alias with this method's receiver. The closure may add any number of nodes (zero or more) via any of the add_* methods.
  • ​dependencies (List[DeviceGraphNode[arena_origin]]): Predecessor node handles that every node added by work should depend on. Defaults to empty (no added predecessors).

Returns:

Self.Node: A handle that successors can depend on to run after everything work added. When work adds two or more nodes, this is a fresh empty node that joins them; when it adds exactly one node, that node is returned directly (no extra empty node); when it adds none, the returned empty node falls back to depending on dependencies so it still chains correctly.

Raises:

Anything work itself raises, or anything raised while adding the join node.

add_output​

def add_output(self, var output: AnyAsyncValueRef)

Add a value as an output for the resulting device graph.

The graph records the output so its backing memory outlives the graph that references it. Ownership of the async value is transferred to the builder.

Args:

num_outputs​

def num_outputs(self) -> Int

Returns the number of outputs registered on the device graph.

Returns:

Int: The number of outputs added via add_output.

View source