nki.isa.nc_n_gather#

nki.isa.nc_n_gather(dst, data, indices, name=None)[source]#

Gather elements from data according to indices using GpSimd Engine.

This instruction performs a gather operation where elements are selected from the input data tile based on flattened indices specified in the indices tile. The free dimensions of data are treated as if they were flattened into a single dimension for indexing purposes, while the partition dimension defines the parallel compute boundary.

The gather operation works independently within each partition. For each partition, the free dimensions of data are conceptually flattened, and elements are gathered according to the corresponding flattened indices from the same partition in indices. If you need to gather elements across partitions (within groups of partitions), consider using nisa.local_gather.

The n in nc_n_gather indicates that this instruction corresponds to n groups of instructions in the underlying ISA, where n = ceil(elems_per_partition / 512).

Alternatively, we could gather elements by calling nisa.dma_copy with an indirect access pattern derived from indices. However, this is less efficient than nc_n_gather, which uses GpSimd Engine to perform local data movement within SBUF, without using DMA engines.

Memory types.

All input and output tiles (data, indices, and dst) must be in SBUF. GpSimd Engine cannot access PSUM (see NeuronCore-v2 Compute Engines for details).

Data types.

The input data tile can be any valid NKI data type (see Supported Data Types for more information). The output dst tile must have the same data type as data. The indices tile must be uint32.

Layout.

The partition dimension of data, indices, and dst must be the same. Within each partition, the free dimensions of data are flattened for indexing. The free dimensions of indices determine the shape of the output dst.

Tile size.

The partition dimension size of data, indices, and dst must be the same and must not exceed 128. The number of elements per partition in dst must match the number of elements per partition in indices. The indices’ values must be within the range [0, data.size / data.shape[0]).

Parameters:

  • dst – output tile containing the gathered elements

  • data – the input tile to gather elements from

  • indices – the indices tile (uint32) specifying which elements to gather