Repository files navigation

Python-first Development for AI Compilers.

• 📥 Installation
• 🔑 Key Features
  • 🏗️ Define IRs with MLC Dataclasses
  • 🐍 Design Python-based Text Formats for IRs
  • 🎯 Test IRs with MLC Structure-Aware Tooling
  • ⚡ Migrate Gradually to C++ with MLC Plugins
• ⛽ Development
  • ⚙️ Editable Build
  • 🎡 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.

pip install -U mlc-python

MLC provides Pythonic dataclasses:

importmlc.dataclassesasmlcd@mlcd.py_class("demo.MyClass")classMyClass(mlcd.PyClass):a:intb:strc:float|Noneinstance=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:mustberealnumber,notstr>>>instance.non_exist=1AttributeError:'MyClass'objecthasnoattribute'non_exist'andno__dict__forsettingnewattributes

Serialization. MLC dataclasses are picklable and JSON-serializable.

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

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 aFunc 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 + bAssign(lhs=e,rhs=Add(d,c)),# e = d + c  ],ret=e)

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

>>>print(mlcp.to_python(f))# Stringify to Pythondeff(a,b,c):d=a+be=d+creturne

• Methodmlc.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__.

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

importmlc.dataclassesasmlcd@mlcd.py_classclassExpr(mlcd.PyClass):def__add__(self,other):returnAdd(a=self,b=other)@mlcd.py_class(structure="nobind")classAdd(Expr):a:Exprb:Expr@mlcd.py_class(structure="var")classVar(Expr):name:str=mlcd.field(structure=None)# excludes `name` from defined structure@mlcd.py_class(structure="bind")classLet(Expr):rhs:Exprlhs:Var=mlcd.field(structure="bind")# `Let.lhs` is the def-sitebody:Expr

Structural equality. Member methodeq_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:Structuralequalitycheckfailedat {root}.rhs.b:Inconsistentbinding.RHShasbeenboundtoadifferentnodewhileLHSisnotbound

Structural hashing. The structure of MLC dataclasses can be hashed viahash_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()>>>assertL1_hash==L2_hash>>>assertL1_hash!=L3_hash

(🚧 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.

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

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

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