IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

stack_allocation

def stack_allocation[LayoutType: TensorLayout, //, dtype: DType, address_space: AddressSpace = AddressSpace.GENERIC, alignment: Int = align_of[dtype]()](var layout: LayoutType) -> TileTensor[dtype, LayoutType, MutUntrackedOrigin, address_space=address_space] where LayoutType.all_dims_known

Allocate a TileTensor on the stack with the given layout.

Creates a stack-allocated buffer sized for the layout and returns a TileTensor pointing to it. The layout must have all dimensions known at compile time.

Constraints:

All layout dimensions must be statically known.

Parameters:

  • ​LayoutType (TensorLayout): The layout type (inferred from layout argument).
  • ​dtype (DType): The data type of tensor elements.
  • ​address_space (AddressSpace): Memory address space (default: GENERIC).
  • ​alignment (Int): Allocation alignment in bytes (default: natural type alignment from align_of[dtype]()). Pass an explicit value when downstream loads/stores require larger alignment (for example, AMD ds_read_b128 requires 16 B, and sub-block swizzles can require alignment up to the sub-block size). Forwarded to the underlying std.memory.stack_allocation.

Args:

  • ​layout (LayoutType): The layout instance defining shape and strides.

Returns:

TileTensor[dtype, LayoutType, MutUntrackedOrigin, address_space=address_space]: A mutable TileTensor backed by stack-allocated memory.

Was this page helpful?