This document is relevant for: Trn2, Trn3

NKI Loops#

This document covers the differences between specialization-time Python loops and NKI compiled loops.

Choosing a Loop #

NKI provides loop constructs that are unrolled eagerly at specialization-time, and loop constructs that run on device. For specialization-time loops, NKI provides Python for and while loops with specialization-time iteration bounds. For compiled loops, NKI provides nki.language.fori_loop and nki.language.while_loop with specialization-time or dynamic iteration bounds. The table below summarizes the key differences between loop constructs:

Loop Construct	Behavior	When to Use
Python `for` loops	Unrolled at specialization time	Small iteration count known at specialization time
Python `while` loops	Unrolled at specialization time	Small iteration count, with a condition known at specialization time
`nki.language.fori_loop()`	Compiled loop	Large iteration count, or iteration count unknown at specialization time
`nki.language.while_loop()`	Compiled loop	Stopping condition known only within the loop, or unknown at specialization time

Earlier versions of NKI used range iterators instead of loop constructs. NKI provided three specialization-time range iterators and one dynamic range iterator.

In NKI 0.7 and later, the static range iterators nki.language.affine_range(), nki.language.sequential_range(), and nki.language.static_range() are deprecated and default to Python range. The dynamic range iterator nki.language.dynamic_range() and Python while loops with a VirtualRegister condition are unsupported. Use NKI loop constructs instead.

Note

To use the dynamic range iterator or Python while loops with a VirtualRegister condition, you must opt in to the NKI meta-programming language from releases before 0.7 by setting the environment variable NKI_COMPILER_FRONTEND=custom. This option will be removed in a future version of NKI. For more information, see nki-migrate-0-6-to-0-7.

The table below shows how each loop construct and range iterator from before NKI 0.7 maps to its equivalent in NKI 0.7 and later:

Before NKI 0.7	NKI 0.7 and later	Behavior	When to Use
Python `for` loops	Python `for` loops (unchanged)	Unrolled at specialization time	Small iteration count known at specialization time
`nki.language.affine_range()` (alias for Python `range`)	Python `for` loops	Unrolled at specialization time	Small iteration count known at specialization time
`nki.language.sequential_range()` (alias for Python `range`)	Python `for` loops	Unrolled at specialization time	Small iteration count known at specialization time
`nki.language.static_range()` (alias for Python `range`)	Python `for` loops	Unrolled at specialization time	Small iteration count known at specialization time
`nki.language.dynamic_range()`	`nki.language.fori_loop()`	Hardware loop, one copy of the body	Iteration count unknown at specialization time
Python `while` loops	Python `while` loops (unchanged)	Unrolled at specialization time	Small iteration count, with a condition known at specialization time
Python `while` loops with a `VirtualRegister` condition	`nki.language.while_loop()`	Hardware loop, one copy of the body	Iteration count unknown at specialization time

Unrolled Loops #

During specialization, Python for and while loops are fully unrolled, turning the loop bodies into straight-line code. Loop unrolling is a transformation technique designed to improve performance at the cost of code size. Unrolling loops helps the compiler eliminate redundant instructions, reducing overhead.

Consider an example using Python loops that print the values from 0 to 4:

# First loop, traditional Python for loop.
for i in range(3):
    print(i)

# Second loop, traditional Python while loop.
j = 3
while j < 5:
    print(j)
    j += 1

After specialization, the loops unroll to the following Python code:

print(0)
print(1)
print(2)
print(3)
print(4)

Implications of Loop Unrolling #

Although loop unrolling can reduce overhead by converting loop bodies into straight-line code, loops over a large number of iterations can greatly increase compile time — each iteration inlines a copy of the loop body, which NKI must trace. For loops with many iterations, consider using compiled loops nki.language.fori_loop or nki.language.while_loop instead.

# Loop over a very large number of iterations.
for i in range(1_000_000):
   nisa.tensor_scalar(dst=acc, data=acc, op0=nl.add, operand0=1.0)

Loop unrolling can also generate many instructions, which leads to larger NEFFs that take longer to load. For example, after specialization, the loop above unrolls into 1 million instructions in the NEFF:

nisa.tensor_scalar(dst=acc, data=acc, op0=nl.add, operand0=1.0)
nisa.tensor_scalar(dst=acc, data=acc, op0=nl.add, operand0=1.0)
# ... many nisa.tensor_scalar instructions
nisa.tensor_scalar(dst=acc, data=acc, op0=nl.add, operand0=1.0)
nisa.tensor_scalar(dst=acc, data=acc, op0=nl.add, operand0=1.0)

As a rule of thumb, use Python loops for small iteration counts, where straight-line code helps performance, and compiled loops for large or unknown iteration counts, where unrolling increases compile time and NEFF size.

Compiled Loops #

The body of a compiled loop is traced once, during specialization, into an NKI IR loop. The compiler lowers the NKI IR loop to NeuronISA. NKI offers two types of compiled loops: nki.language.fori_loop and nki.language.while_loop.

Using nki.language.fori_loop #

nki.language.fori_loop(lower, upper, body_fun, step=1) runs a counted loop, similar to for i in range(lower, upper, step). Instead of a loop body, you pass a function, body_fun, that takes the loop variable i and returns nothing. NKI traces body_fun once. The following example uses nki.language.fori_loop:

import numpy as np

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


@nki.jit
def sum_rows_fori(data, count):
    ROWS, COLS = data.shape

    # Read the run-time row count into a scalar to use as the loop bound.
    count_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
    nisa.dma_copy(dst=count_sb, src=count)
    n = nisa.scalar_load(count_sb)

    total = nl.zeros((1, COLS), dtype=nl.float32, buffer=nl.sbuf)
    row = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.sbuf)

    def body(i):
        # The loop variable selects row i of data at run time.
        nisa.dma_copy(
            dst=row,
            src=data.ap([[COLS, 1], [1, COLS]], scalar_offset=i, indirect_dim=0),
        )
        nisa.tensor_tensor(dst=total, data1=total, data2=row, op=nl.add)

    nl.fori_loop(0, n, body)

    out = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=out, src=total)
    return out


data = np.arange(32, dtype=np.float32).reshape(4, 8)
count = np.array([[3]], dtype=np.int32)
sum_rows_fori(data, count)
# [[24. 27. 30. 33. 36. 39. 42. 45.]]

Parameter Constraints#

lower / upper: Can be a Python int, NkiScalar, or VirtualRegister object. When either lower or upper is a NkiScalar or both lower and upper are Python int, the loop variable is a NkiScalar. Otherwise, the loop variable is a VirtualRegister.
body_fun: A function that takes the loop variable and returns None. Writing to the loop variable does not change the loop. To keep state across iterations, use SBUF or HBM tensors.
step: Must be a positive Python int. Passing a NkiScalar or VirtualRegister raises an AssertionError.

Using nki.language.while_loop #

nki.language.while_loop(init, body_fun) runs while a condition register is not zero, similar to while r != 0. NKI checks the condition before each iteration. If init is zero, the body does not run. The body function takes the current condition register and must return the register to check on the next iteration.

Parameter Constraints#

init: Must be a VirtualRegister. Passing a NkiScalar raises an AssertionError.
body_fun: A function that takes the condition register and returns the register to check on the next iteration. Returning None raises a TypeError.

Interaction with `no_reorder`#

Compiled loops (nki.language.fori_loop, nki.language.while_loop) and dynamic range iterators (nki.language.dynamic_range) inside nki.language.no_reorder() blocks are currently not supported. The following code patterns will produce an error:

# Not supported: nl.fori_loop inside nl.no_reorder
with nl.no_reorder():
    nl.fori_loop(0, n, body)

# Not supported: nl.while_loop inside nl.no_reorder
with nl.no_reorder():
    nl.while_loop(cond, body)

This document is relevant for: Trn2, Trn3

NKI Loops

Contents

NKI Loops#

Choosing a Loop#

Unrolled Loops#

Implications of Loop Unrolling#

Compiled Loops#

Using nki.language.fori_loop#

Parameter Constraints#

Using nki.language.while_loop#

Parameter Constraints#

Interaction with no_reorder#

Choosing a Loop #

Unrolled Loops #

Implications of Loop Unrolling #

Compiled Loops #

Using nki.language.fori_loop #

Using nki.language.while_loop #

Interaction with `no_reorder`#