This document is relevant for: Trn2, Trn3

nki.language.NkiTileGrid

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

A tensor view navigated in tile coordinates.

Each tiled entry in NkiGrid.shape counts tiles, while NkiGrid.tile_shape gives the nominal source-element shape of one tile. Untiled HBM dimensions form a leading prefix and are indexed one slab at a time.

Integer indexing consumes coordinates but retains NkiTileGrid. NkiGrid.data exposes the exact NkiTensor operand for the selected view.

Instances come from nki.language.tiles() or a block-to-tile transition. Creating and indexing the 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_tile_grid(src):
    tile_shape = (128, 256)
    src_tiles = nl.tiles(src, tile_shape=tile_shape)
    dst = nl.ndarray(
        src.shape, dtype=src.dtype, buffer=nl.shared_hbm
    )
    dst_tiles = nl.tiles(dst, tile_shape=tile_shape)

    for i in range(src_tiles.shape[0]):
        for j in range(src_tiles.shape[1]):
            src_tile = src_tiles[i, j]
            tile = nl.ndarray(
                src_tile.element_shape,
                dtype=src.dtype,
                buffer=nl.sbuf,
            )
            nisa.dma_copy(dst=tile.data, src=src_tile.data)
            nisa.dma_copy(
                dst=dst_tiles[i, j].data,
                src=tile.data,
            )
    return dst

See also

NkiBlockGrid: A grid whose cells group multiple tiles. nki.language.tiles(): Create or retile a tile grid.

Methods

Attributes

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.

This document is relevant for: Trn2, Trn3