Mojo struct
RingBuffer
struct RingBuffer[dtype: DType, layout: Layout, pipeline_stages: Int, B_rows: Int, B_cols: Int, W_rows: Int, W_cols: Int, //, consumer_warps: Int, reads_per_warp_block: Int]
Manages access to shared memory tiles using barriers based in shared memory.
Fields
- barrier (
SMemArrayType[AMDSharedMemoryBarrier[consumer_warps], (RingBuffer[consumer_warps, reads_per_warp_block].block_warps * pipeline_stages)]): - smem_buffer (
SMemBuffer[dtype, layout, pipeline_stages, B_rows, B_cols, W_rows, W_cols]):
Implemented traits
AnyType,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = True
BarrierArray
alias BarrierArray = SMemArrayType[AMDSharedMemoryBarrier[consumer_warps], (RingBuffer[consumer_warps, reads_per_warp_block].block_warps * pipeline_stages)]
block_warps
alias block_warps = (B_rows // W_rows)
SmemBufferType
alias SmemBufferType = SMemBuffer[dtype, layout, pipeline_stages, B_rows, B_cols, W_rows, W_cols]
writes_per_warp_block
alias writes_per_warp_block = 1
Methods
__init__
__init__(out self, smem_buffer: SMemBuffer[dtype, layout, pipeline_stages, B_rows, B_cols, W_rows, W_cols])
get_tile
get_tile(mut self, mut phase: Int32, stage: Int, tile_idx: Int) -> LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, B_rows, B_cols]()[0], MutAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _tile_is_masked[pipeline_layout[layout, pipeline_stages](), B_rows, B_cols](), 128, W_rows, W_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](), B_rows, B_cols]() if _tile_is_masked[pipeline_layout[layout, pipeline_stages](), B_rows, B_cols]() else _tile_is_masked[LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, B_rows, B_cols]()[0], W_rows, W_cols](), alignment=128]
Returns:
release_tile
release_tile(mut self, stage: Int, tile_idx: Int)
producer
producer[warps_processed_per_producer: Int, producer_warps: Int](mut self) -> RingBufferProducer[self, RingBuffer[consumer_warps, reads_per_warp_block], warps_processed_per_producer, producer_warps]
Create a producer view of this ring buffer.
Returns:
RingBufferProducer
consumer
consumer[warps_computed_per_consumer: Int, num_consumer_warps: Int](mut self, consumer_warp_id: Int) -> RingBufferConsumer[self, RingBuffer[consumer_warps, reads_per_warp_block], warps_computed_per_consumer, num_consumer_warps]
Create a consumer view of this ring buffer.
Returns:
RingBufferConsumer
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!