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 ofDeviceContextvalues 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:
- device_contexts (
InlineArray[DeviceContext, size]): The per-device contexts to store.
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): OneDeviceContextper device, exactlysizeof 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 theIndexertrait.
Args:
- idx (
I): A runtime index value that conforms to theIndexertrait.
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 returnedInlineArray.
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.