tilelang.language.allocate module#

Memory allocation utilities for Tile-AI programs.

This module provides a set of functions for allocating different types of memory buffers in Tile-AI programs. It wraps TVM’s buffer allocation functionality with convenient interfaces for different memory scopes.

Available allocation functions:

  • alloc_shared: Allocates shared memory buffers for inter-thread communication

  • alloc_local: Allocates local memory buffers for thread-private storage

  • alloc_fragment: Allocates fragment memory buffers for specialized operations

  • alloc_var: Allocates single-element variable buffers

Each function takes shape and dtype parameters and returns a TVM buffer object with the appropriate memory scope.

tilelang.language.allocate.alloc_fragment(shape, dtype, scope='local.fragment')#

Allocate a fragment memory buffer for specialized operations.

Parameters:

  • shape (tuple) – The shape of the buffer to allocate

  • dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)

  • scope (str, optional) – The memory scope. Defaults to “local.fragment”

Returns:

A TVM buffer object allocated in fragment memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_local(shape, dtype, scope='local')#

Allocate a local memory buffer for thread-private storage.

Parameters:

  • shape (tuple) – The shape of the buffer to allocate

  • dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)

  • scope (str, optional) – The memory scope. Defaults to “local”

Returns:

A TVM buffer object allocated in local memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_shared(shape, dtype, scope='shared.dyn')#

Allocate a shared memory buffer for inter-thread communication.

Parameters:

  • shape (tuple) – The shape of the buffer to allocate

  • dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)

  • scope (str, optional) – The memory scope. Defaults to “shared.dyn”

Returns:

A TVM buffer object allocated in shared memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_var(dtype, scope='local.var')#

Allocate a single-element variable buffer.

Parameters:

  • dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)

  • scope (str, optional) – The memory scope. Defaults to “local.var”

Returns:

A TVM buffer object allocated as a single-element variable

Return type:

T.Buffer