tilelang.language.ws_schedule¶

Typed warp-specialization schedule objects.

A WSSchedule is the complete description of how to transform a straight-line kernel into a warp-specialized one. It is built inside the kernel (after buffer allocations, so pipelines can reference the buffers directly), attached with T.annotate_ws_schedule, and consumed by the MaterializeWSSchedule pass.

Classes¶

WSSyncKind

Pipeline synchronization kind.

WSRole

A contiguous warp range with a single duty.

WSPipeline

A full/empty mbarrier pair protecting multi-versioned buffers.

WSInstr

Base class of one step in a role's program.

WSOpRef

Reference to a tile op or child scope by its stable tl.ws_op_id.

WSSync

A pipeline synchronization point.

WSScope

A loop (or the root scope) with per-role instruction lists.

WSSchedule

The complete warp-specialization schedule of one kernel.

Module Contents¶

class tilelang.language.ws_schedule.WSSyncKind¶

Bases: tvm_ffi.dataclasses.Enum

Pipeline synchronization kind.

The producer waits for the empty barrier (ACQUIRE) and signals the full barrier (COMMIT); the consumer waits for the full barrier (WAIT) and signals the empty barrier (RELEASE).

PRODUCER_ACQUIRE: ClassVar[WSSyncKind]¶
PRODUCER_COMMIT: ClassVar[WSSyncKind]¶
CONSUMER_WAIT: ClassVar[WSSyncKind]¶
CONSUMER_RELEASE: ClassVar[WSSyncKind]¶
class tilelang.language.ws_schedule.WSRole(name, *, warps_lo, warps_hi, max_nreg=0)¶

Bases: tvm.ir.Node

A contiguous warp range with a single duty.

Parameters:

  • name (str) – Role name; keys the per-role bodies of every scope.

  • warps_lo (int) – First warp of the role’s range.

  • warps_hi (int) – One past the last warp: the role covers warps lo..hi-1.

  • max_nreg (int) – setmaxnreg budget for the role’s warps; 0 leaves registers untouched.

class tilelang.language.ws_schedule.WSPipeline(name, buffers, depth)¶

Bases: tvm.ir.Node

A full/empty mbarrier pair protecting multi-versioned buffers.

The producer waits for the empty barrier and signals the full barrier; the consumer waits for the full barrier and signals the empty barrier. depth is the number of buffer versions; multiple buffers can share one pipeline.

Parameters:

  • name (str) – Pipeline name; referenced by WSSync instructions.

  • buffers (list[tirx.Buffer]) – The on-chip buffers this pipeline protects (and multi-versions).

  • depth (int) – The number of versions of each buffer.

class tilelang.language.ws_schedule.WSInstr¶

Bases: tvm.ir.Node

Base class of one step in a role’s program.

class tilelang.language.ws_schedule.WSOpRef(id)¶

Bases: WSInstr

Reference to a tile op or child scope by its stable tl.ws_op_id.

Parameters:

id (str)

class tilelang.language.ws_schedule.WSSync(kind, pipeline, stage=0)¶

Bases: WSInstr

A pipeline synchronization point.

Prefer the classmethod constructors:

WSSync.producer_acquire("smem", stage=0)
WSSync.producer_commit("smem", stage=0)
WSSync.consumer_wait("smem", stage=num_stages - 1)
WSSync.consumer_release("smem", stage=num_stages - 1)

Within one role’s scope body, acquire/commit (and wait/release) of a pipeline must pair up at the same stage; entries between them execute at that stage’s iteration offset.

Parameters:
classmethod producer_acquire(pipeline, stage=0)¶

Wait for the empty barrier; binds the stage’s buffer versions.

Parameters:

  • pipeline (str)

  • stage (int)

Return type:

WSSync

classmethod producer_commit(pipeline, stage=0)¶

Signal the full barrier; ends the producer’s span.

Parameters:

  • pipeline (str)

  • stage (int)

Return type:

WSSync

classmethod consumer_wait(pipeline, stage=0)¶

Wait for the full barrier; binds the stage’s buffer versions.

Parameters:

  • pipeline (str)

  • stage (int)

Return type:

WSSync

classmethod consumer_release(pipeline, stage=0)¶

Signal the empty barrier; ends the consumer’s span.

Parameters:

  • pipeline (str)

  • stage (int)

Return type:

WSSync

class tilelang.language.ws_schedule.WSScope(id, bodies)¶

Bases: tvm.ir.Node

A loop (or the root scope) with per-role instruction lists.

Parameters:

  • id (str) – The tl.ws_op_id of the loop this scope schedules, or WSScope.ROOT for the kernel’s implicit root scope.

  • bodies (dict[str, list[WSInstr | str]]) – Role name -> instruction sequence. Plain strings are shorthand for WSOpRef.

ROOT = 'tl.ws_scope_root'¶

The id of the kernel’s implicit root scope.

class tilelang.language.ws_schedule.WSSchedule(num_warps, roles, pipelines, scopes)¶

Bases: tvm.ir.Node

The complete warp-specialization schedule of one kernel.

Parameters:

  • num_warps (int) – Total warp count; overrides the kernel’s thread extent.

  • roles (list[WSRole])

  • pipelines (list[WSPipeline])

  • scopes (list[WSScope])