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_commit_group`
`ptx_wait_group`
`ptx_cp_async_barrier`
`ptx_init_barrier_thread_count`
`ptx_arrive_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_cp_async`(dst_access_ptr, src_access_ptr, num_elems)	TVM intrinsic for ptx async copy from global to shared memory using cp.async
`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_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.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_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.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