[论文解读] DLVM: A modern compiler infrastructure for deep learning systems

Richard Wei, Lane Schwartz+1 more|arXiv (Cornell University)|Nov 8, 2017

Parallel Computing and Optimization Techniques参考文献 4被引用 34

一句话总结

DLVM 是一种现代的、受 LLVM 启发的深度学习编译器基础设施，为张量领域特定语言（DSL）提供了类型安全、模块化且可优化的基础。它通过线性代数中间表示、基于伴随代码生成的系统化自动微分，以及通过 LLVM 针对 GPU 的代码生成，实现了高性能、安全且可组合的深度学习系统。一个 Swift 中的 PoC DSL 展示了端到端的优化与即时编译（JIT）。

ABSTRACT

Deep learning software demands reliability and performance. However, many of the existing deep learning frameworks are software libraries that act as an unsafe DSL in Python and a computation graph interpreter. We present DLVM, a design and implementation of a compiler infrastructure with a linear algebra intermediate representation, algorithmic differentiation by adjoint code generation, domain-specific optimizations and a code generator targeting GPU via LLVM. Designed as a modern compiler infrastructure inspired by LLVM, DLVM is more modular and more generic than existing deep learning compiler frameworks, and supports tensor DSLs with high expressivity. With our prototypical staged DSL embedded in Swift, we argue that the DLVM system enables a form of modular, safe and performant frameworks for deep learning.

研究动机与目标

• 为解决现有深度学习框架的局限性，这些框架通常依赖于不安全的、动态类型的 Python DSL，且缺乏模块化、可组合的编译器设计。
• 提供一种系统化、模块化且可扩展的编译器基础设施，支持功能完整、静态类型的深度学习领域特定语言（DSL）。
• 通过高级优化技术（包括内核融合、代数简化和 AD 检查点）实现高性能执行。
• 将算法自动微分统一为编译器流水线中的第一类转换阶段，从而提升正确性和性能。
• 通过与 LLVM 的集成，支持在多种硬件目标上生成可移植、高效且可扩展的代码。

提出的方法

• 设计并实现一种用于张量计算的领域特定中间表示（IR），采用静态单赋值（SSA）形式并支持控制流，从而支持标准的编译器优化。
• 使用源码转换实现算法自动微分，将伴随代码生成作为第一类编译器阶段，而非依赖操作符重载。
• 与 LLVM 集成以实现低层代码生成，支持透明地针对 CPU 和 GPU（包括通过 NVPTX 支持 NVIDIA GPU）进行编译。
• 在 Swift 中开发一种分阶段 DSL（NNKit），利用 Rep<T> 包装器实现张量计算的静态分析、类型安全和延迟求值。
• 四阶段 JIT 编译流水线：表达式分阶段、形状特化、降低至 DLVM IR，以及函数实例化以在 Swift 中暴露已编译的函数。
• 在一个模块化、可扩展的编译器框架中应用通用和领域特定的优化技术，如内联、循环优化和计算内核融合。

实验结果

研究问题

• RQ1 基于成熟编译技术的现代、模块化编译器基础设施，是否能显著提升深度学习框架的可靠性与性能？
• RQ2 如何将算法自动微分作为第一类转换阶段集成到深度学习编译器中，而非作为库级别的功能？
• RQ3 在系统编程语言（如 Swift）中使用静态类型、分阶段 DSL，能在多大程度上同时实现安全的原型设计与高性能执行？
• RQ4 统一的 IR 和优化流水线是否能有效支持从高层神经网络计算到低层内核生成，并适配多种硬件目标？
• RQ5 与基于图的表示相比，IR 中使用 SSA 形式和控制流图如何提升深度学习 IR 的表达力和优化潜力？

主要发现

• DLVM 通过受 LLVM 启发的系统化编译器基础设施，成功实现了模块化、安全且高性能的深度学习框架。
• 将算法自动微分作为编译器阶段而非库级功能集成，使得梯度计算的优化更加激进且正确。
• 在 Swift 中实现的分阶段 DSL（NNKit）实现了完整的静态分析与类型安全，同时通过 DLVM 实现了高效的 JIT 编译与代码生成。
• DLVM 的 IR 支持高级优化技术，如内核融合、代数简化和 AD 检查点，显著降低了内核启动开销并提升了性能。
• 该系统展示了从高层张量表达式到优化后的 GPU 内核的端到端编译能力，支持 NVIDIA GPU，并可扩展至其他后端（如 HPVM）。
• 原型实现表明，基于编译器的方法能够统一并简化深度学习系统的开发，降低现有框架的复杂性与脆弱性。