This document is relevant for: Trn2, Trn3

About the NKI Compiler#

This topic covers the NKI Compiler and how it interacts with the Neuron Graph Compiler to produce a complete model. The NKI Compiler is responsible for compiling NKI kernels.

Overview#

The NKI language allows kernel writers to have direct, fine grained control over Neuron devices. Through low level APIs that reflect the Neuron instruction set architecture (ISA), NKI empowers developers to take direct control over critical performance optimizations during kernel development. This approach requires a dedicated NKI Compiler, separate from the existing Neuron Graph Compiler, which compiles kernel code while preserving the developer’s optimization choices. To seamlessly integrate NKI into model architectures defined in machine learning frameworks like JAX and PyTorch, the NKI Compiler also works in conjunction with the Neuron Graph compiler.

The diagram below shows the detailed compilation flow inside the Neuron compilers and how they work together to build the overall binary that is executable on Neuron hardware. The NKI Compiler first specializes the kernel by tracing the kernel code with Python. During tracing, NKI ISA API calls emit NKI IR. It then performs a small number of back end transformations on the NKI IR, optimizing resource allocations and instruction scheduling, producing optimized NKI IR that gets integrated back into the overall model.

../../_images/about-nki-3.png

Important

When the Python interpreter encounters a top-level function that has the @nki.jit decorator, NKI specializes the function by tracing it with Python. Because the trace runs in Python, the code that runs during the trace can use all of the Python language.

# this is a Python function that calls 'kernel', which is a NKI kernel
def a_function(x,y,z):
    kernel(x, y, z)

# this is a NKI kernel that will be compiled by the NKI Compiler and
# integrated back into the overall model by the Neuron Graph compiler
@nki.jit
def kernel(x,y,z):
    # this is kernel code

If you set NKI_COMPILER_FRONTEND=custom, NKI uses the NKI meta-programming language from releases before 0.7. This language is a limited subset of Python. Using a Python feature that this language does not support will result in a compile error. This will be removed in a future version of NKI. For more information, see Migrating from NKI 0.6.0 to NKI 0.7.0.

How the NKI Compiler Works with the Graph Compiler#

For each kernel function, Neuron runs the NKI Compiler to produce an artifact for that kernel function. This is similar to compiling a single file with a traditional compiler, such as a C++ compiler.

All of the kernel artifacts are managed by the Neuron SDK. Programmers do not need to manage these files themselves. Similar to prior versions of NKI, programmers mark kernel functions with nki.jit—the NKI Compiler will be invoked automatically when this decorator is encountered during compilation.

The Neuron Graph Compiler (or just the Neuron Compiler) handles the rest of the model, which we refer to as “the compute graph”. The framework, such as PyTorch or Jax, orchestrates the process of building a compute graph from the model definition. When the model includes a call to a NKI kernel function, the NKI Compiler will insert a reference to the compiled artifact into the graph. The Graph Compiler recognizes these references and assembles the final result that can be run on the Trainium Hardware.

Integration#

As described above, both the NKI Compiler and the Neuron Compiler are used to construct the final artifact that can be run on Trainium hardware. The NKI Compiler compiles each NKI kernel function in turn, and the Neuron Compiler compiles the whole model and inserts the NKI kernels based on the references generated by the NKI Compiler.

This insertion of NKI kernels into the graph is done very late in the compilation process. This is different from prior versions of NKI that integrated NKI kernels earlier in the compile process. Insertion later in the process allows the NKI Compiler to provide custom behavior for NKI and give users a more predictable and performant result.

Further reading#

This document is relevant for: Trn2, Trn3