This repository contains an implementation of KLR, a core language and elaborators for machine learning kernels. The goal of KLR is to define a common representation for kernel functions with a precise formal semantics along with translations from common kernel languages to the KLR core language. The initial focus of KLR is the Neuron Kernel Interface, and the Trainium hardware.
The KLR compiler starts with Python code (e.g. NKI kernels), converts the source code to JSON and passes it to the Lean parser. The lean parser converts (aka traces) the Python AST into KLR. As such, we have an external dependency on a Python runtime. To keep these processes as separate as possible, we just use a simple file-IO pipeline;
- Python parser parses kernel.py to JSON (using reflection)
- Python writes kernel.json
- KLR reads kernel.json
- KLR writes klr.json
- Python reads klr.json into a data structure