CUDA Authoring and Support APIs

CUDA Authoring and Support APIs#

CUDA helpers#

The CUDA backend installs the Tx.cuda namespace when it is loaded:

from tvm_tirx.script import tirx as Tx

Tx.cuda.cta_sync()
leader = Tx.cuda.elect_sync()

These helpers cover operations that need multiple instructions, C/C++ expressions, descriptor packing, compiler annotations, or other behavior that is not represented as one table-driven PTX instruction. For a single supported PTX instruction, use Direct PTX Instructions.

The current Tx.cuda surface includes:

Category

Helpers

Synchronization and participation

any_sync, elect_sync, warp_sync, warpgroup_sync, cta_sync, grid_sync, cluster_sync, syncthreads_and, syncthreads_or, thread_rank, ballot_sync, __shfl_sync, __shfl_up_sync, __shfl_down_sync, __shfl_xor_sync, and __activemask

Barriers and memory ordering

mbarrier_wait, mbarrier_wait_acquire_cluster, thread_fence, atomic_add, and atomic_cas

Reductions

warp_reduce, warp_sum, warp_max, warp_min, cta_reduce, cta_sum, cta_max, cta_min, reduce_add_sync_u32, and reduce_min_sync_u32

Descriptors and addresses

wgmma.noop_barrier, wgmma.encode_matrix_descriptor, tcgen05.encode_matrix_descriptor, tcgen05.encode_instr_descriptor, tcgen05.encode_instr_descriptor_block_scaled, runtime_instr_desc, get_tmem_addr, cvta_generic_to_shared, smem_addr_from_uint64, sm100_2sm_leader_smem_addr, and mov_sreg

Loads, calls, and diagnostics

func_call, printf, trap_when_assert_failed, nano_sleep, clock64, and ffs_u32

Numeric conversion and packed math

half2float, bfloat162float, float22half2, half8tofloat8, float8tohalf8, uint_as_float, float_as_uint, make_float2, float2_x, float2_y, fmul2_rn, fadd2_rn, float22bfloat162_rn, float22bfloat162_rn_from_float2, bfloat1622float2, hmin2, hmax2, fp8x4_e4m3_from_float4, and fdividef

Instrumentation and compatibility

iket.mark, iket.range_start, iket.range_end, iket.range_push, iket.range_pop, iket.sentinel_token, iket.official_event, timer_init, timer_start, timer_end, timer_finalize, mma_store, and mma_fill

mma_store and mma_fill remain available for downstream emitters. TVM CUDA codegen does not emit them; express stores and initialization with ordinary tensor assignments.

NVSHMEM namespace#

The same backend installs Tx.nvshmem for my_pe, n_pes, signal_op, wait_until, quiet, fence, and barrier_all. Its getmem_nbi, putmem_nbi, and putmem_signal_nbi operations also provide .warp and .block variants.

CUDA language utilities#

tvm_tirx.backend.cuda.lang provides reusable objects for hand-written CUDA TIRx kernels, including schedulers, pipelines, barriers, and shared-memory pools.

class tvm_tirx.backend.cuda.lang.BaseTileScheduler(*args, **kwargs)#

Base class for tile schedulers with common state and macros.

class tvm_tirx.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False)#

Persistent two-dimensional tile scheduler with group-major or serpentine traversal. init selects the first tile for a cluster, next_tile or next_tile_stride advances it, and valid reports whether the current m_idx / n_idx is in range.

class tvm_tirx.backend.cuda.lang.FlashAttentionLPTScheduler(*args, **kwargs)#

LPT scheduler with L2 swizzle for causal flash attention.

Processes high-work Q blocks (with more KV blocks to attend to) first using Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle for better cache locality across batch*head dimensions.

The LPT aspect comes from reversing m_block order: lower Q blocks have more KV blocks to process due to causal masking, so processing them first balances load.

The scheduler is only applied to non-persistent kernels.

L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.

Parameters:

  • prefix (str) – Prefix for TIR variable names

  • num_batches (int) – Number of batches

  • num_heads (int) – Number of KV heads

  • num_m_blocks (int) – Number of Q blocks (M dimension tiles)

  • num_ctas (int) – Number of CTAs (should equal total_tasks for causal)

  • l2_swizzle (int) – L2 swizzle factor for cache locality

update_current_m_n_idx(linear_idx)#

Convert linear index to (batch, head, m_block) with LPT + L2 swizzle.

init(cta_id)#

Initialize scheduler with CTA ID.

next_tile()#

Advance to the next tile.

Single-tile mode (num_ctas=None, the default): each CTA owns one task; terminate. Persistent mode (num_ctas=N): stride by N, like FlashAttentionLinearScheduler, while keeping the LPT + L2 swizzle index mapping.

valid()#

Check if there are more tiles to process.

class tvm_tirx.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas)#

Linear persistent scheduler over (batch, head, m_block) tasks. init starts from a CTA index, next_tile advances by num_ctas, and valid reports whether work remains.

class tvm_tirx.backend.cuda.lang.MBarrier(*args, **kwargs)#

Mbarrier wrapper with regular mbarrier.arrive.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • depth (int) – Number of barrier slots (one per pipeline stage).

  • phase_offset (int) – XORed into the phase bit on every wait / arrive.

  • leader (Expr, optional) – Boolean predicate selecting the single thread that runs mbarrier.init. Defaults to T.cuda.thread_rank() == 0 – thread 0 of the enclosing CTA, which always picks exactly one thread regardless of which index variables the caller binds. Override only when you want a different CTA-local thread to do the init.

remote_view(rank)#

Create a view of this barrier mapped to another CTA’s shared memory.

The returned view retains the local barrier and target CTA so arrive emits the cluster form. Its mapped buffer remains available through ptr_to for operations that consume a remote shared-memory pointer. init and wait are local-only and reject remote views.

class tvm_tirx.backend.cuda.lang.Pipeline(*args, **kwargs)#

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

Pass barrier-type tags and Pipeline constructs and inits the barriers itself. Tags: "tma" (TMABar), "tcgen05" (TCGen05Bar), "mbar" (MBarrier). The barrier type and arrival count of each event stay explicit at the call site – e.g. Pipeline(pool, n, full="tma", empty="tcgen05", init_empty=NUM_CONSUMER).

Both signals are required: a Pipeline is a pair. For a one-way event (a pure “X happened” signal with no slot to recycle) use a bare barrier (TMABar/TCGen05Bar/MBarrier) directly – it has no empty side.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • stages (int) – Number of pipeline stages (barrier slots).

  • full (str) – Barrier-type tag for the full / empty signal (see above).

  • empty (str) – Barrier-type tag for the full / empty signal (see above).

  • init_full (int) – Expected arrival count for the full / empty barrier.

  • init_empty (int) – Expected arrival count for the full / empty barrier.

  • empty_phase_offset (int) – XORed into the empty barrier’s phase bit on every wait / arrive.

  • leader (Expr, optional) – Propagated to both barriers; defaults to thread 0 of the CTA.

class tvm_tirx.backend.cuda.lang.PipelineState(*args, **kwargs)#

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

This class does not know anything about full/empty barriers. Use it when the kernel manually waits/signals barriers, or when the stage/phase drives a ring not wrapped in a Pipeline.

Parameters:

  • depth (int) – Number of stages in the ring.

  • phase (int, optional) – Initial phase. Omit when initialization should happen later.

class tvm_tirx.backend.cuda.lang.SMEMPool(*args, **kwargs)#

Bump allocator over a contiguous shared memory region.

Parameters:

ptr (Var or None, optional) – If omitted, an alloc_tensor(ty_args=[Tensor([0], "uint8", scope="shared.dyn")]) is created automatically and commit() must be called after all allocations to emit the size annotation. If a Var is provided, the caller manages the backing buffer and commit() is a no-op.

alloc_tcgen05_mma_AB(shape, dtype='float16', swizzle_mode='auto', align=1024)#

Allocate MMA-compatible shared memory with an inferred swizzle layout.

commit(size=None)#

Record the pool resource requirement on the enclosing kernel entry.

Must be called after all alloc() / move_base_to() calls.

Parameters:

size (int, optional) – Explicit shared memory size in bytes. When None (the default), the high-water mark max_offset tracked by the allocator is used.

class tvm_tirx.backend.cuda.lang.TCGen05Bar(*args, **kwargs)#

Barrier signaled by tcgen05 commit.

The caller is responsible for ensuring only one thread issues the commit, e.g. by wrapping the call in if T.cuda.elect_sync(): or by passing pred=T.cuda.elect_sync(). The pred form emits the predicated instruction (@p tcgen05.commit) instead of a branch and lets one elected leader predicate be shared across several commits.

class tvm_tirx.backend.cuda.lang.TMABar(*args, **kwargs)#

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

When tx_count is None, falls back to a remote mbarrier.arrive (matching MBarrier.arrive defaults).

IKET profiling#

NVIDIA IKET annotations and run-iket orchestration.

profile runs an explicit replayable command. run is intended for a script’s __main__ block: the parent process asks run-iket to replay the original script or python -m invocation, while the injected tracker and capture processes execute the supplied callable and then exit.

exception tvm_tirx.backend.cuda.iket.IketProfileError(message: str, *, returncode: int | None = None, command: Sequence[str] = (), output_tail: str = '', timeout: float | None = None)

An error while validating or running an official IKET profile.

class tvm_tirx.backend.cuda.iket.IketProfileResult(output_dir: Path, postprocess: str, json_traces: tuple[Path, ...], perfetto_traces: tuple[Path, ...], html_reports: tuple[Path, ...], command: tuple[str, ...] = ())

Published artifacts from a successful official IKET profile.

property trace_json: Path

Return the only JSON trace requested by this profile.

property perfetto: Path

Return the only Perfetto trace requested by this profile.

property html: Path

Return the only HTML report requested by this profile.

property trace: Any

Return the JSON trace without imposing a schema wrapper.

property launches: Any

Return trace["launches"] without reshaping it.

class tvm_tirx.backend.cuda.iket.IketProfiler(*args, **kwargs)

TIRx annotations compiled for NVIDIA’s official IKET runtime.

compile(mod, target=None, *, tir_pipeline='tirx')

Compile official IKET metadata and NativeDump placeholders.

tvm_tirx.backend.cuda.iket.profile(command: Sequence[str | PathLike[str]], *, output_dir: str | PathLike[str], postprocess: str = 'all', clobber: bool = False, cwd: str | PathLike[str] | None = None, env: Mapping[str, str | PathLike[str]] | None = None, max_ts_cnt_per_warp: int | None = None, keep: bool = False, timeout: float | None = 600.0) → IketProfileResult

Profile a replayable command with NVIDIA’s official run-iket tool.

tvm_tirx.backend.cuda.iket.run(main, *, output_dir: str | PathLike[str], postprocess: str = 'all', clobber: bool = False, cwd: str | PathLike[str] | None = None, env: Mapping[str, str | PathLike[str]] | None = None, max_ts_cnt_per_warp: int | None = None, keep: bool = False, timeout: float | None = 600.0) → IketProfileResult

Replay the active script under run-iket and execute main in its passes.

Module-level code executes once in the parent, tracker, and capture processes. main executes only in tracker and capture. Code following run executes only in the parent after a successful profile.

CUDA-specific transforms#

CUDA-specific TIRx transformations.

tvm_tirx.backend.cuda.transform.BindBackendConfig()

Resolve CUDA function and device-entry compilation settings.

Run after BindTarget and before tile dispatch. Function defaults and entry overrides become target and JSON backend_config attributes; non-CUDA functions are unchanged.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.config import backend_config_json
from tvm_tirx.backend.cuda.backend_config import resolve_backend_config
from tvm_tirx.backend.cuda.transform import BindBackendConfig

target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
config = backend_config_json({"cuda": resolve_backend_config(target=target)})

@T.ir_module
class Before:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr({"target": target})
        with T.device_entry():
            A[0] = 1

@T.ir_module
class Expected:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr({"target": target, "backend_config": config})
        with T.device_entry(attrs={"target": target, "backend_config": config}):
            A[0] = 1

after = BindBackendConfig()(Before)
tvm.ir.assert_structural_equal(after, Expected)
tvm_tirx.backend.cuda.transform.LowerIket()

Lower frontend-only NVIDIA IKET annotations.

Run after SplitHostDevice and before MakePackedAPI. Strip annotations for regular compilation; emit metadata and NativeDump placeholders for IKET-enabled modules. The external run-iket process owns trace collection and postprocessing.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.cuda.transform import LowerIket

@T.ir_module
class Before:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.cuda.iket_mark("store")
        A[0] = 1

@T.ir_module
class Expected:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        A[0] = 1

after = LowerIket()(Before)
tvm.ir.assert_structural_equal(after, Expected)
tvm_tirx.backend.cuda.transform.SpecializeDeviceHelpers()

Specialize private device helpers for each CUDA compilation group.

Run after SplitHostDevice. Kernel target and backend configuration define each group. Called helpers receive those attributes and are cloned when their settings differ; helpers already in the same group can be reused. Unreferenced device-only helpers are omitted from the output module.

tvm_tirx.backend.cuda.transform.SpecializeEntryHelpers()

Specialize helpers for each CUDA device entry’s resolved settings.

Run after BindBackendConfig and before TileDispatch and SplitHostDevice. Clone helpers when entry settings differ, propagating them through the call graph; remove originals that become unreachable.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.cuda.transform import SpecializeEntryHelpers

def settings(arch):
    return {"target": tvm.target.Target({"kind": "cuda", "arch": arch}),
            "backend_config": '{"cuda":{"arch":"' + arch + '"}}'}

host_attrs, entry_attrs = settings("sm_100a"), settings("sm_90a")

@T.ir_module
class Before:
    @T.function(private=True)
    def helper(x: T.int32) -> T.int32:
        T.func_attr(host_attrs)
        return x + 1

    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr(host_attrs)
        with T.device_entry(attrs=entry_attrs):
            A[0] = Before.helper(A[0])

@T.ir_module
class Expected:
    @T.function(private=True)
    def helper_cuda(x: T.int32) -> T.int32:
        T.func_attr(entry_attrs)
        return x + 1

    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr(host_attrs)
        with T.device_entry(attrs=entry_attrs):
            A[0] = Expected.helper_cuda(A[0])

after = SpecializeEntryHelpers()(Before)
tvm.ir.assert_structural_equal(after, Expected)
Contents