Calling CUDA and PTX

This guide shows how to use three interfaces for CUDA programming in TVMScript. Backend intrinsics under T.cuda provide common operations, such as synchronization and warp shuffles. T.ptx lets us specify individual PTX instructions. T.cuda.func_call embeds CUDA C++ source so a kernel can call a custom __device__ function.

Calling Backend Intrinsics

The backend provides helpers for thread synchronization, data movement, reductions, and matrix operations:

Interface	Purpose
T.cuda.cta_sync()	Synchronize the threads in a block.
T.cuda.warp_sync()	Synchronize a full warp.
T.cuda.warpgroup_sync(8)	Synchronize a 128-thread warpgroup using named barrier 8.
T.cuda.__shfl_xor_sync(mask, value, lane_mask, width)	Read a value from another lane selected by XOR.
T.cuda.cta_sum(value, num_warps, scratch)	Reduce a value across a block using a shared scratch pointer.

For example, a warp can sum its 32 input values with T.cuda.__shfl_xor_sync:

from tvm_tirx.script import tirx as T

@T.function
def warp_reduce(A: T.Tensor((32,), "float32", align=16)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
    lane_id = T.cuda.lane_id()
    v = T.alloc_local((1,), "float32")
    i = T.alloc_local((1,), "int32")
    v[0] = A[lane_id]
    i[0] = 16
    while i[0] >= 1:
        v[0] += T.cuda.__shfl_xor_sync(0xFFFFFFFF, v[0], i[0], 32)
        i[0] = i[0] // 2
    A[lane_id] = v[0]

All lanes participate, as selected by 0xFFFFFFFF. The XOR distances 16, 8, 4, 2, and 1 combine all 32 values; every lane receives the sum. The device helper calls CUDA’s __shfl_xor_sync:

v_ptr[0] = v_ptr[0] + __shfl_xor_sync(0xFFFFFFFF, v_ptr[0], i_ptr[0], 32);

The TMA Transfers guide shows a complete barrier initialization, arrival, and phase wait around an asynchronous transfer.

Calling PTX Instructions

T.ptx selects individual PTX instructions. PTX names its destination registers explicitly, so the destination is an argument to the call. This kernel uses add.rn.f32 to add two float32 values in each thread; rn selects round-to-nearest-even:

@T.function
def add_ptx(
    A: T.Tensor((32,), T.float32),
    B: T.Tensor((32,), T.float32),
    C: T.Tensor((32,), T.float32),
):
    with T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32)):
        tid = T.cuda.thread_idx("x")
        result = T.local_scalar("float32")
        T.ptx.add.rn.f32(result, A[tid], B[tid])
        C[tid] = result

result is the writable destination. The PTX call puts the sum in that register, and the next statement stores it in C. The call itself has no return value. A CUDA helper such as __shfl_xor_sync can instead return an expression, as in the preceding reduction.

The indexed spelling T.ptx["add.rn.f32"] constructs the same call. Attribute chains escape Python keywords with a trailing underscore, as in T.ptx.ld.global_.b32; PTX’s :: becomes __, as in T.ptx.mbarrier.init.shared__cta.b64. Indexed names use the exact PTX spelling, which is convenient for longer instruction names.

The PTX reference covers predication and special operands. The PTX TableGen guide explains how instruction definitions generate these interfaces and their CUDA helpers.

Calling CUDA C++ Device Functions

A kernel can also reuse a CUDA C++ device function by including its source and calling it through T.cuda.func_call. The example below defines a small ReLU function, which returns the input when it is positive and zero otherwise:

SRC = r"""
__device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; }
"""


@T.function
def relu_cuda(A: T.Tensor((256,), "float32"), B: T.Tensor((256,), "float32")):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=256))
    tx = T.cuda.thread_idx("x")
    B[tx] = T.cuda.func_call("my_relu", A[tx], SRC, ty="float32")

The generated CUDA source includes my_relu and calls it for each element:

__device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; }
// ...
B_ptr[tx] = my_relu(A_ptr[tx]);