tilelang.language.tir.op¶

Attributes¶

`sum`
`min`
`max`

Functions¶

`call_packed`(*args[, span])	Build expression by call an external packed function.
`call_cpacked`(*args[, span])	Build expression by call an external packed function.
`call_packed_lowered`(*args[, span])	Lowered version of call packed.
`call_cpacked_lowered`(*args[, span])	Lowered version of call c-packed.
`call_intrin`(dtype, func_name, *args[, 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.
`call_llvm_intrin`(dtype, name, *args[, span])	Build expression by calling a llvm intrinsic function
`call_llvm_pure_intrin`(dtype, name, *args[, span])	Build expression by calling a pure llvm intrinsic function
`tvm_check_return`(expected, return_unexpected, nested_call)	Return new on stack dtype[num]
`tvm_stack_alloca`(dtype_str, num)	Return new on stack dtype[num]
`tvm_stack_make_shape`(*args)	Allocate a shape tuple on stack, return the handle
`tvm_stack_make_array`(data, shape, strides, ndim, ...)	Allocate a NDArray(DLTensor) on stack, return the handle
`assume`([cond])	Provide a true statement that can be used for simplifications
`undef`()	Returns an initialized but arbitrary value
`call_tir`(global_var, *args)	Performs a call into another PrimFunc in the same IRModule
`start_profile_intrinsic`(id)	Start profile intrinsic.
`end_profile_intrinsic`(id)	End profile intrinsic.
`tvm_tuple`(*value)	Create a tuple structure in value field of AttrStmt
`tvm_struct_get`(arr, index, field, dtype)	Get struct field value in array
`tvm_struct_set`(arr, index, field, value)	Set value in struct field in array
`address_of`(buffer_load[, span])	Returns the address of an element in the buffer
`lookup_param`(param_name[, span])	Returns the param by name
`tvm_thread_allreduce`(*freduce_args)	Perform allreduce inside threadblock.
`tvm_thread_invariant`(cond)	Mark condition as thread invariant.
`tvm_storage_sync`(storage_scope)	Perform synchronization in specified scope.
`tvm_warp_shuffle`(mask, value, warp_id, width, warp_size)	Exchange value between threads inside a warp.
`tvm_warp_shuffle_up`(mask, value, offset, width, warp_size)	Copy value from a lane with lower (by offset) index relative to caller.
`tvm_warp_shuffle_down`(mask, value, offset, width, ...)	Copy value from a lane with higher (by offset) index relative to caller.
`tvm_warp_activemask`()	Return a 32-bit mask indicates currently active threads in a calling warp.
`type_annotation`(dtype)	Create a type annotation expression
`tvm_access_ptr`(ptype, data, offset, extent, rw_mask)	Get head access address with memory access pattern info
`tvm_throw_last_error`()	Throw TVMGetLastError()
`tvm_load_matrix_sync`(fragment, m, n, k, index, ...)	TVM intrinsic for tensor core load operators
`tvm_mma_sync`(fragment_d, index_d, fragment_a, index_a, ...)	TVM intrinsic for tensor core mma_sync operators
`tvm_bmma_sync`(fragment_d, index_d, fragment_a, ...)	TVM intrinsic for tensor core bmma_sync operators
`tvm_fill_fragment`(fragment, m, n, k, index, value)	TVM intrinsic for tensor core fill_fragment operators
`tvm_store_matrix_sync`(fragment, m, n, k, index, ...)	TVM intrinsic for tensor core store operators
`ptx_mma`(dtype, shape, A_layout, B_layout, A_dtype, ...)	TVM intrinsic for ptx tensor core mma instructions
`ptx_mma_sp`(dtype, shape, A_layout, B_layout, A_dtype, ...)	TVM intrinsic for sparse tensor core ptx instructions
`mma_store`(dtype, m, n, dst_ptr, src_ptr, src_offset, ...)	TVM intrinsic for storing the result of PTX MMA into a destination pointer
`mma_fill`(dtype, local_size, local_ptr, offset)	TVM intrinsic for zero-initalizing an MMA accumulation register
`ptx_ldmatrix`(dtype, trans, num, type, local_ptr, ...)	TVM intrinsic for ptx load matrix from shared memory
`ptx_cp_async`(dtype, shared_ptr, shared_offset, ...)	TVM intrinsic for ptx async copy from global to shared memory using cp.async
`ptx_cp_async_bulk`(dtype, shared_ptr, shared_offset, ...)	TVM intrinsic for ptx async copy from global to shared memory using cp.async.bulk
`ptx_commit_group`()	TVM intrinsic for ptx async copy commit
`ptx_wait_group`(num)	TVM intrinsic for ptx async copy wait
`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_cp_async_barrier`(barrier_id)	TVM intrinsic for ptx async copy barrier using cp.async.mbarrier.arrive
`ptx_init_barrier_thread_count`(barrier_id, thread_count)	TVM intrinsic for ptx barrier initialization of thread count using mbarrier.init
`ptx_arrive_barrier`(barrier_id)	TVM intrinsic for ptx barrier arrival using mbarrier.arrive
`ptx_arrive_barrier_expect_tx`(barrier_id, byte_count)	TVM intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx
`ptx_wait_barrier`(barrier_id)	TVM intrinsic for ptx barrier wait using mbarrier.try_wait
`create_barriers`(barrier_count)	TVM intrinsic to create N barriers
`vectorlow`(dtype, vec)	Get the low level half of the vector
`vectorhigh`(dtype, vec)	Get the high level half of the vector
`vectorcombine`(dtype, vec1, vec2)	Concat two vectors
`ret`(val)	Create a tir return expression
`any`(*args[, span])	Create a new expression of the union of all conditions in the arguments
`all`(*args[, span])	Create a new expression of the intersection of all conditions in the
`trace`(args[, trace_action])	Trace tensor data at the runtime.
`min_value`(dtype[, span])	minimum value of dtype
`max_value`(dtype[, span])	maximum value of dtype
`infinity`(dtype[, span])	infinity value of dtype
`reinterpret`(dtype, value[, span])	infinity value of dtype
`exp`(x)	Take exponential of input x.
`exp2`(x)	Calculate 2**x
`exp10`(x)	Calculate 10**x
`erf`(x)	Take gauss error function of the input x.
`tanh`(x)	Take hyperbolic tanh of input x.
`sigmoid`(x)	Quick function to get sigmoid
`log`(x)	Take log of input x.
`log2`(x)	Take log2 of input x.
`log10`(x)	Take log10 of input x.
`log1p`(x)	Take log(x + 1) with respect to input x.
`tan`(x)	Take tan of input x.
`cos`(x)	Take cos of input x.
`cosh`(x)	Take cosh of input x.
`acos`(x)	Take acos of input x.
`acosh`(x)	Take acos of input x.
`sin`(x)	Take sin of input x.
`sinh`(x)	Take sinh of input x.
`asin`(x)	Take asin of input x.
`asinh`(x)	Take asinh of input x.
`atan`(x)	Take atan of input x.
`atanh`(x)	Take atanh of input x.
`atan2`(x1, x2)	Take arctan2(x1, x2).
`sqrt`(x)	Take square root of input x.
`rsqrt`(x)	Take reciprocal of square root of input x.
`clz`(x)	Count leading zero bits of an integer x.
`floor`(x[, span])	Take floor of float input x.
`ceil`(x[, span])	Take ceil of float input x.
`trunc`(x[, span])	Get truncated value of the input.
`abs`(x[, span])	Get absolute value of the input element-wise.
`bitwise_and`(x, y[, span])	Take bitwise and of two values
`bitwise_not`(x[, span])	Take bitwise not of input value
`bitwise_or`(x, y[, span])	Take bitwise or of two values
`bitwise_xor`(x, y[, span])	Take bitwise xor of two values
`round`(x[, span])	Round elements of the array to the nearest integer.
`nearbyint`(x[, span])	Round elements of the array to the nearest integer.
`nextafter`(x1, x2)	Return the next floating-point value after x1 towards x2.
`hypot`(x1, x2)	Equivalent to sqrt(x12 + x22), element-wise.
`copysign`(x1, x2)	Change the sign of x1 to that of x2, element-wise.
`ldexp`(x1, x2)	Returns x1 * (2 ** x2).
`likely`(cond[, span])	Mark condition as likely.
`isnan`(x[, span])	Check if input value is Nan.
`isnullptr`(x[, span])	Check if input value is nullptr.
`isfinite`(x[, span])	Check if input value is finite.
`isinf`(x[, span])	Check if input value is infinite.
`pow_of_int`(x, y)	Fast power operation than pow(float, float).
`power`(x, y[, span])	x power y
`pow`(x, y[, span])	x power y
`popcount`(x)	Count the number of set bits in input x.
`q_multiply_shift`(x, y, q, s)	Execute a multiplication between two Q-numbers x and y
`q_multiply_shift_per_axis`(x, y, ls, rs, q, ...)	Execute a multiplication between two Q-numbers x and y
`shift_left`(x, y[, span])	Return the result of x left shifted by y bits.
`shift_right`(x, y[, span])	Return the result of x right shifted by y bits.
`fmod`(x, y)	Return the remainder of x divided by y with the same sign as x.
`if_then_else`(cond, t, f[, span])	Conditional selection expression.
`div`(a, b[, span])	Compute a / b as in C/C++ semantics.
`indexdiv`(a, b[, span])	Compute floor(a / b) where a and b are non-negative.
`indexmod`(a, b[, span])	Compute the remainder of indexdiv. a and b are non-negative.
`truncdiv`(a, b[, span])	Compute the truncdiv of two expressions.
`truncmod`(a, b[, span])	Compute the truncmod of two expressions.
`floordiv`(a, b[, span])	Compute the floordiv of two expressions.
`floormod`(a, b[, span])	Compute the floormod of two expressions.
`ceildiv`(lhs, rhs[, span])	Generic ceildiv operator.
`comm_reducer`(fcombine, fidentity[, name])	Create a commutative reducer for reduction.
`TVMBackendAllocWorkspace`(device_type, device_id, ...)	Backend function to allocate temporal workspace
`TVMBackendFreeWorkspace`(device_type, device_id, ptr)	Backend function to free temporal workspace.
`anylist_getitem`(list_handle, index)	Returns an item from any list.
`anylist_resetitem`(list_handle, index)	Reset an item from any list.
`anylist_setitem_call_packed`(list_handle, index, ...)	Set anylist item by result of packed call.
`anylist_setitem_call_cpacked`(list_handle, index, ...)	Set anylist item by result of packed call.
`vscale`()	Get the target's vscale value. It will be lowered to llvm.vscale intrinsic

Module Contents¶

tilelang.language.tir.op.call_packed(*args, span=None)¶

Build expression by call an external packed function.

The argument to packed function can be Expr or Buffer. The argument is the corresponding POD type when Expr is presented.

When the argument is Buffer, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is NDArray.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

See also

te.extern: Create tensor with extern function call.

tilelang.language.tir.op.call_cpacked(*args, span=None)¶

Build expression by call an external packed function.

Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

See also

te.extern: Create tensor with extern function call.

tilelang.language.tir.op.call_packed_lowered(*args, span=None)¶

Lowered version of call packed. The argument to packed function can be Expr or Buffer. The argument is the corresponding POD type when Expr is presented. When the argument is Buffer, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is NDArray.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

See also

te.extern: Create tensor with extern function call.

tilelang.language.tir.op.call_cpacked_lowered(*args, span=None)¶

Lowered version of call c-packed. Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

See also

te.extern: Create tensor with extern function call.

tilelang.language.tir.op.call_intrin(dtype, func_name, *args, 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.call_llvm_intrin(dtype, name, *args, span=None)¶

Build expression by calling a llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
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_llvm_pure_intrin(dtype, name, *args, span=None)¶

Build expression by calling a pure llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
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_check_return(expected, return_unexpected, nested_call)¶

Return new on stack dtype[num] :param expected: The expected return code. :type expected: int :param return_unexpected: The unexpected return code. :type return_unexpected: int :param nested_call: The call expression to check return. :type nested_call: PrimExpr

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

tilelang.language.tir.op.tvm_stack_alloca(dtype_str, num)¶

Return new on stack dtype[num]

Parameters:

dtype_str (str) – The data type of array.
num (int) – The size of array.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_stack_make_shape(*args)¶

Allocate a shape tuple on stack, return the handle

Parameters:: args (int) – The tuple shape.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset)¶

Allocate a NDArray(DLTensor) on stack, return the handle

Parameters:

data (Expr) – The data of array.
shape (Expr) – The shape of array.
strides (Expr) – The strides of array.
ndim (Expr) – The dimensions of array.
arr_dtype (Expr) – The data type of array.
elem_offse (Expr) – The element offset of array.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.assume(cond=None)¶

Provide a true statement that can be used for simplifications

Parameters:: cond (Expr) – The constraint condition.
Returns:: call – The call expression.
Return type:: PrimExpr

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

Returns an initialized but arbitrary value

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

tilelang.language.tir.op.call_tir(global_var, *args)¶

Performs a call into another PrimFunc in the same IRModule

Returns:: call – The call expression.
Return type:: PrimExpr
Parameters:: global_var (tvm.ir.GlobalVar)

tilelang.language.tir.op.start_profile_intrinsic(id)¶

Start profile intrinsic. :param id: The intrinsic id. :type id: int

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

tilelang.language.tir.op.end_profile_intrinsic(id)¶

End profile intrinsic. :param id: The intrinsic id. :type id: int

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

tilelang.language.tir.op.tvm_tuple(*value)¶

Create a tuple structure in value field of AttrStmt

Parameters:: value (Expr) – The value in tuple.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.tvm_struct_get(arr, index, field, dtype)¶

Get struct field value in array

Parameters:

dtype (str) – The date type of the result.
arr (StructType*) – The array of struct.
index (int) – The index of struct.
field (int) – The field of struct.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_struct_set(arr, index, field, value)¶

Set value in struct field in array

Parameters:

arr (StructType*) – The array of struct.
index (int) – The index of struct.
field (int) – The field of struct.
value (Expr) – The value to be set in field.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.address_of(buffer_load, span=None)¶

Returns the address of an element in the buffer

Parameters:

buffer_load (BufferLoad) – The buffer load.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.lookup_param(param_name, span=None)¶

Returns the param by name

Parameters:

param_name (str) – The name of param.
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_thread_allreduce(*freduce_args)¶

Perform allreduce inside threadblock.

Parameters:: freduce_args (Expr) – The args.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.tvm_thread_invariant(cond)¶

Mark condition as thread invariant.

Parameters:: cond (Expr) – The condition.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.tvm_storage_sync(storage_scope)¶

Perform synchronization in specified scope.

Parameters:: storage_scope (str) – The storage scope to perform synchronization.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.tvm_warp_shuffle(mask, value, warp_id, width, warp_size)¶

Exchange value between threads inside a warp.

Parameters:

mask (PrimExpr) – The warp mask indicates active threads inside warp.
value (PrimExpr) – The value to exchange.
warp_id (PrimExpr) – The source lane index to fetch value.
width (PrimExpr) – The width of sub-sections to perform warp shuffle.
warp_size (PrimExpr) – The warp size.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_warp_shuffle_up(mask, value, offset, width, warp_size)¶

Copy value from a lane with lower (by offset) index relative to caller.

Parameters:

mask (PrimExpr) – The warp mask indicates active threads inside warp.
value (PrimExpr) – The value to exchange.
offset (PrimExpr) – The difference between source lane index and destination lane index: offset = dst_lane_idx - src_lane_idx
width (PrimExpr) – The width of sub-sections to perform warp shuffle.
warp_size (PrimExpr) – The warp size.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_warp_shuffle_down(mask, value, offset, width, warp_size)¶

Copy value from a lane with higher (by offset) index relative to caller.

Parameters:

mask (PrimExpr) – The warp mask indicates active threads inside warp.
value (PrimExpr) – The value to exchange.
offset (PrimExpr) – The difference between source lane index and destination lane index: offset = src_lane_idx - dst_lane_idx
width (PrimExpr) – The width of sub-sections to perform warp shuffle.
warp_size (PrimExpr) – The warp size.

Returns:

call – The call expression.

Return type:

PrimExpr

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

Return a 32-bit mask indicates currently active threads in a calling warp.

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

tilelang.language.tir.op.type_annotation(dtype)¶

Create a type annotation expression

Parameters:: dtype (Expr) – The data type.
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.tvm_throw_last_error()¶

Throw TVMGetLastError()

Returns:: ret – The return expression
Return type:: PrimExpr

tilelang.language.tir.op.tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout)¶

TVM intrinsic for tensor core load operators

Parameters:

fragment (Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_mma_sync(fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c)¶

TVM intrinsic for tensor core mma_sync operators

Parameters:

fragment_d (Var) – The wmma fragment_d.
index_d (Expr) – The fragment_d index.
fragment_a (Var) – The wmma fragment_a.
index_a (Expr) – The fragment_a index.
fragment_b (Var) – The wmma fragment_b.
index_b (Expr) – The fragment_b index.
fragment_c (Var) – The wmma fragment_c.
index_c (Expr) – The fragment_c index.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_bmma_sync(fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c)¶

TVM intrinsic for tensor core bmma_sync operators

Parameters:

fragment_d (Var) – The bwmma fragment_d.
index_d (Expr) – The fragment_d index.
fragment_a (Var) – The bwmma fragment_a.
index_a (Expr) – The fragment_a index.
fragment_b (Var) – The bwmma fragment_b.
index_b (Expr) – The fragment_b index.
fragment_c (Var) – The bwmma fragment_c.
index_c (Expr) – The fragment_c index.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_fill_fragment(fragment, m, n, k, index, value)¶

TVM intrinsic for tensor core fill_fragment operators

Parameters:

fragment (Var) – The wmma fragment
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
value (Expr) – The value to be filled in fragment.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.tvm_store_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout)¶

TVM intrinsic for tensor core store operators

Parameters:

fragment (Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_mma(dtype, 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)¶

TVM intrinsic for ptx tensor core mma 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.
saturate (bool) – The optional saturation at the output.
operator (Optional[Literal["xor", "and"]]) – The 1-bit operator.

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.mma_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.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(dtype, trans, num, type, local_ptr, local_offset, smem_ptr, smem_offset)¶

TVM intrinsic for ptx load matrix from shared memory https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix

Parameters:

dtype (str) – The data type of the result.
trans (bool) – The matrix is loaded in column-major format.
num (IntImm) – The number of matrices.
type (Literal[".b16"]) – The data type of the matrices.
local_ptr (Var) – The local pointer variable.
local_offset (Expr) – The offset of local pointer.
smem_ptr (Var) – The shared memory pointer variable.
smem_offset (Expr) – The offset of shared memort pointer.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_cp_async(dtype, shared_ptr, shared_offset, global_ptr, global_offset, bytes)¶

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:

dtype (str) – The data type of the result.
shared_ptr (Var) – The shared memory pointer variable.
shared_offset (Expr) – The offset of shared memory pointer.
global_ptr (Var) – The global memory pointer variable.
global_offset (Expr) – The offset of global memory pointer.
bytes (int) – The data size to copy.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_cp_async_bulk(dtype, shared_ptr, shared_offset, global_ptr, global_offset, bytes, barrier_id)¶

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

Parameters:

dtype (str) – The data type of the result.
shared_ptr (Var) – The shared memory pointer variable.
shared_offset (Expr) – The offset of shared memory pointer.
global_ptr (Var) – The global memory pointer variable.
global_offset (Expr) – The offset of global memory pointer.
bytes (int) – The data size to copy.
barrier_id (int) – The ID of the barrier shared memory pointer.

Returns:

call – The call expression.

Return type:

PrimExpr

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

TVM intrinsic for ptx async copy commit https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-commit-group

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

tilelang.language.tir.op.ptx_wait_group(num)¶

TVM intrinsic for ptx async copy wait https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-wait-group

Parameters:: num (int) – The number of the most recent uncommitted pending cp.async groups to wait.
Returns:: call – The call expression.
Return type:: PrimExpr

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_cp_async_barrier(barrier_id)¶

TVM intrinsic for ptx async copy barrier using cp.async.mbarrier.arrive https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive

Parameters:: barrier_id (int) – The ID of the barrier shared memory pointer.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.ptx_init_barrier_thread_count(barrier_id, thread_count)¶

TVM intrinsic for ptx barrier initialization of thread count using mbarrier.init https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init

Parameters:

barrier_id (int) – The ID of the barrier shared memory pointer.
thread_count (int) – Number of threads expected to arrive at the barrier.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ptx_arrive_barrier(barrier_id)¶

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

Parameters:: barrier_id (int) – The ID of the barrier shared memory pointer.
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.ptx_wait_barrier(barrier_id)¶

TVM intrinsic for ptx barrier wait using mbarrier.try_wait https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait

Parameters:: barrier_id (int) – The ID of the barrier shared memory pointer.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.create_barriers(barrier_count)¶

TVM intrinsic to create N barriers

Parameters:: barrier_count (int) – The number of barriers to create.
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.vectorlow(dtype, vec)¶

Get the low level half of the vector

Parameters:

dtype (str) – The data type of the result.
vec (list) – The input vector.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.vectorhigh(dtype, vec)¶

Get the high level half of the vector

Parameters:

dtype (str) – The data type of the result.
vec (list) – The input vector.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.vectorcombine(dtype, vec1, vec2)¶

Concat two vectors

Parameters:

vec1 (list) – The input vector.
vec2 (list) – The input vector.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.ret(val)¶

Create a tir return expression

Parameters:: val (Expr) – The returned tir expression, whose data type is int, float or void pointer.
Returns:: ret – The return expression
Return type:: PrimExpr

tilelang.language.tir.op.any(*args, span=None)¶

Create a new expression of the union of all conditions in the arguments

Parameters:

args (list) – List of symbolic boolean expressions
span (Optional[Span]) – The location of this operator in the source code.

Returns:

expr – Expression

Return type:

Expr

tilelang.language.tir.op.all(*args, span=None)¶

Create a new expression of the intersection of all conditions in the: arguments

Parameters:

args (list) – List of symbolic boolean expressions
span (Optional[Span]) – The location of this operator in the source code.

Returns:

expr – Expression

Return type:

Expr

tilelang.language.tir.op.trace(args, trace_action='tvm.default_trace_action')¶

Trace tensor data at the runtime.

The trace function allows to trace specific tensor at the runtime. The tracing value should come as last argument. The trace action should be specified, by default tvm.default_trace_action is used.

Parameters:

args (list of Expr or Buffers.) – Positional arguments.
trace_action (str.) – The name of the trace action.

Returns:

call – The call expression.

Return type:

PrimExpr

See also

tvm.tir.call_packed: Creates packed function.

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

minimum value of dtype

Parameters:

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

Returns:

value – The minimum value of dtype.

Return type:

tvm.Expr

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

maximum value of dtype

Parameters:

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

Returns:

value – The maximum value of dtype.

Return type:

tvm.Expr

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(dtype, value, span=None)¶

infinity value of dtype

Parameters:

dtype (str) – The data type.
value (PrimExpr) – The input value.
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.exp(x)¶

Take exponential of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.exp2(x)¶

Calculate 2**x

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.exp10(x)¶

Calculate 10**x

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.erf(x)¶

Take gauss error function of the input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.tanh(x)¶

Take hyperbolic tanh of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.sigmoid(x)¶

Quick function to get sigmoid

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.log(x)¶

Take log of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.log2(x)¶

Take log2 of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.log10(x)¶

Take log10 of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.log1p(x)¶

Take log(x + 1) with respect to input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.tan(x)¶

Take tan of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.cos(x)¶

Take cos of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.cosh(x)¶

Take cosh of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.acos(x)¶

Take acos of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.acosh(x)¶

Take acos of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.sin(x)¶

Take sin of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.sinh(x)¶

Take sinh of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.asin(x)¶

Take asin of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.asinh(x)¶

Take asinh of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.atan(x)¶

Take atan of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.atanh(x)¶

Take atanh of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.atan2(x1, x2)¶

Take arctan2(x1, x2).

Parameters:

x1 (PrimExpr) – Input argument.
x2 (PrimExpr) – Input argument.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.sqrt(x)¶

Take square root of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.rsqrt(x)¶

Take reciprocal of square root of input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.clz(x)¶

Count leading zero bits of an integer x.

Parameters:: x (PrimExpr) – Input 32 or 64 bit integer. The result is undefined if the input is 0.
Returns:: y – The result.
Return type:: PrimExpr

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

Take floor of float input x.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Take ceil of float input x.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Get truncated value of the input.

The truncated value of the scalar x is the nearest integer i which is closer to zero than x is.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Get absolute value of the input element-wise.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Take bitwise and of two values

Parameters:

x (PrimExpr) – Left operand
y (PrimExpr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

PrimExpr

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

Take bitwise not of input value

Parameters:

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

Returns:

res – The result.

Return type:

PrimExpr

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

Take bitwise or of two values

Parameters:

x (PrimExpr) – Left operand
y (PrimExpr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

PrimExpr

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

Take bitwise xor of two values

Parameters:

x (PrimExpr) – Left operand
y (PrimExpr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

PrimExpr

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

Round elements of the array to the nearest integer.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Round elements of the array to the nearest integer. This intrinsic uses llvm.nearbyint instead of llvm.round which is faster but will results different from te.round. Notably nearbyint rounds according to the rounding mode, whereas te.round (llvm.round) ignores that. For differences between the two see: https://en.cppreference.com/w/cpp/numeric/math/round https://en.cppreference.com/w/cpp/numeric/math/nearbyint

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.nextafter(x1, x2)¶

Return the next floating-point value after x1 towards x2.

Parameters:

x1 (PrimExpr) – Input argument.
x2 (PrimExpr) – Input argument.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.hypot(x1, x2)¶

Equivalent to sqrt(x1**2 + x2**2), element-wise.

Parameters:

x1 (PrimExpr) – Input argument.
x2 (PrimExpr) – Input argument.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.copysign(x1, x2)¶

Change the sign of x1 to that of x2, element-wise.

Parameters:

x1 (PrimExpr) – Input argument.
x2 (PrimExpr) – Input argument.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.ldexp(x1, x2)¶

Returns x1 * (2 ** x2).

Parameters:

x1 (PrimExpr) – Input argument.
x2 (PrimExpr) – Input argument.

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.likely(cond, span=None)¶

Mark condition as likely.

Parameters:

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

Returns:

y – The marked expression.

Return type:

PrimExpr

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

Check if input value is Nan.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Check if input value is nullptr.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Check if input value is finite.

Parameters:

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

Returns:

y – The result.

Return type:

PrimExpr

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

Check if input value is infinite.

Parameters:

x (PrimExpr) – Input argument.
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.power(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

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

tilelang.language.tir.op.popcount(x)¶

Count the number of set bits in input x.

Parameters:: x (PrimExpr) – Input argument.
Returns:: y – The result.
Return type:: PrimExpr

tilelang.language.tir.op.q_multiply_shift(x, y, q, s)¶

Execute a multiplication between two Q-numbers x and y followed by a right shift s. The mathematical expression is:

out = round(x*y*2^-s)

More about Q-numbers here: https://en.wikipedia.org/wiki/Q_(number_format) The rounding rule is to the nearest value, rounding half up (i.e., round(x.1) = x and round (x.5) = x+1)

Parameters:

x (PrimExpr) – First Q-number
y (PrimExpr) – Second Q-number
q (PrimExpr) – Number of fractional bits in x and y. Needs to be > 0
s (PrimExpr) – Integer shift

Returns:

y – The result.

Return type:

PrimExpr

tilelang.language.tir.op.q_multiply_shift_per_axis(x, y, ls, rs, q, is_lshift_required, is_rshift_required)¶

Execute a multiplication between two Q-numbers x and y

Parameters:

x (PrimExpr) – First Q-number.
y (PrimExpr) – Second Q-number.
ls (PrimExpr) – Integer left shift.
rs (PrimExpr) – Integer right shift.
q (IntImm) – Number of fractional bits in x and y. Needs to be > 0.
is_lshift_required (IntImm) – Whether we need to do left shift or not.
is_rshift_required (IntImm) – Whether we need to do right shift or not.

Returns:

z – The result.

Return type:

PrimExpr

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

Return the result of x left shifted by y bits.

Parameters:

x (PrimExpr) – Input argument.
y (PrimExpr) – Input argument.

Returns:

z – The result.

Return type:

PrimExpr

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

Return the result of x right shifted by y bits.

Parameters:

x (PrimExpr) – Input argument.
y (PrimExpr) – Input argument.

Returns:

z – The result.

Return type:

PrimExpr

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

Return the remainder of x divided by y with the same sign as x.

Parameters:

x (PrimExpr) – Input argument.
y (PrimExpr) – Input argument.

Returns:

z – The result.

Return type:

PrimExpr

tilelang.language.tir.op.if_then_else(cond, t, f, span=None)¶

Conditional selection expression.

Parameters:

cond (PrimExpr) – The condition
t (PrimExpr) – The result expression if cond is true.
f (PrimExpr) – The result expression if cond is false.
span (Optional[Span]) – The location of this operator in the source.

Returns:

result – The result of conditional expression.

Return type:

Node

Note

Unlike Select, if_then_else will not execute the branch that does not satisfy the condition. You can use it to guard against out of bound access. Unlike Select, if_then_else cannot be vectorized if some lanes in the vector have different conditions.

tilelang.language.tir.op.div(a, b, span=None)¶

Compute a / b as in C/C++ semantics.

Parameters:

a (PrimExpr) – The left hand operand, known to be non-negative.
b (PrimExpr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

Note

When operands are integers, returns truncdiv(a, b, span).

tilelang.language.tir.op.indexdiv(a, b, span=None)¶

Compute floor(a / b) where a and b are non-negative.

Parameters:

a (PrimExpr) – The left hand operand, known to be non-negative.
b (PrimExpr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

Note

Use this function to split non-negative indices. This function may take advantage of operands’ non-negativeness.

tilelang.language.tir.op.indexmod(a, b, span=None)¶

Compute the remainder of indexdiv. a and b are non-negative.

Parameters:

a (PrimExpr) – The left hand operand, known to be non-negative.
b (PrimExpr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

Note

Use this function to split non-negative indices. This function may take advantage of operands’ non-negativeness.

tilelang.language.tir.op.truncdiv(a, b, span=None)¶

Compute the truncdiv of two expressions.

Parameters:

a (PrimExpr) – The left hand operand
b (PrimExpr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

Note

This is the default integer division behavior in C.

tilelang.language.tir.op.truncmod(a, b, span=None)¶

Compute the truncmod of two expressions.

Parameters:

a (PrimExpr) – The left hand operand
b (PrimExpr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

Note

This is the default integer division behavior in C.

tilelang.language.tir.op.floordiv(a, b, span=None)¶

Compute the floordiv of two expressions.

Parameters:

a (PrimExpr) – The left hand operand
b (PrimExpr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

tilelang.language.tir.op.floormod(a, b, span=None)¶

Compute the floormod of two expressions.

Parameters:

a (PrimExpr) – The left hand operand
b (PrimExpr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

PrimExpr

tilelang.language.tir.op.ceildiv(lhs, rhs, span=None)¶

Generic ceildiv operator.

Parameters:

lhs (object) – The left operand.
rhs (object) – The right operand.
span (Optional[Span]) – The location of this operator in the source.

Returns:

op – The result Expr of ceildiv operation.

Return type:

tvm.Expr

tilelang.language.tir.op.comm_reducer(fcombine, fidentity, name='reduce')¶

Create a commutative reducer for reduction.

Parameters:

fcombine (function(Expr -> Expr -> Expr)) – A binary function which takes two Expr as input to return a Expr.
fidentity (function(str -> Expr)) – A function which takes a type string as input to return a const Expr.

Returns:

reducer – A function which creates a reduce expression over axis. There are two ways to use it:

accept (expr, axis, where) to produce an Reduce Expr on specified axis;
simply use it with multiple Exprs.

Return type:

function

Example

n = te.var("n")
m = te.var("m")
mysum = te.comm_reducer(lambda x, y: x+y,
    lambda t: tvm.tir.const(0, dtype=t), name="mysum")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), name="k")
B = te.compute((n,), lambda i: mysum(A[i, k], axis=k), name="B")

tilelang.language.tir.op.TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint)¶

Backend function to allocate temporal workspace

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
nbytes (int) – The size of the space requested.
dtype_code_hint (int) – The type code of the array elements. Only used in certain backends such as OpenGL.
dtype_bits_hint (int) – The type bits of the array elements. Only used in certain backends such as OpenGL.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.TVMBackendFreeWorkspace(device_type, device_id, ptr)¶

Backend function to free temporal workspace.

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
ptr (Var) – The result allocated space pointer.

Returns:

call – The call expression.

Return type:

PrimExpr

tilelang.language.tir.op.anylist_getitem(list_handle, index)¶

Returns an item from any list. list_handle: Var

The handle to anylist

indexint: The index

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

tilelang.language.tir.op.anylist_resetitem(list_handle, index)¶

Reset an item from any list. list_handle: Var

The handle to anylist

indexint: The index

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

tilelang.language.tir.op.anylist_setitem_call_packed(list_handle, index, func_name, *args)¶

Set anylist item by result of packed call. list_handle: Var

The handle to anylist

indexint: The index
func_name: str: The name of the function to be called.

Parameters:: arguments (Extra)
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.anylist_setitem_call_cpacked(list_handle, index, func_name, *args)¶

Set anylist item by result of packed call. list_handle: Var

The handle to anylist

indexint: The index
func_name: str: The name of the function to be called.

Parameters:: arguments (Extra)
Returns:: call – The call expression.
Return type:: PrimExpr

tilelang.language.tir.op.vscale()¶: Get the target’s vscale value. It will be lowered to llvm.vscale intrinsic (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic) :returns: call – Call to the vscale intrinsic :rtype: PrimExpr

tilelang.language.tir.op.sum¶

tilelang.language.tir.op.min¶

tilelang.language.tir.op.max¶