Mojo struct
MmaStage
struct MmaStage[origin: MutOrigin, opc: OutputPipelineConfig]
Unified linear type handle for MMA stage in output pipeline.
Works as both a linear type (direct use) and within context managers.
Lifecycle:
- Created via
output_pipeline.acquire_mma_linear()- waits for epilogue - Use
tmem(),tmem_offset(),mbar()for MMA operations - Must call
release()to signal mma_arrive and advance (compiler-enforced)
Parameters
- origin (
MutOrigin): Origin of the pipeline reference. - opc (
OutputPipelineConfig): Output pipeline configuration (stages, stride, cta_group).
Fields
- pipeline_ptr (
Pointer[OutputTilePipeline[opc], origin]):
Implemented traits
comptime members
Stage
comptime Stage = OutputStage[opc]
TilePipelineType
comptime TilePipelineType = OutputTilePipeline[opc]
Methods
__init__
__init__(out self, pipeline_ptr: Pointer[OutputTilePipeline[opc], origin], stage: OutputStage[opc])
tmem
tmem(self) -> MmaStage[origin, opc].Stage.Tmem
Get the TMEM stage handle.
Returns:
MmaStage
tmem_offset
index
mbar
mbar(self) -> MbarPtr
Get the producer barrier for MMA commit.
Returns:
MbarPtr
release
release(deinit self)
Signal MMA completion and advance to next stage.
This is the only way to destroy this linear type. Internally calls mma_arrive (1-SM) or mma_arrive_multicast (2-SM).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!