Wave Scheduling Pass

Scheduling Types

Wave offers 4 automated scheduling and 1 manual scheduling types:

Value	Category	Description
NONE	—	No scheduling; operations execute in program order
MODULO	Solver-based	Classical modulo scheduling using a resource reservation table and ILP-style heuristics
PREFETCH	Heuristic	Automatic 2-stage software pipeline: global load → shared store → shared load → compute
FOUR_STAGE	Heuristic	Extended 4-stage variant of the prefetch pipeline
PREFETCH_ATTENTION	Heuristic	Attention-kernel-specific variant of prefetch scheduling
MANUAL	Manual	User-defined schedule via the scheduling DSL

Scheduling Parameters attached to nodes

absolute_cycle: int: Absolute clock cycle assignment
cycle: int: Cycle within one initiation interval (absolute_cycle % Initiation Interval)
stage: int: Pipeline stage (absolute_cycle // Initiation Interval)
initiation_interval: int: Initiation Interval for the pipelined loop
prefetch_stage: int | None: For prefetch-style schedules

Manual Scheduling DSL

Node Selection

Operation	Description
`tkw.get_node_by_tag(tag)`	Returns all nodes whose tag attribute matches (string or set of strings)
`tkw.get_node_by_tag_and_type(tag, node_type)`	Returns nodes matching both tag and Python type (e.g., tkw.Read, tkw.Write)
`tkw.get_node_by_tag_and_iteration(tag, iteration)`	After unrolling, returns nodes from a specific unrolled iteration copy
`tkw.filter_nodes(nodes, subgraph=None, node_type=None)`	Filters a node list by pipeline stage (PROLOGUE, KERNEL, EPILOGUE) and/or type

Partitioning

Operation	Description
`tkw.partition_by_address_space(nodes, address_space)`	Splits reads/writes into those matching a given address space and those that do not; returns (matched, unmatched)
`tkw.partition_by_dim(nodes, dim, num_partitions)`	Splits nodes evenly by a loop dimension (e.g., K) into N groups; returns a tuple of lists

Clustering and Reordering

Operation	Description
`tkw.cluster(ops)`	Groups a sequence of operations (including inline sync primitives) into an ordered execution cluster
`tkw.reorder_graph(subgraph_ref, clusters)`	Reorders the specified stage (KERNEL, PROLOGUE, or EPILOGUE) according to the provided cluster ordering

Pipeline Construction

with tkw.pipeline(k_loop) as pipelined_loop:
    pipelined_loop.set_stage([...])  # Stage 0
    pipelined_loop.set_stage([...])  # Stage 1
    # ...

Loop Unrolling

tkw.unroll(loop_node, factor=N)

Insertion Operations

Operation	Description
`tkw.insert_before(target, op)`	Inserts op immediately before target
`tkw.insert_after(target, op)`	Inserts op immediately after target
`tkw.insert_at_start(subgraph_ref, op)`	Inserts op after all placeholder nodes at the start of a subgraph
`tkw.insert_at_end(subgraph_ref, op)`	Inserts op before the output node at the end of a subgraph

Barrier and Synchronization Operations

Operation	Description
`tkw.insert_cond_barrier_before(condition, target)`	Inserts a conditional workgroup barrier before target; condition is a sympy expression or node (e.g., hw.wave_id >= mid_wave)
`tkw.insert_cond_barrier_after(condition, target)`	Inserts a conditional workgroup barrier after target
`tkw.stagger(loop)`	Simple 2-way wave stagger: automatically computes wave conditions from HardwareConstraint and places barriers to divide waves into groups

Synchronization Primitives (Inline Objects)

Object	Description
`tkw.SharedMemoryBarrierSignal(count, ds_wait=False)`	Signal a shared memory barrier; ds_wait=True also inserts a DS wait
`tkw.SharedMemoryBarrierWait(count)`	Wait on a shared memory barrier
`tkw.TensorCounterWait(count)`	Wait on a tensor counter (rocdl.s.wait.tensorcnt)
`tkw.SchedulingBarrier(mask)`	Insert a scheduling barrier (rocdl.sched.barrier) with the given instruction-class mask
`tkw.SetWavePrio(priority)`	Set wave execution priority (0 = low, 1 = high)

Utility

Operation	Description
`tkw.getitem(obj, index)`	Extracts element index from a list or tuple returned by a DSL operation
`tkw.interleave_ops(base_ops, interleaved_ops, intervals, start_offsets, start_after_groups)`	Interleaves groups of operations into a base sequence at specified intervals

Mapping to MLIR Transform Dialect

Re-useable Ops

Upstream op	Wave DSL equivalent
`transform.structured.match ops{["wave.read"]}`	`get_node_by_tag_and_type(..., tkw.Read)` (partial — selects by type)
`transform.structured.match ops{["wave.iterate"]}`	`get_node_by_tag("k_loop")` (partial — selects by type)
`transform.annotate`	Attaching stage/II/cluster annotations to ops
`transform.foreach`	Iterating over a handle set
`transform.get_producer_of_operand`	Graph navigation for dependency analysis
`transform.get_consumers_of_result`	Graph navigation
`transform.loop.unroll`	`tkw.unroll(loop, factor=N)` (post-lowering, on scf.for)
`transform.loop.pipeline`	Automatic pipelining of `scf.for` (post-lowering, optional path)
`transform.apply_registered_pass`	Invoking water passes as part of a transform sequence
`transform.cast`	Handle type narrowing

New Ops

New op	Why no upstream equivalent
`wave.transform.get_by_tag`	transform.structured.match selects by op name only; Wave tags are string attributes on ops, not op names
`wave.transform.filter_by_address_space`	GPU address spaces are Wave/AMDGPU-specific; no upstream notion of this
`wave.transform.partition_by_dim`	Wave dimension symbols are specific to the Wave type system
`wave.transform.set_pipeline_stage`	Attaches Wave-specific pipelining annotations (II, stage index) that LowerWaveToMLIR reads; `transform.annotate` is untyped and insufficient alone
`wave.transform.insert_barrier`	Inserts `wave.shared_memory_barrier` / `wave.scheduling_barrier` ops; GPU-hardware-specific`
`wave.transform.cluster`	Ordered instruction cluster — encodes parallel/sequential execution intent within a stage; no upstream equivalent

Out of Scope

tkw.filter_nodes
tkw.reorder_graph
tkw.insert_before(target, op)
tkw.insert_after(target, op)
tkw.insert_at_start(subgraph_ref, op)
tkw.insert_at_end(subgraph_ref, op)
tkw.insert_cond_barrier_before(condition, target)
tkw.insert_cond_barrier_after(condition, target)
tkw.stagger(loop)
tkw.getitem(obj, index)
tkw.interleave_ops(base_ops, interleaved_ops, intervals, start_offsets, start_after_groups)

Python DSL -> MLIR

Existing Python DSL example from docs/wave/wave_schedule.rst

@wave_schedule.wave_schedule()
def prefetch_schedule():
    k_loop = tkw.get_node_by_tag("k_loop")
    load_a = tkw.get_node_by_tag_and_type("read_a", tkw.Read)
    global_load_a, shared_load_a = tkw.partition_by_address_space(load_a, GLOBAL_ADDRESS_SPACE)
    shared_write_a = tkw.get_node_by_tag_and_type("read_a", tkw.Write)
    mma = tkw.get_node_by_tag("mma")
    with tkw.pipeline(k_loop) as pl:
        pl.set_stage([(global_load_a,), (shared_write_a,)])
        pl.set_stage([(shared_load_a,), (mma,)])

Equivalent Transform Dialect sequence:

module attributes {transform.with_named_sequence} {
  transform.named_sequence @wave_schedule(%module: !transform.any_op) {

    // Select k_loop by tag  [NEW op]
    %k_loop = wave.transform.get_by_tag %module { tag = "k_loop" }
        : (!transform.any_op) -> !transform.op<"wave.iterate">

    // Select read_a read ops by tag and type  [NEW op]
    %read_a = wave.transform.get_by_tag %module { tag = "read_a", op_type = "wave.read" }
        : (!transform.any_op) -> !transform.any_op

    // Partition reads by address space  [NEW op]
    %global_load_a, %shared_load_a = wave.transform.filter_by_address_space %read_a
        { address_space = 1 }
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op)

    // Select the write generated for read_a  [NEW op]
    %shared_write_a = wave.transform.get_by_tag %module { tag = "read_a", op_type = "wave.write" }
        : (!transform.any_op) -> !transform.any_op

    // Select mma  [NEW op]
    %mma = wave.transform.get_by_tag %module { tag = "mma" }
        : (!transform.any_op) -> !transform.any_op

    // Stage 0: global load + shared write  [NEW op]
    wave.transform.set_pipeline_stage %k_loop, %global_load_a
        { stage = 0, cluster = 0, initiation_interval = 2 }
        : (!transform.op<"wave.iterate">, !transform.any_op) -> ()
    wave.transform.set_pipeline_stage %k_loop, %shared_write_a
        { stage = 0, cluster = 1, initiation_interval = 2 }
        : (!transform.op<"wave.iterate">, !transform.any_op) -> ()

    // Stage 1: shared load + mma  [NEW op]
    wave.transform.set_pipeline_stage %k_loop, %shared_load_a
        { stage = 1, cluster = 0, initiation_interval = 2 }
        : (!transform.op<"wave.iterate">, !transform.any_op) -> ()
    wave.transform.set_pipeline_stage %k_loop, %mma
        { stage = 1, cluster = 1, initiation_interval = 2 }
        : (!transform.op<"wave.iterate">, !transform.any_op) -> ()

    transform.yield
  }
}

Potentially re-usable MLIR Transforms

transform.memref.multibuffer
- get factor from WaveCompileOptions.multi_buffer??
transform.apply_patterns.gpu.eliminate_barriers
- naively insert barriers and use this to clean up?
transform.loop.pipeline
- get read_latency from WaveSchedulingSymbol READ_SHARED_DELAY / READ_GLOBAL_DELAY?
- get iteration_interval from WaveCompileOptions.multi_buffer??
- not aware of multi buffering
- no dynamic shapes
- not aware of GPU resources

Pass Pipeline

Transform sequences run via the existing transform-interpreter pass before water-middle-end-lowering.
Running another round of transforms after water-middle-end-lowering should also be possible but is out of scope for this issue.

[water] scheduling pass #1267

Description

Metadata

Metadata

Assignees

Labels

Type

Fields

Projects

Milestone

Relationships

Development

Issue actions