nki.isa.register_store#

nki.isa.register_store(dst, src)[source]#

Store the value from a virtual register into memory (HBM/SBUF).

This instruction writes the scalar value (up to 32-bit) stored in a virtual register to a memory location (HBM or SBUF). The destination must be a tensor with exactly one element (shape [1] or [1, 1]). This enables saving register values back to memory for later use or for output purposes.

The virtual register system allows the NKI compiler to allocate physical registers across different engine sequencers as needed. See nisa.register_alloc for more details on virtual register allocation.

Parameters:

  • dst – the destination tensor with a single element to store the register value

  • src – the source virtual register (allocated via nisa.register_alloc)

Example:

# Store a register value back to memory
counter_reg = nisa.register_alloc(0)
# ... perform operations that modify counter_reg ...
result_tensor = nl.ndarray([1], dtype=nl.int32, buffer=nl.sbuf)
nisa.register_store(result_tensor, counter_reg)