This document is relevant for: Trn2, Trn3

nki.language.NkiBlockGrid

class nki.language.NkiBlockGrid(**kwargs: Any)

A tensor view navigated in block coordinates.

A block groups a fixed number of tiles along every tiled dimension. NkiGrid.shape counts blocks, block_shape gives tiles per block, and NkiGrid.tile_shape gives source elements per tile. This separates a coarser transfer or reuse grain from the tile grain consumed by compute instructions.

Partial final blocks and tiles are retained and clamped to the exact source extent. Creating and indexing a block view allocates no storage and moves no data.

Examples:

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

@nki.jit
def copy_block_grid(src):
    tile_shape = (128, 256)
    block_shape = (2, 2)
    src_blocks = nl.blocks(
        src, tile_shape=tile_shape, block_shape=block_shape
    )
    dst = nl.ndarray(
        src.shape, dtype=src.dtype, buffer=nl.shared_hbm
    )
    dst_blocks = nl.blocks(
        dst, tile_shape=tile_shape, block_shape=block_shape
    )

    for i in range(src_blocks.shape[0]):
        for j in range(src_blocks.shape[1]):
            src_block = src_blocks[i, j]
            block = nl.blocks(
                nl.ndarray(
                    src_block.element_shape,
                    dtype=src.dtype,
                    buffer=nl.sbuf,
                ),
                tile_shape=src_block.tile_shape,
                block_shape=src_block.block_shape,
            )
            nisa.dma_copy(dst=block.data, src=src_block.data)
            nisa.dma_copy(
                dst=dst_blocks[i, j].data,
                src=block.data,
            )
    return dst

See also

nki.language.blocks(): Create or regroup a block grid.

Methods

Attributes

block_shape	Return the nominal number of tiles grouped along each block axis.
buffer	Return the memory region containing the backing tensor.
data	Expose the complete grid selection as one native tensor operand.
dtype	Return the element dtype of the backing tensor.
element_shape	Return the number of selected elements along each source dimension.
index_step_elements	Return each grid-coordinate step in source-element units.
is_data_representable	Return whether data can represent this selection exactly.
is_remainder	Return whether this selection contains a partial boundary tile.
shape	Return the extents of the grid coordinates available for indexing.
tile_shape	Return the nominal source-element shape of one tile.

property block_shape: tuple[int, ...]

Return the nominal number of tiles grouped along each block axis.

A boundary block may contain fewer tiles. Use nki.language.tiles() on a selected block to inspect its exact interior tile counts.

Returns:

The configured tile count per block dimension.

This document is relevant for: Trn2, Trn3