For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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
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.
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.
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.
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.
Returns:
wait_producer_acquire
def wait_producer_acquire(self, tile_idx: Int, stage: Int, phase: Int32)
Producer waits to acquire a tile.
signal_producer_release
def signal_producer_release(mut self, tile_idx: Int, stage: Int)
Producer signals it has released a tile.
wait_consumer_acquire
def wait_consumer_acquire(self, tile_idx: Int, stage: Int, phase: Int32)
Consumer waits to acquire a tile.
signal_consumer_release
def signal_consumer_release(mut self, tile_idx: Int, stage: Int)
Consumer signals it has released a tile.
get_producer_phase_increment
def get_producer_phase_increment(self) -> Int32
Get the phase increment for producers.
Returns:
get_consumer_phase_increment
def get_consumer_phase_increment(self) -> Int32
Get the phase increment for consumers.
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!