tilus.ir.TMemoryTensor

tilus.ir.TMemoryTensor

class tilus.ir.tensor.TMemoryTensor[source]

A tensor that resides in tensor memory (TMEM).

Tensor memory is a dedicated on-chip memory available on Blackwell (SM 10.0+) GPUs, private to the SM’s tensor cores. It is organized as a 2D structure of lanes (rows) and columns, with each cell being 32 bits. The number of lanes (shape[-2]) must be 32, 64, or 128.

dtype

The data type of the tensor elements.

Type:

DataType

shape

The shape of the tensor. Must have at least 2 dimensions. The second-to-last dimension (lanes) must be 32, 64, or 128.

Type:

tuple[int, …]

optional_layout

The layout of the tensor, which is optional. When not provided, the layout will be automatically inferred with compiler pass.

Type:

TMemoryLayout, optional

property layout: TMemoryLayout

Get the layout of the TMemoryTensor.

Returns:

ret – The layout of the TMemoryTensor.

Return type:

TMemoryLayout

Raises:

ValueError – If the TMemoryTensor does not have a layout defined.

Contents