This document is relevant for: Trn2, Trn3

nki.language.tiles#

nki.language.tiles(source: nki.language.NkiTensor | nki.language.NkiTileGrid | nki.language.NkiBlockGrid, *, tile_shape: Optional[tuple[int, ...]] = None) → NkiTileGrid[source]#

Decompose a tensor into a logical grid of fixed-size tiles.

The returned grid is a view over source, indexed by tile coordinate rather than by manually calculated element offsets and strides. A kernel can select src_tiles[i, j] to address one tile, src_tiles[i, :] to address a row of tiles, or src_tiles[:, j] to address a column.

For example, tile_shape=(128, 256) over a (512, 1024) tensor produces a 4 x 4 grid:

                              F (free dimension)
                   256       256       256       256
              +---------+---------+---------+---------+
P         128 |  (0,0)  |  (0,1)  |  (0,2)  |  (0,3)  |
(partition)    +---------+---------+---------+---------+
          128 |  (1,0)  |  (1,1)  |  (1,2)  |  (1,3)  |
              +---------+---------+---------+---------+
          128 |  (2,0)  |  (2,1)  |  (2,2)  |  (2,3)  |
              +---------+---------+---------+---------+
          128 |  (3,0)  |  (3,1)  |  (3,2)  |  (3,3)  |
              +---------+---------+---------+---------+

Note

Grid indexing preserves the grid abstraction. src_tiles[i, j] returns a one-tile NkiTileGrid, not an NkiTensor. Use src_tiles[i, j].data when an NKI instruction requires the underlying tensor operand. Accessing data is also a view operation.

Parameters:

source – The tensor region to tile. It may be a direct NkiTensor, an existing NkiTileGrid, or an NkiBlockGrid. The returned grid shares its storage, dtype, buffer, offsets, and physical strides.
tile_shape – The nominal shape of one tile in source-element units, with at least two dimensions. Every entry must be a positive, exact Python integer. It is required when source is an NkiTensor and optional when source is already a grid. For an HBM tensor, fewer dimensions than source leaves the leading dimensions as batch coordinates. A tile extent may be larger than its source extent, producing one partial tile.

Returns:

An NkiTileGrid over source. Its shape counts addressable tiles, its tile_shape records the nominal tile_shape, and its element_shape records the number of selected elements in each source dimension. Partial boundary tiles expose only valid elements, although independently allocated on-chip storage may reserve nominal tile capacity behind them.

Raises:

TypeError – if source is unsupported, tile_shape is omitted for an NkiTensor, a shape is not a tuple of exact Python integers, or the source buffer is not an NKI MemoryRegion.
ValueError – if tile_shape has fewer than two dimensions, contains a nonpositive extent, has a rank incompatible with the source, the source is an indirect NkiTensor, the source buffer is unsupported, or an existing grid cannot be retiled while preserving its selected addresses.

Examples:

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

@nki.jit
def scale_by_tile(src):
    dst = nl.ndarray(
        src.shape, dtype=src.dtype, buffer=nl.shared_hbm
    )

    src_tiles = nl.tiles(src, tile_shape=(128, 256))
    dst_tiles = nl.tiles(dst, tile_shape=(128, 256))

    for i in range(src_tiles.shape[0]):
        for j in range(src_tiles.shape[1]):
            src_tile = src_tiles[i, j]
            dst_tile = dst_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.tensor_scalar(
                dst=tile.data, data=tile.data, op0=nl.multiply, operand0=2.0
            )
            nisa.dma_copy(dst=dst_tile.data, src=tile.data)

    return dst

Source kinds.

The source can be NkiTensor or NkiTileGrid or NkiBlockGrid.

source              tile_shape          result
------------------  -----------------  ------------------------------
HBM NkiTensor       required           tile grid over HBM
SBUF NkiTensor      required           tile grid over SBUF
PSUM NkiTensor      required           tile grid over PSUM
NkiTileGrid         omitted            same tile-grid view
NkiTileGrid         supplied           address-preserving retile
NkiBlockGrid        omitted            remove block grouping
NkiBlockGrid        supplied           remove blocks, then retile

Both NkiBlockGrid forms remove the block-coordinate level. Omitting tile_shape preserves the existing tile structure. Supplying it rebuilds the tile grid at the requested compute grain. See Retiling section below.

Remainder tiles.

Grid dimensions use ceiling division. When a source extent is not an exact multiple of its tile extent, the final tile is clamped to the remaining elements. For source.shape == (300, 600) and tile_shape == (128, 256):

                      column 0       column 1      column 2
                      256 elements   256 elements   88 elements
                    +--------------+--------------+-----------+
row 0: 128 elements |  128 x 256   |  128 x 256   | 128 x 88 |
                    +--------------+--------------+-----------+
row 1: 128 elements |  128 x 256   |  128 x 256   | 128 x 88 |
                    +--------------+--------------+-----------+
row 2:  44 elements |   44 x 256   |   44 x 256   |  44 x 88 |
                    +--------------+--------------+-----------+

The nominal tile_shape remains (128, 256), while src_tiles[-1, -1].data.shape is (44, 88). The selected boundary tile reports is_remainder == True. Its logical extent and data operand contain no padding.

# check if last row/col tiles are partials
src_tiles[-1].is_remainder
src_tiles[:, -1].is_remainder
src_tiles[-1, -1].is_remainder

nki.language.tiles

Contents

nki.language.tiles#