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

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, ImplicitlyDeletable, 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:

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.

View source