Skip to main content

Mojo struct

ProducerView

@register_passable(trivial) struct ProducerView[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_processed_per_producer: Int]

Producer view of the unified ring buffer.

Fields

  • ring_buffer_ptr (Pointer[ring_buffer_type, origin]):
  • phases (StaticTuple[Int32, (pipeline_stages * warps_processed_per_producer)]):

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

ProducerTileType

comptime ProducerTileType = ProducerTile[origin, ring_buffer_type, warps_processed_per_producer]

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, producer_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 writing by this producer.

Args:

  • stage (Int): Pipeline stage to write to.
  • producer_iteration (Int): Which iteration this producer is on (0 to warps_processed_per_producer - 1).
  • warp_tile_idx (Int): Which tile this producer is responsible for.

Returns:

StaticTuple

release_tiles

release_tiles(mut self, stage: Int, warp_tile_idx: Int)

Signal to consumers that tile is ready.

get_tile

get_tile(mut self, stage: Int, warp_tile_idx: Int, producer_iteration: Int) -> ProducerTile[origin, ring_buffer_type, warps_processed_per_producer]

Get a context manager for accessing a tile.

Args:

  • stage (Int): Pipeline stage.
  • warp_tile_idx (Int): Which tile to access.
  • producer_iteration (Int): Current iteration of this producer.

Returns:

ProducerTile

Was this page helpful?