tvm.backend.cuda

Contents

tvm.backend.cuda#

The CUDA backend — the tile-primitive dispatch, intrinsic builders, the T.cuda / T.ptx script namespaces, and the shared/tensor-memory pools — lives under tvm.backend.cuda, separate from the TIRx frontend (tvm.tirx). Other backends sit alongside it (tvm.backend.rocm and so on).

tvm.backend.cuda#

CUDA-owned TIRx modules.

tvm.backend.cuda.register_backend()#

Register CUDA-owned Python semantics.

tvm.backend.cuda.script_namespace(**kwargs)#

Return the CUDA TVMScript namespace object.

tvm.backend.cuda.script_namespaces(**_)#

Return CUDA-owned TVMScript namespaces.

tvm.backend.cuda.lang#

CUDA-specific TIRx language helpers.

tvm.backend.cuda.op#

CUDA, PTX, and NVSHMEM TIR intrinsic builders.

tvm.backend.cuda.op.const(value, dtype=None, span=None)#

construct a constant

Parameters:

  • value (number) – The content of the constant number.

  • dtype (str or None, optional) – The data type.

  • span (Optional[Span]) – The location of the constant value in the source.

Returns:

const_val – The result expression.

Return type:

tvm.Expr

tvm.backend.cuda.op.bitwise_and(x, y, span=None)#

Take bitwise and of two values

Parameters:

  • x (Expr) – Left operand

  • y (Expr) – Right operand

  • span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

Expr

tvm.backend.cuda.op.call_intrin(dtype: str | Type, func_name, *args, attrs=None, span=None)#

Build expression by calling an intrinsic function.

Intrinsics can be overloaded with multiple data types via the intrinsic translation rule.

Parameters:

  • dtype (str or tvm.ir.Type) – The data type of the result.

  • func_name (str) – The intrinsic function name.

  • args (list) – Positional arguments.

  • attrs (Optional[tvm.ir.Attrs or Dict[str, Object]]) – Additional attributes for the call.

  • span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.tvm_access_ptr(ptype, data, offset, extent, rw_mask)#

Get head access address with memory access pattern info

Parameters:

  • ptype (Expr, PrimType, or str) – The data type of pointer. If a PrimType or str, it is wrapped via type_annotation() so that the lowering rule (which reads args[0].dtype() for the cast type) sees the intended dtype instead of void from a raw StringImm.

  • data (DType*) – The data of pointer.

  • offset (int) – The offset of pointer.

  • extent (int) – The extent of pointer.

  • rw_mask (int) – The read write mask.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_iket_mark(name, payload=None)#

Create an NVIDIA IKET marker annotation.

tvm.backend.cuda.op.cuda_iket_range_start(name, payload=None)#

Create an NVIDIA IKET token-range start annotation.

tvm.backend.cuda.op.cuda_iket_range_end(token, payload=None)#

Create an NVIDIA IKET token-range end annotation.

tvm.backend.cuda.op.cuda_iket_range_push(name, payload=None)#

Create an NVIDIA IKET stack-range push annotation.

tvm.backend.cuda.op.cuda_iket_range_pop()#

Create an NVIDIA IKET stack-range pop annotation.

tvm.backend.cuda.op.cuda_iket_sentinel_token(name)#

Create a no-op NVIDIA IKET range token for warp-uniform control flow.

tvm.backend.cuda.op.cuda_iket_official_event(event_id, source_code='', payload=None)#

Create an NVIDIA IKET official range-end event.

tvm.backend.cuda.op.cuda_func_call(func_name, *args, source_code, return_type='void')#

TVM intrinsic to call a CUDA function. Source code is provided as a string.

Parameters:

  • func_name (str) – The name of the CUDA function.

  • args (Expr) – The arguments to the CUDA function.

  • source_code (str) – The source code of the CUDA function.

  • return_type (str) – The return type of the CUDA function.

tvm.backend.cuda.op.cuda_warp_reduce(value, op, width=32)#

Warp-level butterfly shuffle-XOR reduction.

Reduces value across width adjacent lanes using the specified operation. Codegen emits log2(width) steps of __shfl_xor_sync(0xFFFFFFFF, val, mask) with descending XOR masks.

Parameters:

  • value (Expr) – The per-thread scalar value to reduce.

  • op (str) – Reduction operation: "sum", "max", or "min".

  • width (int) – Number of lanes participating in each reduction group. Must be a power of two in [2, 32]. Defaults to 32 (full warp).

Returns:

call – The reduced value (same dtype as value).

Return type:

Expr

tvm.backend.cuda.op.cuda_warp_sum(value, width=32)#

Convenience wrapper: cuda_warp_reduce(value, "sum", width).

tvm.backend.cuda.op.cuda_warp_max(value, width=32)#

Convenience wrapper: cuda_warp_reduce(value, "max", width).

tvm.backend.cuda.op.cuda_warp_min(value, width=32)#

Convenience wrapper: cuda_warp_reduce(value, "min", width).

tvm.backend.cuda.op.cuda_cta_reduce(value, op, num_warps, scratch)#

CTA-wide reduction via warp shuffle + shared memory.

Two-step reduction: (1) intra-warp shuffle reduction, (2) warp-0 collects per-warp partials from scratch, reduces, broadcasts via __syncthreads(). All CTA threads must participate.

Parameters:

  • value (Expr) – Per-thread scalar value to reduce.

  • op (str) – Reduction operation: "sum", "max", or "min".

  • num_warps (int) – Number of warps in the CTA. Must be a power of two in [1, 32].

  • scratch (Var) – Data pointer to shared-memory scratch space (>= num_warps elements).

Returns:

call – The reduced value broadcast to all threads (same dtype as value).

Return type:

Expr

tvm.backend.cuda.op.cuda_cta_sum(value, num_warps, scratch)#

Convenience wrapper: cuda_cta_reduce(value, "sum", num_warps, scratch).

tvm.backend.cuda.op.cuda_cta_max(value, num_warps, scratch)#

Convenience wrapper: cuda_cta_reduce(value, "max", num_warps, scratch).

tvm.backend.cuda.op.cuda_cta_min(value, num_warps, scratch)#

Convenience wrapper: cuda_cta_reduce(value, "min", num_warps, scratch).

tvm.backend.cuda.op.cuda_warp_sync()#

TVM intrinsic to synchronize threads within the current warp.

This lowers to a CUDA __syncwarp() call.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_cta_sync()#

TVM intrinsic to call CUDA syncthreads (block-wide barrier)

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_grid_sync()#

TVM intrinsic to call CUDA grid-wide sync (cooperative groups)

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_cluster_sync()#

TVM intrinsic to call CUDA cluster-wide barrier sync

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_thread_rank()#

TVM intrinsic that returns cooperative_groups::thread_rank() for the enclosing CTA – the linear thread index within the block.

Useful for building “single thread of CTA” predicates without referencing user-declared scope_id vars. For example, the idiomatic mbarrier.init leader predicate is:

T.cuda.thread_rank() == 0
Returns:

call – The call expression (int32).

Return type:

Expr

tvm.backend.cuda.op.cuda_half2float(src)#

TVM intrinsic to convert half to float

Parameters:

src (Expr) – Source pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_bfloat162float(src)#

TVM intrinsic to convert bfloat16 to float

Parameters:

src (Expr) – Source pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_float22half2(dst, src)#

TVM intrinsic to convert float2 to half2 with rounding

Parameters:

  • dst (Expr) – Destination pointer.

  • src (Expr) – Source pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_trap_when_assert_failed(cond)#

TVM intrinsic to trap when assertion failed (cond == false)

Parameters:

cond (Expr) – Condition to check.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_runtime_instr_desc(desc, sf_id)#

TVM intrinsic to update runtime instruction descriptor

Parameters:

  • desc (Expr) – Pointer to the descriptor (uint32*).

  • sf_id (Expr) – The subfragment id.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_half8tofloat8(src_addr, dst_addr)#

TVM intrinsic to convert 8 half2s to 8 float2s

Parameters:

  • src_addr (Expr) – Source pointer.

  • dst_addr (Expr) – Destination pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_float8tohalf8(src_addr, dst_addr)#

TVM intrinsic to convert 8 float2s to 8 half2s

Parameters:

  • src_addr (Expr) – Source pointer.

  • dst_addr (Expr) – Destination pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_mbarrier_wait(bar, phase)#

Retry mbarrier.try_wait.parity.acquire.cta until it returns true.

Parameters:

  • bar (Var) – The pointer to barrier variable.

  • phase (int) – The phase of the barrier.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_mbarrier_wait_acquire_cluster(bar, phase)#

mbarrier.try_wait.parity.acquire.cluster retry loop.

Cluster-scope acquire wait — used to wait on a barrier that a remote CTA in the cluster arrives on (a group cluster wait).

Parameters:

  • bar (Var) – The pointer to barrier variable.

  • phase (int) – The phase of the barrier.

tvm.backend.cuda.op.ptx_cp_async_legacy(*all_args)#

Legacy ptx_cp_async API taking explicit src/dst offsets.

Signature: (dst_ptr, dst_offset, src_ptr, src_offset, cp_size). Offsets are folded into the pointers via tvm_access_ptr and the call lowers through the raw tirx.s_tir.cp_async_raw op.

T.s_tir.cp_async_raw.legacy runs through _dtype_forward which prepends a dtype= kwarg as a leading positional. The dtype names the element type of the buffer (offsets are in elements of that dtype, not bytes), so this function accepts either 5 or 6 positional args.

tvm.backend.cuda.op.cuda_elect_sync()#

TVM intrinsic to call elect.sync

tvm.backend.cuda.op.cuda_mov_sreg(bits, reg_name)#

TVM intrinsic to tvm instrinsics to fetch PTX pre-defined registers

Parameters:

  • bits (int) – The number of bits of the register.

  • reg_name (str) – The name of the register.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.ptx_legacy_mma(*all_args, operator=None)#

Legacy ptx_mma API.

Signature: (shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index, saturate, operator=None). The accumulator is reused as both input and output (no separate d/c slot). Translation:

  • a_dtype, b_dtype, c_dtype → fork a_type, b_type, c_type (and reuse c_dtype as fork d_type since the accumulator dtype is the output dtype here).

  • (a_ptr, a_offset) and (b_ptr, b_offset) → folded via tvm_access_ptr().

  • (accumulator, c_index) → folded; passed for both d_ptr and c_ptr since the accumulator is reused as the output.

T.ptx_legacy.mma runs through _dtype_forward which prepends a dtype= kwarg as a leading positional, so this function accepts either 13 or 14 positional args.

tvm.backend.cuda.op.mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride)#

Store the result of PTX MMA into a destination pointer.

tvm.backend.cuda.op.mma_store_legacy(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride)#

mma_store with apache-style pointer/offset semantics.

tvm.backend.cuda.op.mma_fill(dtype, local_size, local_ptr, offset)#

Zero-initialize an MMA accumulation register.

tvm.backend.cuda.op.mma_fill_legacy(dtype, local_size, local_ptr, offset)#

mma_fill with apache-style pointer/offset semantics.

tvm.backend.cuda.op.ptx_legacy_ldmatrix(*all_args)#

Legacy ptx_ldmatrix API taking explicit offsets.

Signature: (trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset). Offsets are folded into the pointers via tvm_access_ptr.

T.ptx_legacy.ldmatrix runs through _dtype_forward which prepends a dtype= kwarg as a leading positional naming the buffer element type — offsets are in elements of that dtype, not bytes, so we forward it to tvm_access_ptr for correct scaling.

tvm.backend.cuda.op.cuda_wgmma_encode_matrix_descriptor(desc, addr, ldo, sdo, swizzle)#

TVM intrinsic to create memory descriptor for wgmma instructions

Parameters:

  • desc (Expr) – The pointer to the shared memory descriptor.

  • addr (Expr) – The address of the matrix.

  • ldo (Expr) – The leading dimension offset.

  • sdo (Expr) – The stride dimension offset.

  • swizzle (int) – The swizzle value (CUtensorMapSwizzle_enum).

tvm.backend.cuda.op.cuda_wgmma_noop_barrier(reg)#

TVM intrinsic to call “” : “+{format}”(reg)::”memory”

Parameters:

reg (Expr) – The register to fence.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_tcgen05_encode_matrix_descriptor(desc, addr, ldo, sdo, swizzle)#

TVM intrinsic to create memory descriptor for tcgen05 instructions

Parameters:

  • desc (Expr) – The pointer to the shared memory descriptor.

  • addr (Expr) – The address of the matrix.

  • ldo (Expr) – The leading dimension offset.

  • sdo (Expr) – The stride dimension offset.

  • swizzle (int) – The swizzle value (CUtensorMapSwizzle_enum).

tvm.backend.cuda.op.cuda_tcgen05_encode_instr_descriptor(desc, *, d_dtype, a_dtype, b_dtype, M, N, K, trans_a, trans_b, n_cta_groups=1, neg_a=False, neg_b=False, sat_d=False, is_sparse=False)#

TVM intrinsic to create instruction descriptor for tcgen05 MMA without block scaling

Parameters:

  • desc (Expr) – The pointer to the instruction descriptor.

  • d_dtype (str) – The datatype of resultant matrix D.

  • a_dtype (str) – The datatype of multiplicand matrix A.

  • b_dtype (str) – The datatype of multiplicand matrix B.

  • M (int) – The size of non-reduction dimension of Matrix A.

  • N (int) – The size of non-reduction dimension of Matrix B.

  • K (int) – The size of reduction dimension of Matrix A/B.

  • trans_a (bool) – Whether the multiplicand matrix A is transposed. True for M/N major, False for K major.

  • trans_b (bool) – Whether the multiplicand matrix B is transposed. True for M/N major, False for K major.

  • n_cta_groups (int) – The number of CTA groups involved in the MMA operation.

  • neg_a (bool) – Whether to negate the multiplicand matrix A.

  • neg_b (bool) – Whether to negate the multiplicand matrix B.

  • sat_d (bool) – Whether to saturate the resultant matrix D.

  • is_sparse (bool) – Whether the MMA operation is sparse.

tvm.backend.cuda.op.cuda_tcgen05_encode_instr_descriptor_block_scaled(desc, *, d_dtype, a_dtype, b_dtype, sfa_dtype, sfb_dtype, sfa_tmem_addr, sfb_tmem_addr, M, N, K, trans_a, trans_b, n_cta_groups=1, neg_a=False, neg_b=False, is_sparse=False)#

TVM intrinsic to create instruction descriptor for tcgen05 MMA with block scaling

Parameters:

  • desc (Expr) – The pointer to the instruction descriptor.

  • d_dtype (str) – The datatype of resultant matrix D.

  • a_dtype (str) – The datatype of multiplicand matrix A.

  • b_dtype (str) – The datatype of multiplicand matrix B.

  • sfa_dtype (str) – The datatype of scale factor matrix A.

  • sfb_dtype (str) – The datatype of scale factor matrix B.

  • sfa_tmem_addr (Expr) – The address of the scale factor matrix A in tensor memory, should be uint32_t.

  • sfb_tmem_addr (Expr) – The address of the scale factor matrix B in tensor memory, should be uint32_t.

  • M (int) – The size of non-reduction dimension of Matrix A.

  • N (int) – The size of non-reduction dimension of Matrix B.

  • K (int) – The size of reduction dimension of Matrix A/B.

  • trans_a (bool) – Whether the multiplicand matrix A is transposed. True for M/N major, False for K major.

  • trans_b (bool) – Whether the multiplicand matrix B is transposed. True for M/N major, False for K major.

  • n_cta_groups (int) – The number of CTA groups involved in the MMA operation.

  • neg_a (bool) – Whether to negate the multiplicand matrix A.

  • neg_b (bool) – Whether to negate the multiplicand matrix B.

  • is_sparse (bool) – Whether the MMA operation is sparse.

tvm.backend.cuda.op.timer_init_cuda(profiler_buffer, profiler_tag, profiler_write_offset, num_groups, group_id)#

TVM intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.

Parameters:

  • profiler_buffer (Var) – The buffer to store the profiling result.

  • profiler_tag (Var) – Buffer of length 1 storing the base tag of the current thread.

  • profiler_write_offset (Var) – Buffer of length 1 storing the offset in buffer to write the next profiling result for the current thread.

  • num_groups (int) – The number of groups in the profiler.

  • group_id (Expr) – The group id of the current thread.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.timer_start_cuda(event_type, profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond)#

TVM intrinsic for starting the timer for profiling a specific event, and storing profiling result in a buffer.

Parameters:

  • event_type (Enum) – The event to profile.

  • profiler_buffer (Var) – The buffer to store the profiling result.

  • profiler_tag (Var) – Buffer of length 1 storing the base tag of the current thread.

  • profiler_write_offset (Var) – Buffer of length 1 storing the offset in buffer to write the next profiling result for the current thread.

  • profiler_write_stride (int) – The stride to advance in buffer in the next write.

  • leader_cond (Expr) – The condition to check if the current thread is the leader.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.timer_end_cuda(event_type, profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond)#

TVM intrinsic for ending the timer for profiling a specific event, and storing profiling result in a buffer.

Parameters:

  • event_type (Enum) – The event to profile.

  • profiler_buffer (Var) – The buffer to store the profiling result.

  • profiler_tag (Var) – Buffer of length 1 storing the base tag of the current thread.

  • profiler_write_offset (Var) – Buffer of length 1 storing the offset in buffer to write the next profiling result for the current thread.

  • profiler_write_stride (int) – The stride to advance in buffer in the next write.

  • leader_cond (Expr) – The condition to check if the current thread is the leader.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.timer_finalize_cuda(profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond)#

TVM intrinsic for finalizing the CUDA profiler, and store profiling result in a buffer.

Parameters:

  • profiler_buffer (Var) – The buffer to store the profiling result.

  • profiler_tag (Var) – Buffer of length 1 storing the base tag of the current thread.

  • profiler_write_offset (Var) – Buffer of length 1 storing the offset in buffer to write the next profiling result for the current thread.

  • profiler_write_stride (int) – The stride to advance in buffer in the next write.

  • leader_cond (Expr) – The condition to check if the current thread is the leader.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_atomic_add(res_addr, value)#

TVM intrinsic to call cuda atomic add instruction

Parameters:

  • res_addr (Expr) – The result address.

  • value (Expr) – The value to add.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_thread_fence()#

TVM intrinsic to call cuda thread fence instruction

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_warpgroup_sync(bar_no)#

TVM intrinsic to synchronize a CUDA warpgroup via a named barrier.

Parameters:

bar_no (Expr) – The named barrier id to use for the warpgroup.

Notes

Synchronizes 128 threads in a warpgroup using bar.sync bar_no, 128.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_syncthreads_and(cond)#

TVM intrinsic to call cuda syncthreads_and instruction

Parameters:

cond (Expr) – The condition.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_syncthreads_or(cond)#

TVM intrinsic to call cuda syncthreads_or instruction

Parameters:

cond (Expr) – The condition.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_nano_sleep(time)#

TVM intrinsic to call cuda nano sleep instruction

Parameters:

time (Expr) – The time to sleep.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_printf(fmt, *args)#

TVM intrinsic to call cuda printf instruction

Parameters:

  • fmt (str) – The format string.

  • *args (list) – The arguments to the format string.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_ldg(addr, dtype, *, dst=None, vec='')#

TVM intrinsic to call CUDA C++ __ldg().

Parameters:

  • addr (Expr) – The memory address to load.

  • dtype (str) – The data type of the loaded value.

  • dst (Expr or tuple[Expr], optional) – Destination pointers for vector loads.

  • vec (str) – CUDA vector width. Use "v2" or "v4" together with tuple/list dst.

  • Returns

tvm.backend.cuda.op.cuda_fdividef(x, y)#

TVM intrinsic to call CUDA C++ __fdividef fast float division.

tvm.backend.cuda.op.cuda_get_tmem_addr(addr, row_offset, col_offset)#

TVM intrinsic to call cuda tmem address calculation

Parameters:

  • addr (Expr) – The memory address to calculate.

  • row_offset (Expr) – The row offset to calculate.

  • col_offset (Expr) – The column offset to calculate.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.cuda_cvta_generic_to_shared(ptr)#

Convert a generic pointer to a shared-memory address (uint32).

Wraps __cvta_generic_to_shared(ptr). Used by op-wrappers that precompute the shared-memory address at the wrapper layer instead of inside the asm helper body.

tvm.backend.cuda.op.cuda_smem_addr_from_uint64(cluster_addr)#

Narrow a 64-bit cluster-mapped SMEM address to a 32-bit SMEM address.

Wraps static_cast<unsigned int>(cluster_addr). Used by cp.async.bulk.shared::cluster.* op-wrappers.

tvm.backend.cuda.op.cuda_sm100_2sm_leader_smem_addr(ptr)#

Return the SM100 2SM leader CTA shared-address operand.

The input is a generic pointer to shared memory.

tvm.backend.cuda.op.cuda_any_sync(mask, pred)#

TVM intrinsic for PTX warp-wide any predicate (__any_sync)

Parameters:

  • mask (Expr) – The thread mask (uint32).

  • pred (Expr) – The predicate value (int32).

Returns:

call – The call expression returning 1 if any thread in mask has pred != 0.

Return type:

Expr

tvm.backend.cuda.op.cuda_atomic_cas(ptr, old_val, new_val)#

TVM intrinsic to call cuda atomic cas instruction

Parameters:

  • ptr (Expr) – The pointer to the memory location.

  • old_val (Expr) – The old value.

  • new_val (Expr) – The new value.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_my_pe()#

TVM intrinsic to call nvshmem_my_pe()

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_n_pes()#

TVM intrinsic to call nvshmem_n_pes()

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_getmem_nbi(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_getmem_nbi()

Parameters:

  • dst (Expr) – The pointer to the symmetric address or host/device address of the data object to be updated.

  • src (Expr) – The pointer to the symmetric address of the source data object.

  • nelems (int) – The number of bytes to get per thread.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_nbi(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_putmem_nbi()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the destination data object.

  • src (Expr) – The pointer to the symmetric address or host/device address of the data object to be copied.

  • nelems (int) – The number of bytes to put per thread.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_getmem_nbi_warp(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_getmem_nbi_warp()

Parameters:

  • dst (Expr) – The pointer to the symmetric address or host/device address of the data object to be updated.

  • src (Expr) – The pointer to the symmetric address of the source data object.

  • nelems (int) – The number of bytes to get per warp.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_nbi_warp(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_putmem_nbi_warp()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the destination data object.

  • src (Expr) – The pointer to the symmetric address or host/device address of the data object to be copied.

  • nelems (int) – The number of bytes to put per warp.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_getmem_nbi_block(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_getmem_nbi_block()

Parameters:

  • dst (Expr) – The pointer to the symmetric address or host/device address of the data object to be updated.

  • src (Expr) – The pointer to the symmetric address of the source data object.

  • nelems (int) – The number of bytes to get per block.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_nbi_block(dst, src, nelems, pe)#

TVM intrinsic to call nvshmem_putmem_nbi_block()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the destination data object.

  • src (Expr) – The pointer to the symmetric address or host/device address of the data object to be copied.

  • nelems (int) – The number of bytes to put per block.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_signal_op(sig_addr, signal, sig_op, pe)#

TVM intrinsic to call nvshmem_signal_op()

Parameters:

  • sig_addr (Expr) – The pointer to the symmetric address of the signal word to be updated, must be uint64_t*.

  • signal (uint64_t) – The value used to update sig_addr.

  • sig_op (str) – Operation used to update sig_addr with signal, typical sig_op values are “set” and “add”.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_wait_until(ivar, cmp, cmp_value, type='uint64_t')#

TVM intrinsic to call nvshmem_wait_until()

Parameters:

  • ivar (Expr) – The pointer to the symmetric address of a remotely accessible data object, must be TYPE*.

  • cmp (str) – The compare operator that compares ivar with cmp_value.

  • cmp_value (TYPE) – The value to be compared with ivar.

  • type (str) – The TYPE of ivar and cmp_value.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_quiet()#

TVM intrinsic to call nvshmem_quiet()

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_signal_nbi(dst, src, nelems, sig_addr, signal, sig_op, pe)#

TVM intrinsic to call nvshmem_putmem_signal_nbi()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the data object to be updated on the remote PE.

  • src (Expr) – The pointer to the symmetric address or host/device address of data object containing the data to be copied.

  • nelems (int) – The number of bytes to put per thread.

  • sig_addr (Expr) – The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

  • signal (uint64_t) – The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

  • sig_op (str) – Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_signal_nbi_warp(dst, src, nelems, sig_addr, signal, sig_op, pe)#

TVM intrinsic to call nvshmem_putmem_signal_nbi_warp()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the data object to be updated on the remote PE.

  • src (Expr) – The pointer to the symmetric address or host/device address of data object containing the data to be copied.

  • nelems (int) – The number of bytes to put per warp.

  • sig_addr (Expr) – The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

  • signal (uint64_t) – The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

  • sig_op (str) – Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_putmem_signal_nbi_block(dst, src, nelems, sig_addr, signal, sig_op, pe)#

TVM intrinsic to call nvshmem_putmem_signal_nbi_block()

Parameters:

  • dst (Expr) – The pointer to the symmetric address of the data object to be updated on the remote PE.

  • src (Expr) – The pointer to the symmetric address or host/device address of data object containing the data to be copied.

  • nelems (int) – The number of bytes to put per block.

  • sig_addr (Expr) – The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

  • signal (uint64_t) – The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

  • sig_op (str) – Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

  • pe (int) – The PE number of the remote PE.

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_fence()#

TVM intrinsic to call nvshmem_fence()

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.op.nvshmem_barrier_all()#

TVM intrinsic to call nvshmem_barrier_all()

Returns:

call – The call expression.

Return type:

Expr

tvm.backend.cuda.script#

CUDA TVMScript namespaces.

class tvm.backend.cuda.script.CUDANamespace#

The CUDA intrinsics submodule.

class tvm.backend.cuda.script.NVSHMEMNamespace#

The NVSHMEM intrinsics submodule.

class tvm.backend.cuda.script.PTXLegacyNamespace#

Apache-compatible spellings of instructions the dialect already covers.

They take the historical argument order and are pattern-matched by the passes that lower them, which is why they are not simply deleted: tests inherited from upstream still write them.

class tvm.backend.cuda.script.STIRNamespace#

Nodes the s_tir pipeline’s own passes build.

Nothing here is meant to be written by hand: InjectPTXLDG32 and InjectPTXAsyncCopy construct these, later passes match on them, and codegen turns them into asm. They have a script spelling only so printed IR round-trips.

tvm.backend.cuda.codegen#

Shared machinery behind every CUDA device-intrinsic codegen.

Both dialects build on this layer: tvm.backend.cuda.ptx renders PTX instructions from a table, and tvm.backend.cuda.cpp registers hand-written CUDA C++ device helpers.

  • registry — the op-name → codegen map C++ queries during codegen.

  • schema — device_intrinsic(), the declarative helper registration.

  • header — CUDA header generator and its helper-tag table.

  • types — PTX dtype enum mirroring src/backend/cuda/codegen/ptx.cc.

  • utils — small parsing / validation helpers.

Importing registry and header is load-bearing: their module-level register_global_func decorators publish tirx.intrinsics.cuda.get_codegen and tirx.intrinsics.cuda.header_generator, which codegen_cuda.cc looks up when it emits a kernel. Nothing else imports them for their side effects, so dropping them here fails at kernel-build time, not at import time.

class tvm.backend.cuda.codegen.PTXDataType(value)

A Python equivalent of the provided C++ DataType enum class.

Inherits from IntEnum so that members behave both as enum members and as integers, mirroring the C++ behavior.

see also src/target/source/ptx.cc

tvm.backend.cuda.codegen.device_intrinsic(op_name: str, *, helper_name: str | Callable | None = None, c_signature: str | Callable = '()', body: str | Callable, n_attrs: int = 0, return_type: str | Callable = 'void', tvm_return_type: str | Callable | None = None, templated: bool = False, extra_deps: tuple = ()) → None

Register a CUDA device-helper intrinsic.

Parameters:

  • op_name – Registry key — call_intrin("", "tirx.<op_name>", ...) resolves here. Also used as the default helper name (tvm_builtin_<op_name>) when helper_name is not provided.

  • helper_name – Literal C function name, OR (*args) -> str to compute it from attr values. Defaults to f"tvm_builtin_{op_name}".

  • c_signature – Literal C parameter list including outer parens ("(int x, int y)"), OR (*args) -> str to compute it from attr values. Defaults to "()".

  • body – Literal C body string (already indented), OR (*args) -> str.

  • n_attrs – Number of trailing args that are attrs (consumed by helper_name / c_signature / body callables, NOT forwarded to the helper as call arguments). The first len(args) - n_attrs args are the operand args forwarded to the helper.

  • return_type – C return type. Default "void". Either a literal string or (*args) -> str when the helper return type depends on attrs.

  • tvm_return_type – TVM dtype for the call result, when the helper has a non-void return. Either a literal string ("int32") or (*args) -> str. If omitted and return_type is non-void, it is auto-derived from the _C_TO_TVM_DTYPE table.

  • templated – Prefix the helper with template <typename T>.

  • extra_deps – Helper-tag list (e.g. ("get_tmem_addr",)) forwarded as the second element of the codegen result so the header generator emits the prerequisite snippets.

tvm.backend.cuda.codegen.register_codegen(op, backend='cuda')

Register a codegen function for a given op.

The codegen function should return a cuda_func_call statement, and optionally a list of tags that the codegen function needs.

tvm.backend.cuda.cpp#

Hand-written CUDA C++ device helpers.

The CUDA backend emits device intrinsics two ways. A single PTX instruction is a row in the tvm.backend.cuda.ptx table, rendered by a generic engine. Everything else – anything needing a hand-written __device__ function body – is registered here, grouped by why it needs one:

  • builtins — a CUDA builtin or library call wrapped as an op, no asm.

  • asm — a body that must be hand-written asm: spin-wait loops, barrier pairs, empty compiler-barrier asm, special-register reads.

  • descriptors — wgmma / tcgen05 descriptor bitfield encoding, plus the tcgen05 MMA dtype-kind and shape validation those encoders enforce.

  • instrument — profiler timers, IKET events, printf / trap.

  • nvshmem — NVSHMEM RMA, signal, and collective bindings.

Importing this package is what registers those codegens; the shared registry, device_intrinsic schema, and header generator live in tvm.backend.cuda.codegen.

tvm.backend.cuda.tile_primitive#

tvm.backend.cuda.target_tags#

NVIDIA CUDA target tags.

tvm.backend.cuda.target_tags.register_tag(name: str, config: dict[str, Any], override: bool = False) → Target | None#

Add a user-defined tag into the target tag registry.

Parameters:

  • name (str) – Name of the target, e.g. “nvidia/gtx1080ti”

  • config (Dict[str, Any]) – The config dict used to create the target

  • override (bool) – A boolean flag indicating if overriding existing tags are allowed. If False and the tag has been registered already, an exception will be thrown.

Returns:

target – The target corresponding to the tag None if TVM is built in runtime-only mode.

Return type:

Optional[Target]

Examples

register_tag("nvidia/gtx1080ti", config={
    "kind": "cuda",
    "arch": "sm_61",
})
Contents