nki.compiler.enable_stack_allocator

nki.compiler.enable_stack_allocator(func=None, log_level=50)

Use stack allocator to allocate the psum and sbuf tensors in the kernel.

Must use together with skip_middle_end_transformations.

from neuronxcc import nki

@nki.compiler.enable_stack_allocator
@nki.compiler.skip_middle_end_transformations
@nki.jit
def kernel(...):
  ...