This document is relevant for: Trn2, Trn3

nki.language.fori_loop#

nki.language.fori_loop(lower, upper, body_fun, step=1)[source]#

Structured for loop with dynamic bounds.

Executes body_fun(i) for each iteration value from lower to upper (exclusive) with the given step. The body is a callable that receives the current iteration value.

The loop is roughly equivalent to the following Python code:

for i in range(lower, upper, step):
    body_fun(i)

The body receives the current iteration value i as a VirtualRegister or NkiScalar, passed by value. Read it inside the body (for example, as a scalar_offset into a tensor). The loop controls iteration via lower/upper/step, so writing to i has no effect on the loop; carry any other loop state through SBUF/HBM.

Parameters:

  • lower – start value (int, VirtualRegister, or NkiScalar).

  • upper – end value (exclusive) (int, VirtualRegister, or NkiScalar).

  • body_fun – function (i) -> None called each iteration with the current iteration value. i is a NkiScalar unless either lower or upper is a VirtualRegister, then it will be a VirtualRegister

  • step – step size, must be a compile-time positive integer.

Returns:

None. Side effects in body_fun persist after the loop.

A bound must be 32-bit: the loop compiles into a compare-and-branch that has no 64-bit form, so an nl.int64 scalar cannot bound it. A 64-bit scalar can still be computed inside the body and used as an offset there.

Examples:

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

# nki.language.fori_loop -- counted loop with a runtime (dynamic) upper bound
ub_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=ub_sb, src=ub_input)
ub_reg = nisa.register_alloc()
nisa.register_load(ub_reg, ub_sb)

zeros = nl.zeros((1, N), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=output, src=zeros)
temp = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf)

def body(i):
    # `i` is the loop induction register; read it as a dynamic offset
    nisa.dma_copy(
        dst=temp,
        src=data.ap([[1, 1], [1, 1]], scalar_offset=i, indirect_dim=1),
    )
    nisa.dma_copy(
        dst=output.ap([[1, 1], [1, 1]], scalar_offset=i, indirect_dim=1),
        src=temp,
    )

nl.fori_loop(0, ub_reg, body)

This document is relevant for: Trn2, Trn3