Function sblock_alloc_buffer

python/tvm/tirx/script/builder/ir.py:906–990

SBlock-level buffer allocation function. Parameters ---------- shape : Union[List[PrimExpr], Tuple[PrimExpr], PrimExpr, Integral] The type of the buffer prior to flattening. dtype : str The data type in the content of the buffer. data : Var The pointer to

(
    shape: list[PrimExpr] | tuple[PrimExpr] | PrimExpr | Integral,
    dtype: str = "float32",
    data: Var = None,
    strides: list[PrimExpr] | None = None,
    elem_offset: PrimExpr = None,
    scope: str = "global",
    align: int = -1,
    offset_factor: int = 0,
    buffer_type: str = "default",
    axis_separators: list[int] | None = None,
    layout: str | Layout | None = "default",
    allocated_addr: int | tuple[int, ...] | None = None,
)

def sblock_alloc_buffer(
    shape: list[PrimExpr] | tuple[PrimExpr] | PrimExpr | Integral,
    dtype: str = "float32",
    data: Var = None,
    strides: list[PrimExpr] | None = None,
    elem_offset: PrimExpr = None,
    scope: str = "global",
    align: int = -1,
    offset_factor: int = 0,
    buffer_type: str = "default",
    axis_separators: list[int] | None = None,
    layout: str | Layout | None = "default",
    allocated_addr: int | tuple[int, ...] | None = None,
) -> Buffer:
    """SBlock-level buffer allocation function.

    Parameters
    ----------
    shape : Union[List[PrimExpr], Tuple[PrimExpr], PrimExpr, Integral]
        The type of the buffer prior to flattening.
    dtype : str
        The data type in the content of the buffer.
    data : Var
        The pointer to the head of the data.
    strides : List[PrimExpr]
        The strides of each dimension.
    elem_offset : PrimExpr
        The offset in terms of number of dtype elements (including lanes).
    scope : str
        The optional storage scope of buffer data pointer.
    align : int
        The alignment requirement of data pointer in bytes.
    offset_factor : int
        The factor of elem_offset field.
    buffer_type : str
        The buffer type.
    axis_separators : List[int]
        The separators between input axes when generating flattened output axes.

    layout: Optional[Union[str, Layout]]
        The layout of the buffer.

    allocated_addr: Optional[Union[int, Tuple[int]]]
        The address of the allocated buffer. Might be multi-dimensional.
        There can be pooled storage scopes on some devices. For example,
        the Trainium device has a pooled storage scope for the SRAN buffers. ("trn.sbuf")
        CUDA has a pooled storage scope for the shared memory ("shared.dyn")

    Returns
    -------
    res : Buffer
        The allocated buffer.
    """
    shape = (shape,) if isinstance(shape, PrimExpr | Integral) else shape
    if strides is not None:
        strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides]
    else:
        strides = []