IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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​

AnyType, ImplicitlyDeletable

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

Was this page helpful?