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

DeviceContextArray

struct DeviceContextArray[length: Int]

A fixed-size collection of DeviceContext values.

Used by multi-device custom-op execute methods to receive one DeviceContext per participating device. The graph compiler recognizes this type and synthesizes it from the per-device contexts discovered on the operation, so kernels can index into it like a homogeneous array without the compiler having to introspect a generic Array parameter.

Parameters

• ​length (Int): The number of DeviceContext values in the collection.

Fields

• ​device_contexts (Array[DeviceContext, length]): The underlying storage for the per-device contexts.

Implemented traits

AnyType, Copyable, Deinitable, Movable, Sized

comptime members

size

comptime size = length

The number of DeviceContext values in the collection. Deprecated alias for length.

Deprecated: DeviceContextArray.size is deprecated, use DeviceContextArray.length.

Methods

__init__

def __init__(out self, var device_contexts: Array[DeviceContext, length])

Initialize from an Array of DeviceContext values.

Args:

• ​device_contexts (Array[DeviceContext, length]): The per-device contexts to store.

def __init__[*, __literal_size__: Int](out self: DeviceContextArray[__literal_size__], var *device_contexts: DeviceContext, *, __list_literal__: NoneType = None)

Initialize from a variadic sequence of DeviceContext values.

The graph compiler's multi-device lowering path uses this constructor: it synthesizes DeviceContextArray[length=N](ctx0, ctx1, ..., ctxN-1) directly from the per-device contexts attached to the kernel, so the wrapper avoids forcing callers to assemble an Array themselves.

Parameters:

• ​literal_size (Int): The number of contexts in the literal, inferred from the number of elements given.

Args:

• ​*device_contexts (DeviceContext): One DeviceContext per device, exactly length of them.
• ​list_literal (NoneType): Marker that lets this constructor accept list-literal syntax (var l: DeviceContextArray[N] = [c0, c1]).

Returns:

DeviceContextArray[__literal_size__]

__getitem__

def __getitem__[I: Indexer, //](self, idx: I) -> DeviceContext

Access a DeviceContext using a runtime index value.

Parameters:

• ​I (Indexer): A type that conforms to the Indexer trait.

Args:

• ​idx (I): A runtime index value that conforms to the Indexer trait.

Returns:

DeviceContext: The DeviceContext at the specified index.

__getitem_param__

def __getitem_param__[index: Int](self) -> DeviceContext

Access a DeviceContext at a compile-time known index.

Parameters:

• ​index (Int): A compile-time integer index.

Returns:

DeviceContext: The DeviceContext at the specified index.

__len__

def __len__(self) -> Int

Get the number of DeviceContext values in the collection.

Returns:

Int: The size of the collection as specified by the length parameter.

filter_gpu_contexts

def filter_gpu_contexts[num_gpu_devices: Int](self) -> Array[DeviceContext, num_gpu_devices]

Filters CPU contexts out and returns the GPU contexts in order.

Some kernels receive a DeviceContextArray that mixes GPU contexts with CPU contexts carrying host-side pointers. Most kernels only want the GPU contexts in launch order, packed into a fixed-size Array.

Parameters:

• ​num_gpu_devices (Int): The expected number of GPU contexts. Used as the size of the returned Array.

Returns:

Array[DeviceContext, num_gpu_devices]: An Array of size num_gpu_devices containing the GPU contexts in their original order.

Raises:

If the number of GPU contexts in the list is not equal to num_gpu_devices.