This document is relevant for: Trn2, Trn3

nki.isa.VirtualRegister#

class nki.isa.VirtualRegister(state=None, frozen=False)[source]#

A virtual register on engine.

Allocated via nisa.register_alloc() and manipulated via nisa.register_move(), nisa.register_load(), nisa.register_store().

Virtual registers represent registers on engine and are used for various APIs such as loading and storing constants from tensors, as the return value of nki.collective and nki.isa APIs, and for dynamic addressing.

In addition to NKI APIs, virtual registers can be used to represent dynamic loop bounds for for loops using dynamic_range, and while loops.

Note

nl.dynamic_range and bare while reg: are supported by the parser frontend. Under the tracer frontend a register’s value is not known while the kernel is being built, so using a register in Python control flow (if/while) or converting it with int() / range() raises TypeError. For the tracer, use nl.fori_loop for counted loops and nki.experimental.while_loop for data-dependent loops.

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

# Using a register in a dynamic for loop.
reg = nisa.register_alloc(5)
for _ in nl.dynamic_range(reg):
    tile = nl.load(input_tensor[0:128, 0:512])
    result = nl.multiply(tile, tile)
    nl.store(out_tensor[0:128, 0:512], result)

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

# Using a register in a dynamic while loop.
cond_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=cond_sb, src=...)

# Load condition into register
reg = nisa.register_alloc()
nisa.register_load(reg, cond_sb)

while reg:
    ...
    nisa.dma_copy(dst=cond_sb, src=...)
    nisa.register_load(reg, cond_sb)

Methods

This document is relevant for: Trn2, Trn3