This document is relevant for: Trn2, Trn3

nki.isa.VirtualRegister#

class nki.isa.VirtualRegister[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 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.

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