tvm_tirx.backend.cuda.lang

Contents

tvm_tirx.backend.cuda.lang#

CUDA-specific TIRx language helpers.

Classes#

MBarrier(*args, **kwargs)

Mbarrier wrapper with regular mbarrier.arrive.

Pipeline(*args, **kwargs)

A full/empty mbarrier pair for a software-pipelined data flow.

PipelineState(*args, **kwargs)

Tracks stage and phase for a software-pipelined ring buffer.

BaseTileScheduler(*args, **kwargs)

Base class for tile schedulers with common state and macros.

ClusterPersistentScheduler2D(*args, **kwargs)

Tile scheduler for cluster-based persistent kernels.

FlashAttentionLPTScheduler(*args, **kwargs)

LPT scheduler with L2 swizzle for causal flash attention.

FlashAttentionLinearScheduler(*args, **kwargs)

Linear 3D scheduler for flash attention (batch, head, m_block).

SMEMPool(*args, **kwargs)

Bump allocator over a contiguous shared memory region.

TCGen05Bar(*args, **kwargs)

Barrier signaled by tcgen05 commit.

TMABar(*args, **kwargs)

Barrier signaled by TMA (mbarrier.arrive.expect_tx).
Contents