Skip to main content

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, ImplicitlyDestructible

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]()[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]()[0], warp_rows, warp_cols](), alignment=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__​

__init__(out self)

get_tiles​

get_tiles(self, stage: Int, warp_tile_idx: Int) -> RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type].WarpTileTupleType

Get tiles from shared memory.

Returns:

RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type].WarpTileTupleType

producer​

producer[warps_processed_per_producer: Int](mut self) -> ProducerView[origin_of(self), RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type], warps_processed_per_producer]

Create a producer view of this ring buffer.

Returns:

ProducerView[origin_of(self), RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type], warps_processed_per_producer]

consumer​

consumer[warps_computed_per_consumer: Int](mut self) -> ConsumerView[origin_of(self), RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type], warps_computed_per_consumer]

Create a consumer view of this ring buffer.

Returns:

ConsumerView[origin_of(self), RingBuffer[dtype, layout, pipeline_stages, block_rows, block_cols, warp_rows, warp_cols, reads_per_warp_block, tile_buffers, sync_strategy_type], warps_computed_per_consumer]

get_staged_idx​

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

Get the staged index for a tile and stage.

Returns:

Int

wait_producer_acquire​

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

Producer waits to acquire a tile.

signal_producer_release​

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

Producer signals it has released a tile.

wait_consumer_acquire​

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

Consumer waits to acquire a tile.

signal_consumer_release​

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

Consumer signals it has released a tile.

get_producer_phase_increment​

get_producer_phase_increment(self) -> Int32

Get the phase increment for producers.

Returns:

Int32

get_consumer_phase_increment​

get_consumer_phase_increment(self) -> Int32

Get the phase increment for consumers.

Returns:

Int32

Models

Cookbook

Puzzles

Blog

Community

Contact