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

RingBuffer

struct RingBuffer[dtype: DType, layout: Layout, pipeline_stages: Int, block_rows: Int, block_cols: Int, warp_rows: Int, warp_cols: Int, reads_per_warp_block: Int, tile_buffers: Int, sync_strategy_type: SyncStrategy]

Ring buffer for coordinating producer-consumer warps in matrix multiplication.

Parameters

• ​dtype (DType): Data type of elements.
• ​layout (Layout): Memory layout for shared memory tiles.
• ​pipeline_stages (Int): Number of stages for software pipelining.
• ​block_rows (Int): Number of rows in block-level tiles.
• ​block_cols (Int): Number of columns in block-level tiles.
• ​warp_rows (Int): Number of rows in warp-level tiles.
• ​warp_cols (Int): Number of columns in warp-level tiles.
• ​reads_per_warp_block (Int): How many consumer warps read each tile.
• ​tile_buffers (Int): Number of separate tile buffers (usually 1).
• ​sync_strategy_type (SyncStrategy): Synchronization strategy (SingleCounterSync or SplitCounterSync).

Fields

• ​smem_buffers (RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type].SMemBuffersType):
• ​sync_strategy (sync_strategy_type):

Implemented traits

AnyType, Deinitable, Movable

comptime members

block_warps

comptime block_warps = (block_rows // warp_rows)

SMemBuffersType

comptime SMemBuffersType = StaticTuple[SMemBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols], tile_buffers]

SmemBufferType

comptime SmemBufferType = SMemBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols]

total_tiles

comptime total_tiles = (RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type].block_warps * pipeline_stages)

WarpTileTupleType

comptime WarpTileTupleType = StaticTuple[LayoutTensor[dtype, LayoutTensor._compute_tile_layout[warp_rows, warp_cols]()[Int(0)], MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), linear_idx_type=_get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), masked=_tile_is_masked[pipeline_layout[layout, pipeline_stages](), block_rows, block_cols]() or _tile_is_masked[LayoutTensor._compute_tile_layout[block_rows, block_cols]()[Int(0)], warp_rows, warp_cols](), alignment=Int(128)], tile_buffers]

WarpTileType

comptime WarpTileType = RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type].SmemBufferType.WarpTileType

Methods

__init__

def __init__(out self)

get_tiles

def get_tiles(self, stage: Int, warp_tile_idx: Int) -> Self.WarpTileTupleType

Get tiles from shared memory.

Args:

• ​stage (Int): Pipeline stage to read tiles from.
• ​warp_tile_idx (Int): Index of the warp tile within the stage to read.

Returns:

Self.WarpTileTupleType

producer

def producer[warps_processed_per_producer: Int](mut self) -> ProducerView[origin_of(self), Self, warps_processed_per_producer]

Create a producer view of this ring buffer.

Parameters:

• ​warps_processed_per_producer (Int): Number of warps each producer processes per pipeline stage.

Returns:

ProducerView[origin_of(self), Self, warps_processed_per_producer]

consumer

def consumer[warps_computed_per_consumer: Int](mut self) -> ConsumerView[origin_of(self), Self, warps_computed_per_consumer]

Create a consumer view of this ring buffer.

Parameters:

• ​warps_computed_per_consumer (Int): Number of warps each consumer computes per pipeline stage.

Returns:

ConsumerView[origin_of(self), Self, warps_computed_per_consumer]

get_staged_idx

def get_staged_idx(self, tile_idx: Int, stage: Int) -> Int

Get the staged index for a tile and stage.

Args:

• ​tile_idx (Int): Index of the warp tile within the block.
• ​stage (Int): Pipeline stage of the tile.

Returns:

Int

wait_producer_acquire

def wait_producer_acquire(self, tile_idx: Int, stage: Int, phase: Int32)

Producer waits to acquire a tile.

Args:

• ​tile_idx (Int): Index of the warp tile within the block to acquire.
• ​stage (Int): Pipeline stage of the tile to acquire.
• ​phase (Int32): Synchronization phase counter for this tile and stage.

signal_producer_release

def signal_producer_release(mut self, tile_idx: Int, stage: Int)

Producer signals it has released a tile.

Args:

• ​tile_idx (Int): Index of the warp tile within the block to release.
• ​stage (Int): Pipeline stage of the tile to release.

wait_consumer_acquire

def wait_consumer_acquire(self, tile_idx: Int, stage: Int, phase: Int32)

Consumer waits to acquire a tile.

Args:

• ​tile_idx (Int): Index of the warp tile within the block to acquire.
• ​stage (Int): Pipeline stage of the tile to acquire.
• ​phase (Int32): Synchronization phase counter for this tile and stage.

signal_consumer_release

def signal_consumer_release(mut self, tile_idx: Int, stage: Int)

Consumer signals it has released a tile.

Args:

• ​tile_idx (Int): Index of the warp tile within the block to release.
• ​stage (Int): Pipeline stage of the tile to release.

get_producer_phase_increment

def get_producer_phase_increment(self) -> Int32

Get the phase increment for producers.

Returns:

Int32

get_consumer_phase_increment

def get_consumer_phase_increment(self) -> Int32

Get the phase increment for consumers.

Returns:

Int32