Skip to content

leanprover/KLR

Repository files navigation

Kernel Language Representation (KLR)

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.

Interop

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;

  1. Python parser parses kernel.py to JSON (using reflection)
  2. Python writes kernel.json
  3. KLR reads kernel.json
  4. KLR writes klr.json
  5. Python reads klr.json into a data structure

Steps to make a new version/wheel

  1. Bump the build or minor version in
  • interop/pyproject.toml (Deployment to PyPI will fail if you forget this.)
  • Main.klrCmd (Nothing will break if you don't, but we'd like to be consistent)
  1. Create a git tag of the form v1.2.3 and push it to KLR repo

This should trigger a build that uploads the artifacts to pypi.

About

A formalization of ML kernel languages

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Contributors 4

  •  
  •  
  •  
  •