Skip to main content

Mojo struct

OutputStage

struct OutputStage[opc: OutputPipelineConfig]

Acquired output stage with TMEM handle and pipeline reference.

Fields

  • index (UInt32):
  • tmem (OutputStage[opc].Tmem):
  • pipeline (OutputStage[opc].Pipeline):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

cta_group

comptime cta_group = opc.cta_group

num_stages

comptime num_stages = opc.num_stages

Pipeline

comptime Pipeline = ProducerConsumerPipeline[OutputStage[opc].num_stages]

stage_stride

comptime stage_stride = opc.stage_stride_cols

Tmem

comptime Tmem = TmemStage[opc]

Methods

__init__

__init__(index: UInt32, tmem: TmemStage[opc], pipeline: ProducerConsumerPipeline[OutputStage[opc].num_stages]) -> Self

from_raw

static from_raw(pipeline: ProducerConsumerPipeline[OutputStage[opc].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?