This document is relevant for: Trn1, Trn2, Trn3

nki.language.dynamic_range#

nki.language.dynamic_range(start, stop=None, step=1)[source]#

Create a sequence for dynamic loop iteration.

Create a sequence of numbers for use as dynamic loop iterators in NKI. The loop runs on device with dynamic bounds.

Parameters:

  • start – start value (or stop if stop is None), can be VirtualRegister.

  • stop – stop value (exclusive), can be VirtualRegister.

  • step – step size, must be a compile-time positive integer (not VirtualRegister).

Returns:

an iterator yielding integer values from start to stop.

Examples:

import nki.language as nl

# nki.language.dynamic_range
for _ in nl.dynamic_range(1):
    tile = nl.load(input_tensor[0:128, 0:512])
    result = nl.multiply(tile, tile)
    nl.store(out_tensor[0:128, 0:512], result)

This document is relevant for: Trn1, Trn2, Trn3