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