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

DeviceContextList

struct DeviceContextList[size: 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 InlineArray parameter.

Parameters

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

Fields

  • device_contexts (InlineArray[DeviceContext, size]): The underlying storage for the per-device contexts.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, Sized

Methods

__init__

def __init__(out self, device_contexts: InlineArray[DeviceContext, size])

Initialize from an InlineArray of DeviceContext values.

Args:

def __init__(out self, 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 DeviceContextList[size=N](ctx0, ctx1, ..., ctxN-1) directly from the per-device contexts attached to the kernel, so the wrapper avoids forcing callers to assemble an InlineArray themselves.

Args:

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

__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 size parameter.

filter_gpu_contexts

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

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

Some kernels receive a DeviceContextList 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 InlineArray.

Parameters:

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

Returns:

InlineArray[DeviceContext, num_gpu_devices]: An InlineArray 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.