Mojo struct
ConsumerTile
@register_passable(trivial)
struct ConsumerTile[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, //, origin: MutOrigin, ring_buffer_type: AnyStruct[RingBuffer[consumer_warps, reads_per_warp_block]], warps_computed_per_consumer: Int, num_consumer_warps: Int]
Context manager for consumer access to a single ring buffer tile.
Fields
- consumer_view_ptr (
Pointer[RingBufferConsumer[origin, ring_buffer_type, warps_computed_per_consumer, num_consumer_warps], origin]): - tile (
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]): - stage (
Int): - warp_tile_idx (
Int):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
ConsumerViewPtrType
alias ConsumerViewPtrType = Pointer[RingBufferConsumer[origin, ring_buffer_type, warps_computed_per_consumer, num_consumer_warps], origin]
ConsumerViewType
alias ConsumerViewType = RingBufferConsumer[origin, ring_buffer_type, warps_computed_per_consumer, num_consumer_warps]
Methods
__init__
__init__(consumer_view_ptr: Pointer[RingBufferConsumer[origin, ring_buffer_type, warps_computed_per_consumer, num_consumer_warps], origin], stage: Int, local_tile_count: Int, warp_tile_idx: Int) -> Self
__enter__
__enter__(mut self) -> 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]
Return the acquired tile for use.
Returns:
__exit__
__exit__(mut self)
Release the tile back to producers.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!