tilelang.language.tir.op¶

Attributes¶

`call_packed`
`call_cpacked`
`call_packed_lowered`
`call_cpacked_lowered`
`call_llvm_intrin`
`call_llvm_pure_intrin`
`tvm_stack_alloca`
`tvm_stack_make_shape`
`tvm_stack_make_array`
`assume`
`undef`
`start_profile_intrinsic`
`end_profile_intrinsic`
`tvm_tuple`
`tvm_struct_get`
`tvm_struct_set`
`address_of`
`lookup_param`
`tvm_thread_allreduce`
`tvm_thread_invariant`
`tvm_storage_sync`
`tvm_warp_shuffle`
`tvm_warp_shuffle_up`
`tvm_warp_shuffle_down`
`tvm_warp_activemask`
`tvm_throw_last_error`
`tvm_load_matrix_sync`
`tvm_mma_sync`
`tvm_bmma_sync`
`tvm_fill_fragment`
`tvm_store_matrix_sync`
`ptx_mma`
`mma_store`
`ptx_cp_async_bulk`
`ptx_commit_group`
`ptx_wait_group`
`ptx_cp_async_barrier`
`ptx_init_barrier_thread_count`
`ptx_arrive_barrier`
`ptx_wait_barrier`
`create_barriers`
`vectorlow`
`vectorhigh`
`vectorcombine`
`ret`
`min_value`
`max_value`
`exp`
`exp2`
`exp10`
`erf`
`tanh`
`sigmoid`
`log`
`log2`
`log10`
`log1p`
`tan`
`cos`
`cosh`
`acos`
`acosh`
`sin`
`sinh`
`asin`
`asinh`
`atan`
`atanh`
`atan2`
`sqrt`
`rsqrt`
`clz`
`floor`
`ceil`
`trunc`
`abs`
`bitwise_and`
`bitwise_not`
`bitwise_or`
`bitwise_xor`
`nearbyint`
`nextafter`
`hypot`
`copysign`
`ldexp`
`likely`
`isnan`
`isnullptr`
`isfinite`
`isinf`
`popcount`
`q_multiply_shift`
`q_multiply_shift_per_axis`
`shift_left`
`shift_right`
`fmod`
`if_then_else`
`truncdiv`
`truncmod`
`floordiv`
`floormod`
`ceildiv`
`TVMBackendAllocWorkspace`
`TVMBackendFreeWorkspace`
`anylist_getitem`
`anylist_resetitem`
`anylist_setitem_call_packed`
`anylist_setitem_call_cpacked`
`vscale`

Functions¶

`extract_lane`(vector, lane[, span])	Extract one scalar lane from a fixed-width vector expression.
`call_intrin`(dtype, func_name, *args[, annotations, span])	Build expression by calling an intrinsic function.
`call_pure_extern`(dtype, func_name, *args[, span])	Build expression by calling a pure extern function.
`call_extern`(dtype, func_name, *args[, span])	Build expression by calling a extern function.
`tvm_access_ptr`(ptype, data, offset, extent, rw_mask)	Get head access address with memory access pattern info
`ptx_mma_sp`(dtype, shape, A_layout, B_layout, A_dtype, ...)	TVM intrinsic for sparse tensor core ptx instructions
`ptx_mma_block_scale`(accum_dtype, shape, A_layout, ...)	TVM intrinsic for SM120a warp-level NVF4 block-scaled MMA.
`ptx_wgmma_ss`(dtype, wgmma_prefix, a_is_k_major, ...)	TVM intrinsic for ptx tensor core wmma instructions
`ptx_wgmma_rs`(dtype, wgmma_prefix, b_is_k_major, ...)
`ptx_wgmma_sp_ss`(dtype, wgmma_prefix, a_is_k_major, ...)
`ptx_wgmma_sp_rs`(dtype, wgmma_prefix, b_is_k_major, ...)
`ptx_tcgen05_mma_ss`(kind_dtype, desc_a, A_offset, ...)	TVM intrinsic for tcgen05.mma shared-memory x shared-memory instructions.
`ptx_tcgen05_mma_ts`(kind_dtype, A_ptr, A_offset, ...[, ...])	TVM intrinsic for tcgen05.mma tensor-memory x shared-memory instructions.
`ptx_tcgen05_mma_blockscaled_ss`(kind_dtype, desc_a, ...)	TVM intrinsic for tcgen05.mma block-scaled (mxf8f6f4.block_scale) instructions.
`mma_fill`(dtype, local_size, local_ptr, offset)	TVM intrinsic for zero-initalizing an MMA accumulation register
`ptx_ldmatrix`(trans, num, src_access_ptr, dst_access_ptr)	TileLang intrinsic for ptx load matrix from shared memory
`ptx_cp_async`(dst_access_ptr, src_access_ptr, num_elems)	TVM intrinsic for ptx async copy from global to shared memory using cp.async
`tvm_mfma`(dtype, shape, A_layout, B_layout, A_dtype, ...)	TVM intrinsic for amd matrix core mfma instructions
`tvm_mfma_store`(dtype, m, n, dst_ptr, src_ptr, ...)	TVM intrinsic for storing the result of PTX MMA into a destination pointer
`tvm_rdna_wmma`(dtype, shape, A_layout, B_layout, ...)	TVM intrinsic for amd matrix core mfma instructions
`tvm_rdna_wmma_store`(dtype, m, n, dst_ptr, src_ptr, ...)	TVM intrinsic for storing the result of PTX MMA into a destination pointer
`ptx_fence_barrier_init`()	TVM intrinsic for ptx fence barrier initialization.
`ptx_arrive_barrier_expect_tx`(barrier_id, byte_count)	TVM intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx
`infinity`(dtype[, span])	infinity value of dtype
`reinterpret`(value, dtype[, span])	Reinterpret cast a value to dtype.
`round`(x[, rounding_mode, span])	Round elements of the array to the nearest integer.
`pow_of_int`(x, y)	Fast power operation than pow(float, float).
`pow`(x, y[, span])	x power y

Module Contents¶

tilelang.language.tir.op.call_packed¶

tilelang.language.tir.op.call_cpacked¶

tilelang.language.tir.op.call_packed_lowered¶

tilelang.language.tir.op.call_cpacked_lowered¶

tilelang.language.tir.op.call_llvm_intrin¶

tilelang.language.tir.op.call_llvm_pure_intrin¶

tilelang.language.tir.op.tvm_stack_alloca¶

tilelang.language.tir.op.tvm_stack_make_shape¶

tilelang.language.tir.op.tvm_stack_make_array¶

tilelang.language.tir.op.assume¶

tilelang.language.tir.op.undef¶

tilelang.language.tir.op.start_profile_intrinsic¶

tilelang.language.tir.op.end_profile_intrinsic¶

tilelang.language.tir.op.tvm_tuple¶

tilelang.language.tir.op.tvm_struct_get¶

tilelang.language.tir.op.tvm_struct_set¶

tilelang.language.tir.op.address_of¶

tilelang.language.tir.op.lookup_param¶

tilelang.language.tir.op.tvm_thread_allreduce¶

tilelang.language.tir.op.tvm_thread_invariant¶

tilelang.language.tir.op.tvm_storage_sync¶

tilelang.language.tir.op.tvm_warp_shuffle¶

tilelang.language.tir.op.tvm_warp_shuffle_up¶

tilelang.language.tir.op.tvm_warp_shuffle_down¶

tilelang.language.tir.op.tvm_warp_activemask¶

tilelang.language.tir.op.tvm_throw_last_error¶

tilelang.language.tir.op.tvm_load_matrix_sync¶

tilelang.language.tir.op.tvm_mma_sync¶

tilelang.language.tir.op.tvm_bmma_sync¶

tilelang.language.tir.op.tvm_fill_fragment¶

tilelang.language.tir.op.tvm_store_matrix_sync¶

tilelang.language.tir.op.ptx_mma¶

tilelang.language.tir.op.mma_store¶

tilelang.language.tir.op.ptx_cp_async_bulk¶

tilelang.language.tir.op.ptx_commit_group¶

tilelang.language.tir.op.ptx_wait_group¶

tilelang.language.tir.op.ptx_cp_async_barrier¶

tilelang.language.tir.op.ptx_init_barrier_thread_count¶

tilelang.language.tir.op.ptx_arrive_barrier¶

tilelang.language.tir.op.ptx_wait_barrier¶

tilelang.language.tir.op.create_barriers¶

tilelang.language.tir.op.vectorlow¶

tilelang.language.tir.op.vectorhigh¶

tilelang.language.tir.op.vectorcombine¶

tilelang.language.tir.op.ret¶

tilelang.language.tir.op.min_value¶

tilelang.language.tir.op.max_value¶

tilelang.language.tir.op.exp¶

tilelang.language.tir.op.exp2¶

tilelang.language.tir.op.exp10¶

tilelang.language.tir.op.erf¶

tilelang.language.tir.op.tanh¶

tilelang.language.tir.op.sigmoid¶

tilelang.language.tir.op.log¶

tilelang.language.tir.op.log2¶

tilelang.language.tir.op.log10¶

tilelang.language.tir.op.log1p¶

tilelang.language.tir.op.tan¶

tilelang.language.tir.op.cos¶

tilelang.language.tir.op.cosh¶

tilelang.language.tir.op.acos¶

tilelang.language.tir.op.acosh¶

tilelang.language.tir.op.sin¶

tilelang.language.tir.op.sinh¶

tilelang.language.tir.op.asin¶

tilelang.language.tir.op.asinh¶

tilelang.language.tir.op.atan¶

tilelang.language.tir.op.atanh¶

tilelang.language.tir.op.atan2¶

tilelang.language.tir.op.sqrt¶

tilelang.language.tir.op.rsqrt¶

tilelang.language.tir.op.clz¶

tilelang.language.tir.op.floor¶

tilelang.language.tir.op.ceil¶

tilelang.language.tir.op.trunc¶

tilelang.language.tir.op.abs¶

tilelang.language.tir.op.bitwise_and¶

tilelang.language.tir.op.bitwise_not¶

tilelang.language.tir.op.bitwise_or¶

tilelang.language.tir.op.bitwise_xor¶

tilelang.language.tir.op.nearbyint¶

tilelang.language.tir.op.nextafter¶

tilelang.language.tir.op.hypot¶

tilelang.language.tir.op.copysign¶

tilelang.language.tir.op.ldexp¶

tilelang.language.tir.op.likely¶

tilelang.language.tir.op.isnan¶

tilelang.language.tir.op.isnullptr¶

tilelang.language.tir.op.isfinite¶

tilelang.language.tir.op.isinf¶

tilelang.language.tir.op.popcount¶

tilelang.language.tir.op.q_multiply_shift¶

tilelang.language.tir.op.q_multiply_shift_per_axis¶

tilelang.language.tir.op.shift_left¶

tilelang.language.tir.op.shift_right¶

tilelang.language.tir.op.fmod¶

tilelang.language.tir.op.if_then_else¶

tilelang.language.tir.op.truncdiv¶

tilelang.language.tir.op.truncmod¶

tilelang.language.tir.op.floordiv¶

tilelang.language.tir.op.floormod¶

tilelang.language.tir.op.ceildiv¶

tilelang.language.tir.op.TVMBackendAllocWorkspace¶

tilelang.language.tir.op.TVMBackendFreeWorkspace¶

tilelang.language.tir.op.anylist_getitem¶

tilelang.language.tir.op.anylist_resetitem¶

tilelang.language.tir.op.anylist_setitem_call_packed¶

tilelang.language.tir.op.anylist_setitem_call_cpacked¶

tilelang.language.tir.op.vscale¶

tilelang.language.tir.op.extract_lane(vector, lane, span=None)¶

Extract one scalar lane from a fixed-width vector expression.

Parameters:

vector (PrimExpr) – The vector expression to extract from.
lane (int or IntImm) – The zero-based lane index. The index must be known at compile time.
span (Optional[Span]) – The location of this expression in the source code.

Returns:

result – A scalar expression with the vector’s element dtype.

Return type:

PrimExpr

tilelang.language.tir.op.call_intrin(dtype, func_name, *args, annotations=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) – The data type of the result.
func_name (str) – The intrinsic function name.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.call_pure_extern(dtype, func_name, *args, span=None)¶

Build expression by calling a pure extern function.

Parameters:

dtype (str) – The data type of the result.
func_name (str) – The extern function name.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.call_extern(dtype, func_name, *args, span=None)¶

Build expression by calling a extern function.

Parameters:

dtype (str) – The data type of the result.
func_name (str) – The extern function name.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_access_ptr(ptype, data, offset, extent, rw_mask)¶

Get head access address with memory access pattern info

Parameters:

ptype (Expr) – The data type of pointer.
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:

PrimExpr

tilelang.language.tir.op.ptx_mma_sp(dtype, shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index, metadata, meta_index, sparse_selector, saturate)¶

TVM intrinsic for sparse tensor core ptx instructions https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

Parameters:

dtype (str) – The data type of the result.
shape (str) – The shape of mma fragment.
A_layout (Literal["row", "col"]) – The layout of multiplicand fragment A.
B_layout (Literal["row", "col"]) – The layout of multiplicand fragment B.
A_dtype (str) – The data type of multiplicand fragment A.
B_dtype (str) – The data type of multiplicand fragment B.
C_dtype (str) – The data type of accumulator fragment C.
multiplicand_a (Var) – The multiplicand fragment A variable.
a_index (Expr) – The index of multiplicand fragment A.
multiplicand_b (Var) – The multiplicand fragment B variable.
b_index (Expr) – The index of multiplicand fragment B.
accumulator (Var) – The accumulator fragment C variable.
c_index (Expr) – The index of accumulator fragment C.
metadata (Expr) – The metadata of operand.
meta_index (Expr) – The metadata index of operand.
sparse_selector (Expr) – The sparse selector indicating the thread that stores the metadata.
saturate (bool) – The optional saturation at the output.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_mma_block_scale(accum_dtype, shape, A_layout, B_layout, kind, scale_vec_size, A_dtype, B_dtype, scale_type, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index, scale_a, scale_b, scale_a_byte_id=0, scale_a_thread_id=0, scale_b_byte_id=0, scale_b_thread_id=0)¶: TVM intrinsic for SM120a warp-level NVF4 block-scaled MMA.

tilelang.language.tir.op.ptx_wgmma_ss(dtype, wgmma_prefix, a_is_k_major, b_is_k_major, a_dtype_abbrv, b_dtype_abbrv, accum_dtype_abbrv, A_desc, A_offset, B_desc, B_offset, C_data, C_offset, scale_out, scale_in_a, scale_in_b)¶: TVM intrinsic for ptx tensor core wmma instructions https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-wmma

tilelang.language.tir.op.ptx_wgmma_rs(dtype, wgmma_prefix, b_is_k_major, a_dtype_abbrv, b_dtype_abbrv, accum_dtype_abbrv, A_buf, A_offset, B_desc, B_offset, C_data, C_offset, scale_out, scale_in_a, scale_in_b)¶

tilelang.language.tir.op.ptx_wgmma_sp_ss(dtype, wgmma_prefix, a_is_k_major, b_is_k_major, a_dtype_abbrv, b_dtype_abbrv, accum_dtype_abbrv, A_desc, A_offset, E_data, E_offset, sparse_selector, B_desc, B_offset, C_data, C_offset, scale_out, scale_in_a, scale_in_b)¶

tilelang.language.tir.op.ptx_wgmma_sp_rs(dtype, wgmma_prefix, b_is_k_major, a_dtype_abbrv, b_dtype_abbrv, accum_dtype_abbrv, A_buf, A_offset, E_buf, E_offset, sparse_selector, B_desc, B_offset, C_data, C_offset, scale_out, scale_in_a, scale_in_b)¶

tilelang.language.tir.op.ptx_tcgen05_mma_ss(kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3, enable_ws=False, enable_2cta=False, ws=None, warp_specialized=None, variant=None)¶

TVM intrinsic for tcgen05.mma shared-memory x shared-memory instructions.

Expects 14 or 15 positional arguments: (kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset,

desc_val, scale_out, mask0, mask1, mask2, mask3[, enable_ws]).

Aliases: you can also pass ws or warp_specialized (booleans) instead of enable_ws. Alternatively, use variant=”ws” (or “default”). - kind_dtype: instruction kind selector (e.g., T.float16 for kind::f16,

“tf32” for kind::tf32, “int8” for kind::i8, “float8_e4m3” for kind::f8f6f4).

tilelang.language.tir.op.ptx_tcgen05_mma_ts(kind_dtype, A_ptr, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3, enable_2cta=False)¶

TVM intrinsic for tcgen05.mma tensor-memory x shared-memory instructions.

Expects 13 positional arguments: (kind_dtype, A_ptr, A_offset, desc_b, B_offset, C_ptr, C_offset,

desc_val, scale_out, mask0, mask1, mask2, mask3).

kind_dtype: instruction kind selector (e.g., T.float16 for kind::f16, “tf32” for kind::tf32, “int8” for kind::i8, “float8_e4m3” for kind::f8f6f4).

tilelang.language.tir.op.ptx_tcgen05_mma_blockscaled_ss(kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, sfa_ptr, sfa_offset, sfb_ptr, sfb_offset, reserved0=0, reserved1=0, enable_2cta=False)¶

TVM intrinsic for tcgen05.mma block-scaled (mxf8f6f4.block_scale) instructions.

Block-scaled TCGEN05 is explicit-async and carries an explicit enable_2cta flag, analogous to the regular SS/TS TCGEN05 intrinsics. There is no fallback path if 2CTA is requested.

Positional args: kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, sfa_ptr, sfa_offset, sfb_ptr, sfb_offset, reserved0, reserved1, enable_2cta.

tilelang.language.tir.op.mma_fill(dtype, local_size, local_ptr, offset)¶

TVM intrinsic for zero-initalizing an MMA accumulation register

Parameters:

dtype (str) – The data type of the result.
local_size (IntImm) – The number of elements.
local_ptr (Var) – The destination pointer variable.
offset (Expr) – The destination offset.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_ldmatrix(trans, num, src_access_ptr, dst_access_ptr)¶

TileLang intrinsic for ptx load matrix from shared memory

Uses tl.ptx_ldmatrix which expects access pointers created via T.access_ptr (i.e. tl.access_ptr wrapping a BufferLoad).

https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix

Parameters:

trans (bool) – The matrix is loaded in column-major format.
num (IntImm) – The number of matrices (2 or 4).
src_access_ptr (PrimExpr) – A tl.access_ptr pointing to the source (shared memory) buffer.
dst_access_ptr (PrimExpr) – A tl.access_ptr pointing to the destination (local/register) buffer.

Returns:

call – The call expression (handle-typed).

Return type:

PrimExpr

tilelang.language.tir.op.ptx_cp_async(dst_access_ptr, src_access_ptr, num_elems, predicate=None)¶

TVM intrinsic for ptx async copy from global to shared memory using cp.async https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async

Parameters:

dst_access_ptr (PrimExpr) – The destination (shared memory) access pointer created by tvm_access_ptr. Should include pointer, offset, extent, and write access flag (rw_mask=2).
src_access_ptr (PrimExpr) – The source (global memory) access pointer created by tvm_access_ptr. Should include pointer, offset, extent, and read access flag (rw_mask=1).
num_elems (int or PrimExpr) –
The number of logical elements to copy.

For TileLang’s tl.ptx_cp_async frontend op, the final PTX byte width is derived later from num_elems * element_bits(access_ptr) and must eventually land on a legal cp.async width of 4, 8, or 16 bytes.
predicate (PrimExpr, optional) – Optional predicate condition for conditional cp.async. When provided, the copy will only be performed if the predicate evaluates to true. Otherwise, the destination will be filled with zeros (default behavior of cp.async).

Returns:

call – The call expression.

Return type:

PrimExpr

Examples

>>> # Copy 16 uint8 elements (= 16 bytes) from global to shared memory
>>> T.ptx_cp_async(
...     T.tvm_access_ptr(T.type_annotation(T.uint8), A_shared.data, 0, 16, 2),  # dst
...     T.tvm_access_ptr(T.type_annotation(T.uint8), B_global.data, 0, 16, 1),  # src
...     16  # num_elems
... )
>>>
>>> # Predicated cp.async (only copy if condition is true)
>>> T.ptx_cp_async(
...     T.tvm_access_ptr(T.type_annotation(T.uint8), A_shared.data, 0, 16, 2),
...     T.tvm_access_ptr(T.type_annotation(T.uint8), B_global.data, 0, 16, 1),
...     16,
...     predicate=guard  # only copy if guard is true
... )

tilelang.language.tir.op.tvm_mfma(dtype, shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index)¶

TVM intrinsic for amd matrix core mfma instructions https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma

Parameters:

dtype (str) – The data type of the result.
shape (str) – The shape of mma fragment.
A_layout (Literal["row", "col"]) – The layout of multiplicand fragment A.
B_layout (Literal["row", "col"]) – The layout of multiplicand fragment B.
A_dtype (str) – The data type of multiplicand fragment A.
B_dtype (str) – The data type of multiplicand fragment B.
C_dtype (str) – The data type of accumulator fragment C.
multiplicand_a (Var) – The multiplicand fragment A variable.
a_index (Expr) – The index of multiplicand fragment A.
multiplicand_b (Var) – The multiplicand fragment B variable.
b_index (Expr) – The index of multiplicand fragment A.
accumulator (Var) – The accumulator fragment C variable.
c_index (Expr) – The index of accumulator fragment C.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_mfma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride)¶

TVM intrinsic for storing the result of PTX MMA into a destination pointer

Parameters:

dtype (str) – The data type of the result.
m (IntImm) – The shape of mma fragment.
n (IntImm) – The shape of mma fragment.
dst_ptr (Var) – The destination pointer variable.
src_ptr (Var) – The source pointer variable.
src_offset (Expr) – The source offset.
dst_stride (Var) – The destination stride.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_rdna_wmma(dtype, shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index)¶

TVM intrinsic for amd matrix core mfma instructions https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma

Parameters:

dtype (str) – The data type of the result.
shape (str) – The shape of mma fragment.
A_layout (Literal["row", "col"]) – The layout of multiplicand fragment A.
B_layout (Literal["row", "col"]) – The layout of multiplicand fragment B.
A_dtype (str) – The data type of multiplicand fragment A.
B_dtype (str) – The data type of multiplicand fragment B.
C_dtype (str) – The data type of accumulator fragment C.
multiplicand_a (Var) – The multiplicand fragment A variable.
a_index (Expr) – The index of multiplicand fragment A.
multiplicand_b (Var) – The multiplicand fragment B variable.
b_index (Expr) – The index of multiplicand fragment A.
accumulator (Var) – The accumulator fragment C variable.
c_index (Expr) – The index of accumulator fragment C.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_rdna_wmma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride)¶

TVM intrinsic for storing the result of PTX MMA into a destination pointer

Parameters:

dtype (str) – The data type of the result.
m (IntImm) – The shape of mma fragment.
n (IntImm) – The shape of mma fragment.
dst_ptr (Var) – The destination pointer variable.
src_ptr (Var) – The source pointer variable.
src_offset (Expr) – The source offset.
dst_stride (Var) – The destination stride.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_fence_barrier_init()¶

TVM intrinsic for ptx fence barrier initialization.

Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.ptx_arrive_barrier_expect_tx(barrier_id, byte_count)¶

TVM intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-expect-tx-operation

Parameters:

barrier_id (int) – The ID of the barrier shared memory pointer.
byte_count (int) – Increases the tx count of the mbarrier object to track completion of additional async transactions.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.infinity(dtype, span=None)¶

infinity value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The infinity value of dtype.

Return type:

tvm.Expr

tilelang.language.tir.op.reinterpret(value, dtype, span=None)¶

Reinterpret cast a value to dtype.

Parameters:

value (PrimExpr) – The input value.
dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The reinterpret cast value of dtype.

Return type:

tvm.Expr

tilelang.language.tir.op.round(x, rounding_mode='ties-to-even', span=None)¶

Round elements of the array to the nearest integer.

Parameters:

x (PrimExpr) – Input argument.
rounding_mode (str) – Rounding mode to use. Supported values are "ties-to-even" and "ties-away-from-zero". "ties-to-even" is the default and matches the existing TileLang/TVM semantics.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.pow_of_int(x, y)¶

Fast power operation than pow(float, float).

Parameters:

x (PrimExpr) – Base value
y (int) – Exponent value

Return type:

tvm.ir.PrimExpr

tilelang.language.tir.op.pow(x, y, span=None)¶

x power y

Parameters:

x (PrimExpr) – Input argument.
y (PrimExpr) – The exponent
span (Optional[Span]) – The location of this operator in the source code.

Returns:

z – The result.

Return type:

PrimExpr