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

Context

struct Context[params: ContextParams]

Per-block dispatch bundle. Comptime tier identity and widths flow through params; runtime per-block plumbing (split-K scratch + position) sits in var fields. The body forwards ctx to rowwise.reduce / rowwise.once and treats it as opaque, reading only the plainly-named members (ctx.axis, ctx.accumulator_width, ...). The _-prefixed members (ctx._tier, ctx._partials_base, ...) are scaffolder internals.

On CPU the runtime fields hold sentinels via Context.empty() and are never read.

Parameters

• ​params (ContextParams): Comptime tier parameters (axis, BLOCK_SIZE, tier discriminator, SIMD widths, target).

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

accumulator_width

comptime accumulator_width = Int(1) if (params > Int(1)) else params.simd_width

Reduce-op accumulator width a body should use for its M[dtype, accumulator_width] state: 1 on the tiled tier (each output lane is a separate accumulator), else the full simd_width. Centralized here so every reduction picks the same width.

axis

comptime axis = params.axis

Axis being reduced (mirrors Self.params.axis).

BLOCK_SIZE

comptime BLOCK_SIZE = params.BLOCK_SIZE

Threads per block (mirrors Self.params.BLOCK_SIZE).

emit_tile_width

comptime emit_tile_width = params.emit_tile_width

Rows-per-thread (mirrors Self.params.emit_tile_width).

simd_width

comptime simd_width = params.simd_width

SIMD width (mirrors Self.params.simd_width).

target

comptime target = params.target

Backend target (mirrors Self.params.target).

Methods

__init__

def __init__(partials_base: Pointer[UInt8, MutUntrackedOrigin], counters_base: Pointer[Int32, MutUntrackedOrigin], blocks_per_row: Int32, block_in_row: Int32, row_idx: Int32, is_last_block: Bool, phase: Int32 = Int32(0)) -> Self

Initializes a Context with the per-block runtime values.

Args:

• ​partials_base (Pointer[UInt8, MutUntrackedOrigin]): Split-K partials scratch (or sentinel).
• ​counters_base (Pointer[Int32, MutUntrackedOrigin]): Split-K per-row counters (or sentinel).
• ​blocks_per_row (Int32): Blocks cooperating on this row (or 0).
• ​block_in_row (Int32): This block's position in its row (or 0).
• ​row_idx (Int32): The row this block is reducing (or 0).
• ​is_last_block (Bool): Initialized to False; set by rowwise.reduce after the atomic finish (GPU).
• ​phase (Int32): Per-element-output split-K launch phase (or 0).

element_alignment

static def element_alignment[dtype: DType, ws: Int]() -> Int

Store/load alignment for a width-ws tile of dtype on this Context's target — element-natural on CPU, SIMD-natural on GPU.

A @staticmethod so it's usable in a comptime initializer: the receiver in ctx.element_alignment[dtype, ws]() is only a type carrier (runtime value never read), which comptime allows. Delegates to the module-level tile_alignment, the single source of truth shared with the public-op wrappers' loads.

Parameters:

• ​dtype (DType): The tile's element dtype.
• ​ws (Int): The tile's SIMD width.

Returns:

Int: The alignment in elements -- matches element_alignment in graph_compiler/extensibility/tensor_arg_traits.mojo and managed_tensor_slice.mojo, which multiply this by align_of[dtype]() to get the byte alignment passed to the backend. Named to match that convention explicitly, since a caller that instead hands this straight to a raw byte-level primitive (e.g. TileTensor.raw_load/raw_store, which wants bytes) without that multiplication under-claims alignment -- silently harmless for 4-byte-and-wider dtypes, a real miscompilation risk for narrower ones.

empty

static def empty() -> Self

Returns a Context with zeroed runtime fields. Used by CPU bodies and non-split-K GPU kernels.

Returns:

Self: A Context[params] with sentinel pointers and zero counters.