This document is relevant for: Trn2, Trn3

nki.language.alloc_tiles#

nki.language.alloc_tiles(*, tile_shape: tuple[int, ...], element_shape: tuple[int, ...], buffer: MemoryRegion, dtype: str, address: Optional[tuple] = None, name: str = '') → NkiTileGrid[source]#

Allocate independent, uninitialized storage organized as a grid of tiles.

The returned NkiTileGrid owns one backing allocation and provides both tile navigation and an aggregate NkiGrid.data operand.

Parameters:

  • tile_shape – The positive source-element shape of one tile. It must be a tuple containing at least two exact Python integers and applies to the trailing dimensions of element_shape.

  • element_shape – The exact positive logical element extent.

  • buffer – Destination MemoryRegion, such as nl.shared_hbm, nl.sbuf, or nl.psum.

  • dtype – Destination NKI element dtype.

  • address – Optional manual address (partition_offset, free_offset) with the same contract as nki.language.ndarray().

  • name – the name of the tensor, used in scheduling.

Returns:

A new NkiTileGrid over independent, uninitialized storage. Its shape counts tiles and its boundary selections expose their exact clamped extents.

Raises:

  • TypeError – if a shape is not a tuple of exact Python integers or buffer is not an NKI MemoryRegion.

  • ValueError – if a shape entry is nonpositive, tile_shape has fewer than two dimensions or greater rank than element_shape, or buffer is unsupported.

  • AssertionError – if the native allocation violates a nki.language.ndarray() constraint, including PSUM address alignment.

Examples:

import nki
import nki.isa as nisa
import nki.language as nl

@nki.jit
def initialize_tile_grid():
    staging = nl.alloc_tiles(
        tile_shape=(128, 256),
        element_shape=(256, 512),
        buffer=nl.sbuf,
        dtype=nl.bfloat16,
    )
    dst = nl.ndarray(
        (256, 512), dtype=nl.bfloat16, buffer=nl.shared_hbm
    )
    dst_tiles = nl.tiles(dst, tile_shape=(128, 256))

    nisa.memset(dst=staging.data, value=0.0)
    nisa.dma_copy(dst=dst_tiles.data, src=staging.data)
    return dst

See also

nki.language.tiles(): View existing storage as a tile grid. nki.language.alloc_blocks(): Allocate tiles grouped into blocks.

This document is relevant for: Trn2, Trn3