This document is relevant for: Trn2, Trn3

nki.language.alloc_blocks#

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

Allocate new storage organized as a grid of tile blocks.

alloc_blocks creates the same tile geometry as nki.language.alloc_tiles(), then adds a block-coordinate level above it. tile_shape defines the compute grain, block_shape defines the number of tiles in one transfer or reuse grain, and element_shape remains the exact logical tensor extent.

Block grouping does not change physical tile order or add per-block allocations. HBM uses the compact logical extent as backing storage; SBUF and PSUM use the same dense rank-two capacity layout as nki.language.alloc_tiles().

The returned NkiBlockGrid owns one independent uninitialized allocation.

Parameters:

tile_shape – The positive source-element shape of one compute tile. It must have at least two dimensions and applies to the trailing dimensions of element_shape.
block_shape – The positive tile counts grouped along each block dimension. Its rank must equal tile_shape.
element_shape – The exact positive logical element extent. Any dimensions preceding tile_shape remain untiled.
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 NkiBlockGrid over independent, uninitialized storage. Partial boundary blocks and tiles retain exact logical extents; aggregate data access requires one exact native access pattern.

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, block_shape and tile_shape ranks differ, 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_block_grid():
    staging = nl.alloc_blocks(
        tile_shape=(128, 128),
        block_shape=(2, 2),
        element_shape=(512, 512),
        buffer=nl.sbuf,
        dtype=nl.bfloat16,
    )
    dst = nl.ndarray(
        (512, 512), dtype=nl.bfloat16, buffer=nl.shared_hbm
    )
    dst_blocks = nl.blocks(
        dst,
        tile_shape=(128, 128),
        block_shape=(2, 2),
    )

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

nki.language.alloc_blocks

Contents

nki.language.alloc_blocks#