triton.experimental.gluon.language.amd.warp_pipeline_stage

class triton.experimental.gluon.language.amd.warp_pipeline_stage(self, label=None, *, priority: int | None = None, **_internal)

Marks a warp-pipeline stage inside a Gluon kernel.

Within a @gluon.jit kernel, each with gl.amd.warp_pipeline_stage(...) block defines a distinct stage of a warp pipeline. All operations inside the block belong to the same pipeline cluster and execute as a unit relative to other stages.

The optional string label (for example, "load" or "compute") identifies the stage in diagnostics without affecting program semantics.

priority is an optional integer that expresses the scheduling priority of the warp that executes the stage. The priority applies to the entire cluster. Valid values range from 0 (lowest) to 3 (highest), matching the operand range of s_setprio. If omitted, the priority resets to zero when another stage in the loop uses an explicit priority; otherwise, no priority instruction is emitted. N.B., This is a performance hint to the hardware scheduler. Its effect depends on the dynamic interaction between warp instruction streams across different warps. It is optional and should be used judiciously, only when explicit scheduling guidance is beneficial.

The following schematic shows the stage boundaries; it is not intended as an optimal kernel.

Example

from triton.experimental import gluon
from triton.experimental.gluon import language as gl

@gluon.jit
def warp_pipelined_kernel(a_ptr, b_ptr, c_ptr, K: gl.constexpr):
    acc = 0.0

    for k in range(0, K):
        # Stage 0: prefetch tiles.
        with gl.amd.warp_pipeline_stage("load", priority=3):
            a = gl.load(a_ptr + k)
            b = gl.load(b_ptr + k)

        # Stage 1: prepare MFMA operands.
        with gl.amd.warp_pipeline_stage("prep"):
            a_tile = a  # Convert to the required dot-operand layout.
            b_tile = b  # Convert to the required dot-operand layout.

        # Stage 2: compute.
        with gl.amd.warp_pipeline_stage("compute", priority=0):
            acc += a_tile * b_tile

    gl.store(c_ptr, acc)

__init__(self, label=None, *, priority: int | None = None, **_internal)

Methods

__init__(self[, label, priority])

Attributes

label
priority