mlc-python

MLC-Python

Python-first Development for AI Compilers.

• :inbox_tray: Installation
• :key: Key Features
  • :building_construction: Define IRs with MLC Dataclasses
  • :snake: Design Python-based Text Formats for IRs
  • :dart: Test IRs with MLC Structure-Aware Tooling
  • :zap: Migrate Gradually to C++ with MLC Plugins
• :fuelpump: Development
  • :gear: Editable Build
  • :ferris_wheel: Create Wheels

MLC is a Python-first toolkit that streamlines the development of AI compilers, runtimes, and compound AI systems with its Pythonic dataclasses, structure-aware tooling, and Python-based text formats.

Beyond pure Python, MLC natively supports zero-copy interoperation with C++ plugins, and enables a smooth engineering practice transitioning from Python to hybrid or Python-free development.

:inbox_tray: Installation

pip install -U mlc-python

:key: Key Features

:building_construction: Define IRs with MLC Dataclasses

MLC provides Pythonic dataclasses:

import mlc.dataclasses as mlcd

@mlcd.py_class("demo.MyClass")
class MyClass(mlcd.PyClass):
  a: int
  b: str
  c: float | None

instance = MyClass(12, "test", c=None)

Type safety. MLC dataclass enforces strict type checking using Cython and C++.

>>> instance.c = 10; print(instance)
demo.MyClass(a=12, b='test', c=10.0)

>>> instance.c = "wrong type"
TypeError: must be real number, not str

>>> instance.non_exist = 1
AttributeError: 'MyClass' object has no attribute 'non_exist' and no __dict__ for setting new attributes

Serialization. MLC dataclasses are picklable and JSON-serializable.

>>> MyClass.from_json(instance.json())
demo.MyClass(a=12, b='test', c=None)

>>> import pickle; pickle.loads(pickle.dumps(instance))
demo.MyClass(a=12, b='test', c=None)

:snake: Design Python-based Text Formats for IRs

Printer. MLC looks up method __ir_print__ to convert IR nodes to Python AST:

[Example]. Copy the toy IR definition to REPL and then create a Func node below:

>>> a, b, c, d, e = Var("a"), Var("b"), Var("c"), Var("d"), Var("e")
>>> f = Func("f", [a, b, c],
  stmts=[
    Assign(lhs=d, rhs=Add(a, b)),  # d = a + b
    Assign(lhs=e, rhs=Add(d, c)),  # e = d + c
  ],
  ret=e)

• Method mlc.printer.to_python converts an IR node to Python-based text;

>>> print(mlcp.to_python(f))  # Stringify to Python
def f(a, b, c):
  d = a + b
  e = d + c
  return e

• Method mlc.printer.print_python further renders the text with proper syntax highlighting. [Screenshot]

>>> mlcp.print_python(f)  # Syntax highlighting

AST Parser. MLC has a concise set of APIs for implementing parser with Python's AST module, including:

• Inspection API that obtains source code of a Python class or function and the variables they capture;
• Variable management APIs that help with proper scoping;
• AST fragment evaluation APIs;
• Error rendering APIs.

[Example]. With MLC APIs, a parser can be implemented with 100 lines of code for the Python text format above defined by __ir_printer__.

:dart: Test IRs with MLC Structure-Aware Tooling

By annotating IR definitions with structure, MLC supports structural equality and structural hashing to detect structural equivalence between IRs:

Define a toy IR with `structure`.

import mlc.dataclasses as mlcd

@mlcd.py_class
class Expr(mlcd.PyClass):
  def __add__(self, other):
    return Add(a=self, b=other)

@mlcd.py_class(structure="nobind")
class Add(Expr):
  a: Expr
  b: Expr

@mlcd.py_class(structure="var")
class Var(Expr):
  name: str = mlcd.field(structure=None) # excludes `name` from defined structure

@mlcd.py_class(structure="bind")
class Let(Expr):
  rhs: Expr
  lhs: Var = mlcd.field(structure="bind") # `Let.lhs` is the def-site
  body: Expr

Structural equality. Member method eq_s compares the structural equality (alpha equivalence) of two IRs represented by MLC's structured dataclass.

>>> x, y, z = Var("x"), Var("y"), Var("z")
>>> L1 = Let(rhs=x + y, lhs=z, body=z)  # let z = x + y; z
>>> L2 = Let(rhs=y + z, lhs=x, body=x)  # let x = y + z; x
>>> L3 = Let(rhs=x + x, lhs=z, body=z)  # let z = x + x; z
>>> L1.eq_s(L2)
True
>>> L1.eq_s(L3, assert_mode=True)
ValueError: Structural equality check failed at {root}.rhs.b: Inconsistent binding. RHS has been bound to a different node while LHS is not bound

Structural hashing. The structure of MLC dataclasses can be hashed via hash_s, which guarantees if two dataclasses are alpha-equivalent, they will share the same structural hash:

>>> L1_hash, L2_hash, L3_hash = L1.hash_s(), L2.hash_s(), L3.hash_s()
>>> assert L1_hash == L2_hash
>>> assert L1_hash != L3_hash

:zap: Migrate Gradually to C++ with MLC Plugins

(🚧 Under construction)

MLC seamlessly supports zero-copy bidirectional interoperabilty with C++ plugins with no extra dependency. By gradually migrating classes and methods one at a time, a pure Python prototype can be transitioned to hybrid or Python-free development.

:fuelpump: Development

:gear: Editable Build

pip install --verbose --editable ".[dev]"
pre-commit install

:ferris_wheel: Create Wheels

This project uses cibuildwheel to build cross-platform wheels. See .github/workflows/wheels.ym for more details.

export CIBW_BUILD_VERBOSITY=3
export CIBW_BUILD="cp3*-manylinux_x86_64"
python -m pip install pipx
pipx run cibuildwheel==2.22.0 --output-dir wheelhouse