Mojo struct
EpilogueStage
struct EpilogueStage[_mlir_origin: LITMutOrigin, //, origin: MutOrigin, num_stages: Int, stage_stride_cols: Int, cta_group: Int]
Unified linear type handle for epilogue stage in output pipeline.
Works as both a linear type (direct use) and within context managers.
Lifecycle:
- Created via
output_pipeline.acquire_epilogue_linear()- waits for MMA - Use
tmem(),tmem_offset()for reading MMA results - Must call
release()to advance (compiler-enforced)
Parameters
- origin (
MutOrigin): Origin of the pipeline reference. - num_stages (
Int): Number of pipeline stages. - stage_stride_cols (
Int): TMEM column stride between stages. - cta_group (
Int): CTA group size (1 or 2).
Fields
- pipeline_ptr (
Pointer[EpilogueStage[origin, num_stages, stage_stride_cols, cta_group].TilePipelineType, origin]):
Implemented traits
comptime members
Stage
comptime Stage = OutputStage[num_stages, stage_stride_cols, cta_group]
TilePipelineType
comptime TilePipelineType = OutputTilePipeline[num_stages, stage_stride_cols, cta_group]
Methods
__init__
__init__(out self, pipeline_ptr: Pointer[EpilogueStage[origin, num_stages, stage_stride_cols, cta_group].TilePipelineType, origin], stage: OutputStage[num_stages, stage_stride_cols, cta_group])
tmem
tmem(self) -> EpilogueStage[origin, num_stages, stage_stride_cols, cta_group].Stage.Tmem
Get the TMEM stage handle.
Returns:
EpilogueStage
tmem_offset
index
release
release(deinit self)
Free stage for MMA reuse and advance to next stage.
This is the only way to destroy this linear type.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!