Repository navigation
NKI docs: iterations of a dynamic loop (fori_loop / while_loop / dynamic_range) never overlap; each ends with an all-engine barrier #1411
Description
Describe the issue
The NKI dynamic-loop documentation (the deep dive, the nl.fori_loop API page and the 0.6.0 migration guide) explains bounds, registers and allocation. It does not say that the compiler synchronizes every engine at the end of every iteration. NKI's own source states it. nki/compiler/_baremetal_driver.py in NKI 0.6.0: "lower_control: inserts Drain+Barrier+SemaReset+Barrier before loop body terminators to synchronize all engines and reset semaphore state between iterations. See NKI-1699." Nothing of iteration i + 1 (in particular its DMA loads) starts before iteration i has finished on every engine. A dynamic loop therefore cannot pipeline across iterations the way an unrolled loop does, and each iteration pays its own pipeline fill and drain. That is a large and non-obvious performance cost for exactly the kernels that need dynamic loops (data-dependent trip counts, compact NEFFs).
Measured on trn1.2xlarge, Neuron SDK 2.32.0, nki 0.6.0+31049202112.g85070674, neuronx-cc 2.27.5334.0+f702b353, 2026-10-05. The kernel is loop_repro.py below: each iteration loads a [128, 2048] fp32 tile (1 MiB) from HBM, applies nisa.activation(exp) and a nisa.tensor_scalar, and stores it. Device time is neuron-bench exec -n 1000 -w 50 --fixed-nc-count 1 --disable-hw-profile nc_latency p50 on the NEFF from NKI_ARTIFACTS_DIR. Every variant's output was checked against NumPy (rel. error 7.8e-6).
| loop (64 iterations) | device time p50 | per iteration |
|---|---|---|
for i in range(64) (unrolled) |
483 us | 7.5 us |
nl.fori_loop(0, 64, body) |
1147 us | 17.9 us |
for i in nl.dynamic_range(64) |
1147 us | 17.9 us |
nl.fori_loop(0, 16, body4), 4 tiles per iteration (unrolled inside the body) |
653 us | 10.2 us per tile |
From neuron-explorer profiles of the 16-iteration NEFFs (exp-to-exp start interval): range 4.25 us median (the next iteration's load overlaps the current compute), fori_loop with a constant bound 17.90 us, fori_loop with a register bound loaded from an input tensor 17.92 us, nl.while_loop 18.27 us. The fori_loop profile holds 439 DRAIN and 90 EVENT_SEMAPHORE_RANGE_CLEAR instructions, against 40 and 5 for the unrolled loop: about 25 drains and 5 semaphore range clears per iteration.
The 4-tiles-per-iteration variant shows the workaround: the loads of one tile overlap the compute of the previous tile inside the body, and the barrier is paid once per group of tiles. That kind of guidance belongs in the deep dive.
Suggested documentation (PR linked): a section "Iterations Do Not Overlap" in the dynamic loops deep dive, with this example and the workaround, and a note on the nl.fori_loop API page.
Reproduction (loop_repro.py):
"""Does a dynamic loop overlap iteration i + 1's DMA with iteration i's compute? Minimal repro.
NEURON_PLATFORM_TARGET_OVERRIDE=trn1 NKI_ARTIFACTS_DIR=<dir> python loop_repro.py --mode <mode> --iters N
Every iteration i loads a [128, F] fp32 tile of src[i] from HBM, applies exp on the scalar engine and a multiply-add
on the vector engine, and stores it to out[i]. The modes run the same body:
static: `for i in range(N)` (unrolled at trace time)
fori: nl.fori_loop(0, N, body) with a Python int bound
fori_reg: nl.fori_loop(0, reg, body) with the bound loaded into a register from an input tensor
dyn_range: `for i in nl.dynamic_range(N)` (deprecated in NKI 0.6.0)
while: nl.while_loop on a register counter loaded from an input tensor
fori_u4: nl.fori_loop(0, N // 4, body4), four tiles per iteration unrolled inside the body
The script checks the output against NumPy; time the NEFF afterwards (neuron-bench / neuron-explorer).
"""
import argparse
import numpy as np
import nki
import nki.isa as nisa
import nki.language as nl
F32 = nl.float32
MODES = {"static": 0, "fori": 1, "fori_reg": 2, "fori_u4": 3, "dyn_range": 4, "while": 5}
@nki.jit
def loop_kernel(src, nreg, mode: int):
# src [N, 128 * F] fp32 -> out [N, 128 * F] = exp(src) * 2 + 1
N, PF = src.shape
F = PF // 128
out = nl.ndarray((N, PF), dtype=F32, buffer=nl.shared_hbm)
if mode == 0:
for i in range(N):
t = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.dma_copy(dst=t, src=src.ap(pattern=[[F, 128], [1, F]], offset=i * PF))
e = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.activation(dst=e, op=nl.exp, data=t)
y = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=y, data=e, op0=nl.multiply, operand0=2.0, op1=nl.add, operand1=1.0,
engine=nisa.vector_engine)
nisa.dma_copy(dst=out.ap(pattern=[[F, 128], [1, F]], offset=i * PF), src=y)
else:
def body(i):
t = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.dma_copy(dst=t, src=src.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=i, indirect_dim=0))
e = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.activation(dst=e, op=nl.exp, data=t)
y = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=y, data=e, op0=nl.multiply, operand0=2.0, op1=nl.add, operand1=1.0,
engine=nisa.vector_engine)
nisa.dma_copy(dst=out.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=i, indirect_dim=0), src=y)
def body4(j):
# four iterations per loop pass, unrolled: src viewed as [N / 4, 4 * 128 * F], sub-tile u at offset u * PF
for u in range(4):
t = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.dma_copy(dst=t, src=src4.ap(pattern=[[F, 128], [1, F]], offset=u * PF, scalar_offset=j,
indirect_dim=0))
e = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.activation(dst=e, op=nl.exp, data=t)
y = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=y, data=e, op0=nl.multiply, operand0=2.0, op1=nl.add, operand1=1.0,
engine=nisa.vector_engine)
nisa.dma_copy(dst=out4.ap(pattern=[[F, 128], [1, F]], offset=u * PF, scalar_offset=j, indirect_dim=0),
src=y)
def wbody(r):
# while_loop body: the tile index lives in idx (SBUF), loaded into register ri for the DMA offsets
nisa.register_load(ri, idx)
t = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.dma_copy(dst=t, src=src.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=ri, indirect_dim=0))
e = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.activation(dst=e, op=nl.exp, data=t)
y = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=y, data=e, op0=nl.multiply, operand0=2.0, op1=nl.add, operand1=1.0,
engine=nisa.vector_engine)
nisa.dma_copy(dst=out.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=ri, indirect_dim=0), src=y)
nisa.tensor_scalar(dst=idx, data=idx, op0=nl.add, operand0=1, engine=nisa.vector_engine)
nisa.tensor_scalar(dst=cnt, data=cnt, op0=nl.subtract, operand0=1, engine=nisa.vector_engine)
nisa.register_load(r, cnt)
return r
if mode == 4:
for i in nl.dynamic_range(N):
t = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.dma_copy(dst=t, src=src.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=i, indirect_dim=0))
e = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.activation(dst=e, op=nl.exp, data=t)
y = nl.ndarray((128, F), dtype=F32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=y, data=e, op0=nl.multiply, operand0=2.0, op1=nl.add, operand1=1.0,
engine=nisa.vector_engine)
nisa.dma_copy(dst=out.ap(pattern=[[F, 128], [1, F]], offset=0, scalar_offset=i, indirect_dim=0),
src=y)
elif mode == 5:
cnt = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=cnt, src=nreg)
idx = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.memset(dst=idx, value=0)
ri = nisa.register_alloc()
rw = nisa.register_alloc()
nisa.register_load(rw, cnt)
nl.while_loop(rw, wbody)
elif mode == 3:
src4 = src.reshape((N // 4, 4 * PF))
out4 = out.reshape((N // 4, 4 * PF))
nl.fori_loop(0, N // 4, body4)
elif mode == 1:
nl.fori_loop(0, N, body)
else:
n_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=n_sb, src=nreg)
r = nisa.register_alloc()
nisa.register_load(r, n_sb)
nl.fori_loop(0, r, body)
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--mode", choices=list(MODES), default="fori")
ap.add_argument("--iters", type=int, default=16)
ap.add_argument("--free", type=int, default=2048, help="F: free size of each [128, F] fp32 tile")
args = ap.parse_args()
rng = np.random.default_rng(0)
src = rng.standard_normal((args.iters, 128 * args.free)).astype(np.float32) * 0.5
nreg = np.array([[args.iters]], dtype=np.int32)
out = np.asarray(loop_kernel(src, nreg, mode=MODES[args.mode]))
want = np.exp(src) * 2 + 1
err = float(np.abs(out - want).max() / np.abs(want).max())
print(f"mode={args.mode} iters={args.iters} F={args.free} rel max|err|={err:.2e} {'PASS' if err < 1e-4 else 'FAIL'}")
if __name__ == "__main__":
main()export NEURON_PLATFORM_TARGET_OVERRIDE=trn1 NKI_DISABLE_COMPILE_CACHE=1
for m in static fori fori_u4 dyn_range; do
NKI_ARTIFACTS_DIR=/tmp/loop-$m python loop_repro.py --mode $m --iters 64
(cd /tmp/loop-$m && neuron-bench exec -n 1000 -w 50 --fixed-nc-count 1 --disable-hw-profile -o bench kernel.neff)
done
# device time: nc_latency in /tmp/loop-<mode>/bench/*/info.json
# (fori_reg and while read their trip count from an input, so profile them with neuron-explorer capture and real
# inputs, e.g. `neuron-explorer capture -n kernel.neff -s p.ntff src src.npy nreg nreg.npy`: neuron-bench's zero
# inputs give them zero iterations)
Links
- https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/deep-dives/nki-dynamic-loops.html
- https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/api/generated/nki.language.fori_loop.html
- https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/migration/index.html (0.5 to 0.6 dynamic loop migration)
Activity
github-actions commentedon Oct 5, 2026
Hi @didhd, Thank you for filing the issue! We will take a look and get back to you.
ealiaj-aws commentedon Oct 15, 2026
Hi @didhd, thank you for the contribution. This has been addressed and there will be an update in the near future.