sueszli/llvm-to-air
GitHub: sueszli/llvm-to-air
首个开源的 LLVM IR 到 Apple AIR 编译器后端，实现了 Apple Silicon GPU 的端到端编译流水线。
Stars: 6 | Forks: 0
▌ ▌ ▌ ▌▙▗▌ ▐ ▞▀▖▜▘▛▀▖
▌ ▌ ▚▗▘▌▘▌ ▜▀ ▞▀▖ ▙▄▌▐ ▙▄▘
▌ ▌ ▝▞ ▌ ▌ ▐ ▖▌ ▌ ▌ ▌▐ ▌▚
▀▀▘▀▀▘▘ ▘ ▘ ▀ ▝▀ ▘ ▘▀▘▘ ▘
// Apple Silicon GPU 的逆向工程编译器栈
//
// 即将在 xDSL 中推出：
// https://github.com/xdslproject/xdsl/blob/main/xdsl/backend/mps/__init__.py
MLIR 是跨异构硬件实现可移植性能的正确抽象，但它没有 Apple Silicon GPU 后端 [^1]。这是一个问题，因为 (1) Apple Silicon 的市场份额增长迅速，(2) 这些机器大多拥有闲置的强大 GPU，(3) 从蛋白质折叠到 ML 训练，世界需要更多的算力。
Mojo 证明了通过 MLIR->LLVM->AIR->MetalLib 的路径来针对 Apple GPU 是可行的，但他们的实现是闭源的 [^2]。
本项目逆向工程了这一缺失的环节，并提供了 LLVM IR 到 AIR lowering pass 的开源实现。
+----------------+ +----------------+ +----------------+
| Frontend |----->| MLIR Dialect |----->| LLVM Bitcode |
| | | | | (Open Source) |
+----------------+ +----------------+ +----------------+
|
[ src/llvm_to_air.py ]
|
+----------------+ +----------------+ +--------v-------+
| Apple GPU |<-----| Metallib |<-----| AIR Bitcode |
| (M-Series) | | (Binary) | | (Proprietary) |
+----------------+ +----------------+ +----------------+
核心贡献是 `src/llvm_to_air.py`，它接收 LLVM IR 并将其 lowering 到 Apple 的中间表示 AIR。这使得从高级 MLIR 方言到 Apple Silicon GPU 上可执行代码的完整编译流水线成为可能。我使用 xDSL 用 Python 编写了整个编译器栈，使其易于访问和修改。
友情提示：这是实验性的且脆弱的。AIR 是闭源且无文档的，所以这里的一切都是逆向工程的。但它确实有效，据我所知，这是第一个针对 Apple Silicon 的端到端开源栈。
## 性能
mandelbrot 基准测试显示，相比原生 Python 实现，实现了 1150 倍的加速 🔥。
$ uv run demo_mandelbrot.py
```
mandelbrot benchmark (1,048,576 pixels)
results (avg latency ms):
gpu : 2.47 ms
numba : 188.56 ms
numpy : 1519.57 ms
numpy+numba : 1820.99 ms
plain : 2840.38 ms
relative to vanilla python:
gpu : 1150.23x faster
numba : 15.06x faster
numpy : 1.87x faster
numpy+numba : 1.56x faster
```
## Lisp 前端
还有一个微型的 Common Lisp 子集作为前端。
$ uv run demo_linalg.py
```
(print
(add
(matmul
(tensor (2 3) (-1.0 2.0 -3.0 4.0 -5.0 6.0))
(tensor (3 2) (7.0 8.0 9.0 10.0 11.0 12.0))
)
(tensor (2 2) (100.0 100.0 100.0 100.0))
)
)
Tensor(2 x 2):
78.000000 76.000000
149.000000 154.000000
```
## 参考文献
[^1] MLIR: https://discourse.llvm.org/t/rfc-mps-dialect-in-mlir/77102
[^2] Mojo: https://forum.modular.com/t/apple-silicon-gpu-support-in-mojo/2295
标签：AIR, Apple Silicon, GPU编译器, LLVM, Metal, Metallib, MLIR, M系列芯片, Python, xDSL, 云资产清单, 代码生成, 底层开发, 开源编译器, 异构计算, 无后门, 机器学习基础设施, 渗透测试工具, 编译器中间表示, 编译器栈, 苹果芯片, 逆向工具, 逆向工程