060、Enzyme自动微分框架与MLIR的集成
060、Enzyme自动微分框架与MLIR的集成一、一个让我熬夜到凌晨三点的梯度问题去年秋天,我在给一个边缘计算设备写模型推理加速器。模型里有个自定义算子,用C++手写了前向,反向梯度死活算不对。手动推导链式法则?矩阵维度一多,脑子就炸。用PyTorch的autograd?算子太底层,hook进去性能损失30%。当时团队里有个从LLVM社区回来的兄弟,甩给我一句话:“试试Enzyme,直接对LLVM IR做自动微分,不用碰源码。”那是我第一次认真看Enzyme。后来发现,这玩意儿不仅能对LLVM IR做AD,还能跟MLIR深度集成——在MLIR的IR层面直接注入梯度计算逻辑,比传统基于tracing或source-to-source的方案优雅得多。今天这篇笔记,就聊聊Enzyme与MLIR集成的技术细节,以及我在实际集成中踩过的坑。二、Enzyme到底在做什么传统自动微分框架,比如PyTorch的autograd,是在Python层面构建计算图,然后反向传播。Enzyme不一样,它工作在编译器中间表示层面——直接对LLVM IR做微分。这意味着你不需要修改源码,不需要定义梯度函数,只要给一个函数入口,Enzyme就能自动生成它的梯度函数。举个例子,你写了个C函数:doublefoo