This document is relevant for: Trn2, Trn3

nki.language.while_loop#

nki.language.while_loop(init, body_fun)[source]#

Structured while loop with a register condition.

Loops while the condition register is nonzero, checking the condition before each iteration (a true while, not do-while: if init is zero the body never runs).

The loop is roughly equivalent to the following Python code:

r = init
while r != 0:
    r = body_fun(r)

The body receives the current condition value r as a VirtualRegister, passed by value. Read it inside the body (for example, materialize it with register_store and use it as a scalar_offset). The loop’s next condition is the register the body returns; carry any other loop state through SBUF/HBM.

Parameters:

  • init – Initial condition register (VirtualRegister).

  • body_fun – function (r: VirtualRegister) -> VirtualRegister called each iteration with the current condition value; returns the next condition register.

Returns:

None. Side effects in body_fun persist after the loop.

Examples:

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

# nki.language.while_loop -- flag-driven loop with a register condition
val = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=val, src=data)
acc = nl.zeros((1, 1), dtype=nl.float32, buffer=nl.sbuf)

count_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=count_sb, src=count_input)
one_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.memset(dst=one_sb, value=1)

reg = nisa.register_alloc()
nisa.register_load(reg, count_sb)

def body(r):
    nisa.tensor_tensor(dst=acc, data1=acc, data2=val, op=nl.add)
    nisa.tensor_tensor(dst=count_sb, data1=count_sb, data2=one_sb, op=nl.subtract)
    nisa.register_load(r, count_sb)
    return r

nl.while_loop(reg, body)

This document is relevant for: Trn2, Trn3