This document is relevant for: Trn2, Trn3

nki.language.NkiGrid#

class nki.language.NkiGrid(**kwargs: Any)[source]#

Base class for tensor views organized into tiles or blocks.

NkiGrid is the common interface for NkiTileGrid and NkiBlockGrid. A grid describes an existing tensor region through logical tile or block coordinates while retaining its dtype, buffer, and underlying storage.

Users do not construct NkiGrid directly. Grid views are created via nki.language.tiles() and nki.language.blocks():

tile_grid = nl.tiles(src, tile_shape=(128, 256))
block_grid = nl.blocks(
    src,
    tile_shape=(128, 256),
    block_shape=(2, 2),
)

nki.language.tiles() returns an NkiTileGrid, and nki.language.blocks() returns an NkiBlockGrid. Both concrete objects are instances of NkiGrid. Refer to those factory functions for the tiling, block-grouping, and boundary behavior of the returned view.

See also

nki.language.tiles(): Create a tile-coordinate view. nki.language.blocks(): Group tiles into block coordinates.

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.

property buffer: MemoryRegion#

Return the memory region containing the backing tensor.

Returns:: The MemoryRegion backing the view.

property data: NkiTensor#

Expose the complete grid selection as one native tensor operand.

Grid navigation retains tile or block coordinates. data is the explicit transition to the NkiTensor accepted by NKI instructions. The returned tensor preserves the selection’s offsets, strides, and physical gaps.

Returns:: An exact native tensor operand over the current grid region.
Raises:: ValueError – if one native access pattern cannot represent the selected region, for example across an unsupported nonuniform partition remainder or an overly complex access pattern. Split the grid into representable regions in that case. Use is_data_representable to test this before accessing data.

Note

data.shape describes the native access-pattern traversal. It may differ from element_shape and from the private backing allocation shape. For example, a complete two-dimensional tile grid can expose (P, tile_rows, tile_columns, F), while a contiguous free-dimension tail can compact to (P, valid_free_elements). Capacity reserved behind an SBUF or PSUM remainder is never exposed.

Examples:

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

@nki.jit
def copy_grid_as_one_operand(src):
    # src is tile-aligned 2D, for example (256, 512).
    src_tiles = nl.tiles(src, tile_shape=(128, 128))
    staging_tiles = nl.tiles(
        nl.ndarray(
            src_tiles.element_shape,
            dtype=src.dtype,
            buffer=nl.sbuf,
        ),
        tile_shape=src_tiles.tile_shape,
    )
    dst = nl.ndarray(
        src.shape, dtype=src.dtype, buffer=nl.shared_hbm
    )
    dst_tiles = nl.tiles(dst, tile_shape=(128, 128))

    nisa.dma_copy(dst=staging_tiles.data, src=src_tiles.data)
    nisa.dma_copy(dst=dst_tiles.data, src=staging_tiles.data)
    return dst

property dtype: str#

Return the element dtype of the backing tensor.

Returns:: The grid’s NKI element dtype.

property element_shape: tuple[int, ...]#

Return the number of selected elements along each source dimension.

Returns:: The selected element counts in source-dimension order.

property index_step_elements: tuple[int, ...]#

Return each grid-coordinate step in source-element units.

Returns:: The source-element stride of each current grid coordinate.

property is_data_representable: bool#

Return whether data can represent this selection exactly.

Returns:: True when one native tensor operand can describe the complete selection.

property is_remainder: bool#

Return whether this selection contains a partial boundary tile.

Returns:: True when at least one selected tile extent is partial.

property shape: tuple[int, ...]#

Return the extents of the grid coordinates available for indexing.

Tile-grid dimensions count tiles and block-grid dimensions count blocks. Untiled leading HBM dimensions appear first. Counts include partial boundary cells. Integer indexing consumes a coordinate level; slicing preserves it with the selected count.

Returns:: One positive count per currently indexable grid dimension.

property tile_shape: tuple[int, ...]#

Return the nominal source-element shape of one tile.

Returns:: The configured per-tile shape in source-element units.

This document is relevant for: Trn2, Trn3

nki.language.NkiGrid

Contents

nki.language.NkiGrid#