We accidentally built an LLVM compiler for Jax
2 days ago
- PennyLane的Catalyst量子编译器意外地支持纯经典JAX代码，将其编译为MLIR和LLVM，不包含量子指令。
- 团队选择JAX是因为它可以追踪Python函数，支持NumPy和SciPy API，并降级到MLIR，同时添加了原生Python控制流和动态形状。
- 目标是将整个混合量子-经典工作流编译成独立的可执行文件，而不仅仅是量子指令。
- Catalyst绕过XLA，直接将JAX表示降级为标准MLIR方言，然后到LLVM，从而能够使用LLVM工具（如Enzyme）进行反向传播。
- 好处包括无需机器学习运行时依赖的独立AOT二进制文件、支持动态形状数组、原生Python控制流以及易于自定义MLIR passes。
- 这种方法可能对边缘机器学习、自定义硬件和编译器基础设施有用，尽管在标准深度学习工作负载上无法超越XLA。