tilelang.utils.language

Functions

is_global(buffer)	Check if the buffer is in the global memory scope.
is_shared(buffer[, allow_dynamic])	Check if the buffer is in the shared memory scope.
is_shared_dynamic(buffer)	Check if the buffer is in the dynamic shared memory scope.
is_local(buffer)	Check if the buffer is in the local memory scope.
is_fragment(buffer)	Check if the buffer is a fragment (e.g., for matrix multiplication operations).
get_buffer_elems(buffer)	Get the number of elements in the buffer.
array_reduce(array)	Reduce an array of integers to a single integer.
retrieve_func_from_module(ir_module)	Retrieve the single PrimFunc from an IRModule.
get_buffer_region_from_load(buffer_load)	Get the buffer region from a buffer load.

Module Contents

tilelang.utils.language.is_global(buffer)

Check if the buffer is in the global memory scope.

Parameters:

buffer (Buffer) – The TVM buffer to check.

Returns:

True if the buffer is in global memory, False otherwise.

Return type:

bool

tilelang.utils.language.is_shared(buffer, allow_dynamic=True)

Check if the buffer is in the shared memory scope.

Parameters:

• buffer (Buffer) – The TVM buffer to check.

• allow_dynamic (bool)

Returns:

True if the buffer is in shared memory, False otherwise.

Return type:

bool

tilelang.utils.language.is_shared_dynamic(buffer)

Check if the buffer is in the dynamic shared memory scope.

Parameters:

buffer (Buffer) – The TVM buffer to check.

Returns:

True if the buffer is in dynamic shared memory, False otherwise.

Return type:

bool

tilelang.utils.language.is_local(buffer)

Check if the buffer is in the local memory scope.

Parameters:

buffer (Buffer) – The TVM buffer to check.

Returns:

True if the buffer is in local memory, False otherwise.

Return type:

bool

tilelang.utils.language.is_fragment(buffer)

Check if the buffer is a fragment (e.g., for matrix multiplication operations).

Parameters:

buffer (Buffer) – The TVM buffer to check.

Returns:

True if the buffer is a fragment, False otherwise.

Return type:

bool

tilelang.utils.language.get_buffer_elems(buffer)

Get the number of elements in the buffer.

Parameters:

buffer (tvm.tir.Buffer)

Return type:

int

tilelang.utils.language.array_reduce(array)

Reduce an array of integers to a single integer.

Parameters:

array (List[int]) – The array of integers to reduce.

Returns:

The reduced integer.

Return type:

int

tilelang.utils.language.retrieve_func_from_module(ir_module)

Retrieve the single PrimFunc from an IRModule.

Parameters:

ir_module (IRModule) – The TVM IRModule to extract the function from. The module should contain exactly one global function.

Returns:

The single function contained in the module.

Return type:

PrimFunc

Raises:

• ValueError – If ir_module is not an IRModule.

• AssertionError – If the module contains more than one global function.

tilelang.utils.language.get_buffer_region_from_load(buffer_load)

Get the buffer region from a buffer load.

May encounter buffer load like C[0:128, 0:32], ref to pull request for buffer wise op: https://github.com/apache/tvm/pull/14693 convert load to region

Parameters:

buffer_load (tvm.tir.BufferLoad)

Return type:

Optional[tvm.tir.BufferRegion]