Skip to main content

/

Docs

Nightly

v25.1

Log in

Mojo struct

PipelineState

@register_passable(trivial) struct PipelineState[num_stages: Int]

Manages state for a multi-stage pipeline with circular buffer semantics.

PipelineState provides a mechanism for tracking the current stage in a multi-stage pipeline, particularly useful for double or triple buffering in GPU tensor operations. It maintains an index that cycles through the available stages, a phase bit that toggles when the index wraps around, and a monotonically increasing count.

This struct is commonly used with TMA operations to coordinate the use of multiple buffers in a pipeline fashion, allowing for overlapping computation and data transfer.

Parameters​

  • ​num_stages (Int): The number of stages in the pipeline (e.g., 2 for double buffering, 3 for triple buffering).

Implemented traits​

AnyType, Copyable, ExplicitlyCopyable, Movable, UnknownDestructibility

Methods​

__init__​

__init__() -> Self

Initialize a PipelineState with default values.

Creates a new PipelineState with index 0, phase 0, and count 0.

__init__(index: Int, phase: Int, count: Int) -> Self

Initialize a PipelineState with specific values.

Creates a new PipelineState with the specified index, phase, and count.

Args:

  • ​index (Int): The initial stage index.
  • ​phase (Int): The initial phase value (0 or 1).
  • ​count (Int): The initial count value.

index​

index(self) -> Int

Get the current stage index.

Returns:

The current index value, which ranges from 0 to num_stages-1.

phase​

phase(self) -> SIMD[uint32, 1]

Get the current phase bit.

Returns:

The current phase value (0 or 1), which toggles when the index wraps around.

step​

step(mut self)

Advance the pipeline state to the next stage.

Increments the index and count. When the index reaches num_stages, it wraps around to 0 and toggles the phase bit.

This function is used to move to the next buffer in a multi-buffer pipeline, implementing circular buffer semantics.