Mojo struct
OutputStage
@register_passable(trivial)
struct OutputStage[num_stages: Int, stage_stride: Int, cta_group: Int]
Acquired output stage with TMEM handle and pipeline reference.
Fields
- index (
UInt32): - tmem (
OutputStage[num_stages, stage_stride, cta_group].Tmem): - pipeline (
OutputStage[num_stages, stage_stride, cta_group].Pipeline):
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
Pipeline
comptime Pipeline = ProducerConsumerPipeline[num_stages]
Tmem
comptime Tmem = TmemStage[num_stages, stage_stride, cta_group]
Methods
__init__
__init__(index: UInt32, tmem: TmemStage[num_stages, stage_stride, cta_group], pipeline: ProducerConsumerPipeline[num_stages]) -> Self
from_raw
static from_raw(pipeline: ProducerConsumerPipeline[num_stages], stage_index: UInt32, tmem_offset: UInt32) -> Self
Create OutputStage from raw pipeline, stage index, and TMEM offset.
Useful when not using OutputTilePipeline's consumer() context manager.
Args:
- pipeline (
ProducerConsumerPipeline): The ProducerConsumerPipeline for barrier signaling. - stage_index (
UInt32): Current pipeline stage index. - tmem_offset (
UInt32): Pre-computed TMEM offset for this stage.
Returns:
Self: OutputStage with the given parameters.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!