This document is relevant for: Trn2, Trn3

nki.language.ds#

nki.language.ds(start: Union[int, NkiTensor, VirtualRegister, NkiScalar], size: int)[source]#

A size-element window starting at start.

Used inside an index expression: t[nl.ds(start, size), :] is equivalent to t[start:start + size, :] for a compile-time start. size must be a compile-time integer (the access pattern needs a static element count); start may be an int or a runtime index (see NkiTensor.select()). A runtime start is only spelled this way, because a slice needs both bounds while the kernel is built.

t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t[:, nl.ds(8, 16)]        # static: columns 8..23

idx = nl.ndarray((1, 1), dtype=nl.uint32, buffer=nl.sbuf)
t[:, nl.ds(idx, 16)]      # runtime: 16 columns starting at idx

With a runtime start the window slides at execution time and the resulting view is indirect (see NkiTensor.is_indirect()).

Parameters:

  • start – window start — compile-time int or a runtime NkiTensor / VirtualRegister / NkiScalar

  • size – window length, must be a compile-time integer

Returns:

a DynamicSlice to use inside an index expression

This document is relevant for: Trn2, Trn3