tilelang.language.allocate

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.

Functions

alloc_shared(shape, dtype[, scope])

Allocate a shared memory buffer for inter-thread communication.

alloc_local(shape, dtype[, scope])

Allocate a local memory buffer for thread-private storage.

alloc_fragment(shape, dtype[, scope])

Allocate a fragment memory buffer for specialized operations.

alloc_var(dtype[, scope])

Allocate a single-element variable buffer.

alloc_barrier(arrive_count)

Allocate a barrier buffer.

Module Contents

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_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_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_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

tilelang.language.allocate.alloc_barrier(arrive_count)

Allocate a barrier buffer.

Parameters:

arrive_count (int) – The number of threads that need to arrive at the barrier

Returns:

A TVM buffer object allocated as a barrier

Return type:

T.Buffer