tvm_tirx.backend.cuda.op

Contents

tvm_tirx.backend.cuda.op#

CUDA, PTX, and NVSHMEM TIR intrinsic builders.

Classes#

TCGen05InstrDescriptorAttrs(d_dtype, ...[, ...])

Static options carried by tcgen05_encode_instr_descriptor().

TCGen05InstrDescriptorBlockScaledAttrs(...)

Static options carried by tcgen05_encode_instr_descriptor_block_scaled().

TensorMapEncodeTiledAttr(descriptor_dtype, rank)

Fixed options for tensormap_encode_tiled().

Functions#

iket_mark(name[, payload])

Create an NVIDIA IKET marker annotation.

iket_range_start(name[, payload])

Create an NVIDIA IKET token-range start annotation.

iket_range_end(token[, payload])

Create an NVIDIA IKET token-range end annotation.

iket_range_push(name[, payload])

Create an NVIDIA IKET stack-range push annotation.

iket_official_event(event_id[, source_code, ...])

Create an NVIDIA IKET official range-end event.

func_call(func_name, *args[, ty, loc])

Call a CUDA function with its source code as the final operand.

warp_reduce(value, op[, width, ty, loc])

Warp-level butterfly shuffle-XOR reduction.

warp_sum(value[, width])

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

warp_max(value[, width])

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

warp_min(value[, width])

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

cta_reduce(value, op, num_warps, scratch, *)

CTA-wide reduction via warp shuffle + shared memory.

cta_sum(value, num_warps, scratch)

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

cta_max(value, num_warps, scratch)

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

cta_min(value, num_warps, scratch)

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

wait_until(dst, ptr, predicate[, scope, ...])

Read the global word at ptr into dst until predicate holds, and leave the exit value there.

ptx_cp_async_legacy(dst_ptr, dst_offset, ...)

Fold element offsets into the raw cp.async pointer operands.

mov_sreg(bits, reg_name, *[, ty, loc])

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

ptx_legacy_mma(shape, a_layout, b_layout, ...)

Construct legacy MMA with element offsets and an optional bit operator.

bmma_sync(fragment_d, index_d, fragment_a, ...)

TVM intrinsic for tensor core bmma_sync operators

mma_store(m, n, dst_ptr, src_ptr, ...)

Store the result of PTX MMA into a destination pointer.

mma_fill(local_size, local_ptr, offset, *, ty)

Zero-initialize an MMA accumulation register.

ptx_legacy_ldmatrix(trans, num, dtype, ...)

Load a matrix with explicit pointer and element-offset operands.

tcgen05_encode_instr_descriptor(desc, *, ...)

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

tcgen05_encode_instr_descriptor_block_scaled(...)

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

timer_init(profiler_buffer, profiler_tag, ...)

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

timer_start(event_type, profiler_buffer, ...)

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

timer_end(event_type, profiler_buffer, ...)

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

timer_finalize(profiler_buffer, ...[, ty, loc])

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

atomic_add(res_addr, value, *[, ty, loc])

TVM intrinsic to call cuda atomic add instruction

sm100_2sm_leader_smem_addr(ptr)

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

atomic_cas(ptr, old_val, new_val, *[, ty, loc])

TVM intrinsic to call cuda atomic cas instruction

nvshmem_my_pe(*[, ty, loc])

TVM intrinsic to call nvshmem_my_pe()

nvshmem_n_pes(*[, ty, loc])

TVM intrinsic to call nvshmem_n_pes()

nvshmem_getmem_nbi(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_getmem_nbi()

nvshmem_putmem_nbi(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_putmem_nbi()

nvshmem_getmem_nbi_warp(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_getmem_nbi_warp()

nvshmem_putmem_nbi_warp(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_putmem_nbi_warp()

nvshmem_getmem_nbi_block(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_getmem_nbi_block()

nvshmem_putmem_nbi_block(dst, src, nelems, pe, *)

TVM intrinsic to call nvshmem_putmem_nbi_block()

nvshmem_signal_op(sig_addr, signal, sig_op, ...)

TVM intrinsic to call nvshmem_signal_op()

nvshmem_wait_until(ivar, cmp, cmp_value[, ...])

TVM intrinsic to call nvshmem_wait_until()

nvshmem_quiet(*[, ty, loc])

TVM intrinsic to call nvshmem_quiet()

nvshmem_putmem_signal_nbi(dst, src, nelems, ...)

TVM intrinsic to call nvshmem_putmem_signal_nbi()

nvshmem_putmem_signal_nbi_warp(dst, src, ...)

TVM intrinsic to call nvshmem_putmem_signal_nbi_warp()

nvshmem_putmem_signal_nbi_block(dst, src, ...)

TVM intrinsic to call nvshmem_putmem_signal_nbi_block()

nvshmem_fence(*[, ty, loc])

TVM intrinsic to call nvshmem_fence()

nvshmem_barrier_all(*[, ty, loc])

TVM intrinsic to call nvshmem_barrier_all()

tensormap_encode_tiled(*args, ...[, ...])

Encode a tiled tensor map using runtime pointers and shape operands.

any_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.any_sync.

ballot_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.ballot_sync.

bfloat1622float2(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.bfloat1622float2.

bfloat162float(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.bfloat162float.

block_dim(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.block_dim.

block_idx(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.block_idx.

clock64(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.clock64.

cluster_cta_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cluster_cta_id.

cluster_dim(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cluster_dim.

cluster_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cluster_id.

cluster_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cluster_sync.

cta_pair_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cta_pair_id.

cta_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.cta_sync.

cvta_generic_to_shared(*args[, attrs, ...])

Construct a call to tirx.cuda.cvta_generic_to_shared.

decl_tmem(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.decl_tmem.

elect_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.elect_sync.

fadd2_rn(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.fadd2_rn.

fdividef(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.fdividef.

ffs_u32(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.ffs_u32.

float22bfloat162_rn(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.float22bfloat162_rn.

float22bfloat162_rn_from_float2(*args[, ...])

Construct a call to tirx.cuda.float22bfloat162_rn_from_float2.

float22half2(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.float22half2.

float2_x(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.float2_x.

float2_y(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.float2_y.

float8tohalf8(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.float8tohalf8.

float_as_uint(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.float_as_uint.

fmul2_rn(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.fmul2_rn.

fp8x4_e4m3_from_float4(*args[, attrs, ...])

Construct a call to tirx.cuda.fp8x4_e4m3_from_float4.

get_tmem_addr(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.get_tmem_addr.

grid_dim(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.grid_dim.

grid_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.grid_sync.

half2float(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.half2float.

half8tofloat8(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.half8tofloat8.

hmax2(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.hmax2.

hmin2(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.hmin2.

iket_range_pop(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.iket_range_pop.

iket_sentinel_token(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.iket_sentinel_token.

lane_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.lane_id.

linear_thread_id(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.linear_thread_id.

make_float2(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.make_float2.

mbarrier_wait(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.mbarrier_wait.

mbarrier_wait_acquire_cluster(*args[, ...])

Construct a call to tirx.cuda.mbarrier_wait_acquire_cluster.

nano_sleep(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.nano_sleep.

printf(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.printf.

reduce_add_sync_u32(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.reduce_add_sync_u32.

reduce_min_sync_u32(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.reduce_min_sync_u32.

runtime_instr_desc(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.runtime_instr_desc.

smem_addr_from_uint64(*args[, attrs, ...])

Construct a call to tirx.cuda.smem_addr_from_uint64.

syncthreads_and(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.syncthreads_and.

syncthreads_or(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.syncthreads_or.

tcgen05_encode_matrix_descriptor(*args[, ...])

Construct a call to tirx.cuda.tcgen05_encode_matrix_descriptor.

thread_fence(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.thread_fence.

thread_idx(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.thread_idx.

thread_in_warpgroup(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.thread_in_warpgroup.

thread_rank(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.thread_rank.

trap_when_assert_failed(*args[, attrs, ...])

Construct a call to tirx.cuda.trap_when_assert_failed.

uint_as_float(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.uint_as_float.

warp_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.warp_id.

warp_in_warpgroup(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.warp_in_warpgroup.

warp_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.warp_sync.

warpgroup_id(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.warpgroup_id.

warpgroup_sync(*args[, attrs, ty_args, loc, ty])

Construct a call to tirx.cuda.warpgroup_sync.

wgmma_encode_matrix_descriptor(*args[, ...])

Construct a call to tirx.cuda.wgmma_encode_matrix_descriptor.

wgmma_noop_barrier(*args[, attrs, ty_args, ...])

Construct a call to tirx.cuda.wgmma_noop_barrier.
Contents