This document is relevant for: Trn2, Trn3

nki.language.blocks#

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

Create a block view over a tensor where each block contains multiple tiles.

blocks() adds a second level of grouping above the tile grid. A block is a rectangular group of tiles intended as a coarser transfer or reuse grain, while tile_shape remains the compute grain presented to nisa.* operations. When one native access pattern can represent the complete block, block.data allows it to participate in one aggregate operation.

block_shape is expressed in tile counts. For a 4 x 4 tile grid, block_shape=(2, 2) creates a 2 x 2 block grid:

block grid: 2 x 2
+-------------------+-------------------+
| tile (0,0) (0,1)  | tile (0,2) (0,3)  |
| tile (1,0) (1,1)  | tile (1,2) (1,3)  |
+-------------------+-------------------+
| tile (2,0) (2,1)  | tile (2,2) (2,3)  |
| tile (3,0) (3,1)  | tile (3,2) (3,3)  |
+-------------------+-------------------+

Blocks are logical views over the source. Creating, regrouping, slicing, or indexing a block grid allocates no storage, copies no data, and emits no instruction.

Note

Indexing preserves NkiBlockGrid. Call nki.language.tiles() on a selected block to expose its interior tile coordinates, then use NkiGrid.data when an instruction requires one native tensor operand.

A block containing nonuniform boundary tiles may need to be split into uniform tile regions first.

Parameters:

source – A direct NkiTensor, NkiTileGrid, or NkiBlockGrid whose storage and selected addresses are shared by the returned view.
tile_shape – The nominal source-element shape of one compute tile. Required for an NkiTensor; optional for an existing grid.
block_shape – The positive tile counts grouped along each block dimension. It must be a tuple with the same rank as the effective tile_shape.

Returns:

An NkiBlockGrid over the same storage. Its shape counts blocks, while tile_shape and block_shape preserve both levels of the hierarchy. Partial boundary blocks and tiles are retained at their exact logical extents; aggregate data access still requires one exact native access pattern.

Raises:

TypeError – if source is unsupported, tile_shape is omitted for an NkiTensor, or a shape is not a tuple of exact Python integers.
ValueError – if a shape entry is nonpositive, block_shape and tile_shape have different ranks, the source cannot be tiled, or retiling or grouping cannot preserve the selected addresses.

Examples:

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

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

    for i in range(src_blocks.shape[0]):
        for j in range(src_blocks.shape[1]):
            src_block = src_blocks[i, j]
            dst_block = dst_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,
            )

            # copy entire block in single instruction
            nisa.dma_copy(dst=block.data, src=src_block.data)

            # demote to tile-grid to access individual tiles
            block_tiles = nl.tiles(block)
            for k in range(block_tiles.shape[0]):
                for p in range(block_tiles.shape[1]):
                    tile = block_tiles[k, p].data
                    nisa.tensor_scalar(
                        dst=tile,
                        data=tile,
                        op0=nl.multiply,
                        operand0=2.0,
                    )

            nisa.dma_copy(
                dst=dst_block.data,
                src=block.data,
            )

    return dst

Source kinds.

The source can be NkiTensor or NkiTileGrid or NkiBlockGrid.

source              tile_shape          behavior
------------------  -----------------  ------------------------------
NkiTensor           required           tile, then group
NkiTileGrid         omitted            group existing tiles
NkiTileGrid         supplied           retile, then group
NkiBlockGrid        omitted            remove old groups, regroup
NkiBlockGrid        supplied           remove groups, retile, regroup

Remainder blocks.

Tile counts are ceiling-divided from the source extent, and block counts are ceiling-divided from the tile counts. A final block can therefore contain fewer tiles, and its final tile can contain fewer source elements. Both levels are clamped to the exact source extent.

nki.language.blocks

Contents

nki.language.blocks#