nki.isa.register_alloc#

nki.isa.register_alloc(x=None)[source]#

Allocate a virtual register and optionally initialize it with an integer value x.

Each engine sequencer (Tensor/Scalar/Vector/GpSimd/Sync Engine) within a NeuronCore maintains its own set of physical registers for scalar operations (64x 32-bit registers per engine sequencer in NeuronCore v2-v4). The nisa.register_alloc API conceptually allocates a register within a virtual register space. Users do not need to expliclity free a register through nisa APIs. The NKI compiler handles physical register allocation (and deallocation) across the appropriate engine sequencers based on the dynamic program flow.

NKI provides the following APIs to manipulate allocated registers:

  • nisa.register_move: Move a constant value into a register

  • nisa.register_load: Load a scalar (32-bit) value from HBM/SBUF into a register

  • nisa.register_store: Store register contents to HBM/SBUF

In the current NKI release, these registers are primarily used to specify dynamic loop boundaries and while loop conditions. The NKI compiler compiles such dynamic looping constructs to branching instructions executed by engine sequencers. For additional details, see nl.dynamic_range. For more information on engine sequencer and its capabilities, see Trainium/Inferentia2 architecture guide.

Parameters:

  • dst – a virtual register object

  • x – optional integer value to initialize the register with