Mojo struct
ConsumerView
@register_passable(trivial)
struct ConsumerView[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, //, origin: MutOrigin, ring_buffer_type: AnyStruct[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: Int]
Consumer view of the unified ring buffer.
Fields
- ring_buffer_ptr (
Pointer[ring_buffer_type, origin]): - phases (
StaticTuple[Int32, (pipeline_stages * warps_computed_per_consumer)]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
ConsumerTileType
comptime ConsumerTileType = ConsumerTile[origin, ring_buffer_type, warps_computed_per_consumer]
RingBufferPtrType
comptime RingBufferPtrType = Pointer[ring_buffer_type, origin]
Methods
__init__
__init__(ring_buffer_ptr: Pointer[ring_buffer_type, origin]) -> Self
__enter__
__enter__(mut self) -> Self
Context manager entry.
__exit__
__exit__(mut self)
Context manager exit.
acquire_tiles
acquire_tiles(mut self, stage: Int, consumer_iteration: Int, warp_tile_idx: Int) -> StaticTuple[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, block_rows, block_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](), block_rows, block_cols](), 128, 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]() if _tile_is_masked[pipeline_layout[layout, pipeline_stages](), block_rows, block_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, block_rows, block_cols]()[0], warp_rows, warp_cols](), alignment=128], tile_buffers]
Acquire tiles for reading by this consumer.
Args:
- stage (
Int): Pipeline stage to read from. - consumer_iteration (
Int): Which iteration this consumer is on (0 to warps_computed_per_consumer-1). - warp_tile_idx (
Int): Which tile this consumer wants to read.
Returns:
release_tiles
release_tiles(mut self, stage: Int, warp_tile_idx: Int)
Signal to producers that tile is free.
get_tile
get_tile(mut self, stage: Int, consumer_iteration: Int, warp_tile_idx: Int) -> ConsumerTile[origin, ring_buffer_type, warps_computed_per_consumer]
Get a context manager for accessing a tile.
Args:
- stage (
Int): Pipeline stage. - consumer_iteration (
Int): Current iteration of this consumer. - warp_tile_idx (
Int): Which tile to access.
Returns:
ConsumerTile
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!