nki.isa.register_load#

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

Load a scalar value from memory (HBM or SBUF) into a virtual register.

This instruction reads a single scalar value (up to 32-bit) from a memory location (HBM or SBUF) and stores it in the specified virtual register. The source must be a NKI tensor with exactly one element (shape [1] or [1, 1]). This enables dynamic loading of values computed at runtime into registers for use in control flow operations.

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 virtual register (allocated via nisa.register_alloc)

  • src – the source tensor containing a single scalar value to load

Example:

# Load a computed value into a register
computed_bound = nl.ones([1], dtype=nl.int32, buffer=nl.sbuf)  # bound of 1 in SBUF
loop_reg = nisa.register_alloc()
nisa.register_load(loop_reg, computed_bound)