论文
Nova:面向深度学习的端到端MLIR编译器
Nova: An End-to-End MLIR Compiler for Deep Learning
摘要
大规模深度学习模型的性能高度依赖于高层数学运算能否有效映射到底层物理硬件。高层张量框架虽然为模型设计提供灵活抽象,但即时执行模式缺乏全图视野及精细的硬件、内存控制,难以充分利用硬件。为弥合这一差距,我们设计Nova,一个自动化端到端JIT编译器,目标是控制硬件映射:跨算子边界融合、优化复杂内存层次,并将执行调优到寄存器级别。Nova捕获即时执行,将前向与反向计算统一到单一值语义方言中,以进行全图优化;再用Analytic Configurator按算术强度确定性推导执行调度,将搜索时间降至零。结构哈希运行时依据计算结构直接合成细粒度计算核。在RTX 3060的评测中,Nova在多数形状的TF32矩阵乘法上达到或略高于cuBLAS与XLA,保持小于5e-4的相对误差。在4200万参数模型上,相较PyTorch和XLA,吞吐量最多分别提升10.6%和4.4%,且不牺牲数值精度。相较PyTorch,内存占用最多下降29%,使Nova能在同一张12GB消费级GPU上以每秒17900个Token训练1.44亿参数模型,而PyTorch出现内存不足。
