Skip to main content

Mojo struct

BlockScaledInputProducer

@register_passable(trivial) struct BlockScaledInputProducer[origin: MutOrigin, a_type: DType, b_type: DType, sfa_type: DType, sfb_type: DType, a_tile_layout: Layout, b_tile_layout: Layout, sfa_tile_layout: Layout, sfb_tile_layout: Layout, num_pipeline_stages: Int, num_group_stages: Int, k_group_size: Int]

Producer view for TMA Load warp (block-scaled input pipeline).

Fields

  • pipeline_ptr (Pointer[BlockScaledInputProducer[origin, a_type, b_type, sfa_type, sfb_type, a_tile_layout, b_tile_layout, sfa_tile_layout, sfb_tile_layout, num_pipeline_stages, num_group_stages, k_group_size].TilePipelineType, origin]):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable

comptime members

__copyinit__is_trivial

comptime __copyinit__is_trivial = True

__del__is_trivial

comptime __del__is_trivial = True

__moveinit__is_trivial

comptime __moveinit__is_trivial = True

TilePipelineType

comptime TilePipelineType = BlockScaledTilePipeline[a_type, b_type, sfa_type, sfb_type, a_tile_layout, b_tile_layout, sfa_tile_layout, sfb_tile_layout, num_pipeline_stages, num_group_stages, k_group_size]

Methods

__enter__

__enter__(mut self) -> Self

__exit__

__exit__(mut self)

drain

drain(mut self)

Drain pipeline to prevent CTA exit while peer is still working.

acquire

acquire(mut self) -> BlockScaledProducerStage[origin, a_type, b_type, sfa_type, sfb_type, a_tile_layout, b_tile_layout, sfa_tile_layout, sfb_tile_layout, num_pipeline_stages, num_group_stages, k_group_size]

Acquire next stage, waiting for slot availability.

Returns:

BlockScaledProducerStage

Was this page helpful?