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.

alloc_reducer(shape, dtype[, op, replication])

Allocate a reducer 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

tilelang.language.allocate.alloc_reducer(shape, dtype, op='sum', replication=None)

Allocate a reducer buffer.

Modifications needs to conform with op, such as op=”sum” requires reducer[…] += … and op=”max” requires reducer[…] = T.max(reducer[…], …).

Only after T.fill with proper initializer the reduction may begin; only after T.finalize_reducer the partial results will be available.

For op=”sum”, filled value must be 0; for min and max, the filled initializer will become max or min clamper correspondingly. You may want to use T.max_value for min and T.min_value for max.

Parameters:

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

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

  • op (str) – The reduce operation corresponded with the reducer

  • replication (str | None) – Replication strategy, can be “all” or “none”. Defaults to not specified, and the compiler will do whatever it want.

Returns:

A TVM buffer object allocated in thread-private storage, available to reduce values in T.Parallel loops.

Return type:

T.Buffer