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

Richard Wei, Vikram Adve+1 more|arXiv (Cornell University)|Nov 8, 2017

Parallel Computing and Optimization Techniques被引用 2

一句话总结

DLVM 是一种现代的、模块化的深度学习编译器基础设施，采用线性代数中间表示（IR），通过伴随代码生成实现自动微分，支持领域特定优化，并基于 LLVM 实现 GPU 代码生成。DLVM 以 Swift 中的分阶段领域特定语言（DSL）构建，可实现安全、高性能且可组合的深度学习框架，具备强大的类型安全性和表达能力。

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 以运行时解释执行的方式构建。
• 设计一种受 LLVM 启发的模块化、可扩展的编译器基础设施，支持具有强类型安全性的高级张量 DSL。
• 利用 LLVM 的优化和代码生成管道，实现高效、正确且可移植的 GPU 代码生成。
• 通过伴随代码生成支持算法微分，确保梯度计算的正确性和性能。
• 证明在 Swift 中的分阶段 DSL 可在深度学习系统中同时实现高表达性和高性能。

提出的方法

• 设计一种原生表达张量操作并支持高级优化的线性代数中间表示（IR）。
• 通过伴随代码生成实现算法微分，系统性地从正向计算推导出梯度代码。
• 集成领域特定优化，如计算融合、布局转换和分块，以提升性能。
• 利用 LLVM 的基础设施进行代码生成，实现可移植且高效的 GPU 内核编译。
• 在 Swift 中构建分阶段 DSL，以实现具有完整类型安全性和模块化的深度学习模型表达。
• 采用模块化编译器管道，使每个阶段（前端、IR、优化、代码生成）均可组合且可重用。

实验结果

研究问题

• RQ1 基于线性代数 IR 的编译器基础设施是否能在深度学习工作负载中同时实现高性能和高表达性？
• RQ2 与其它自动微分技术相比，伴随代码生成在正确性和性能方面表现如何？
• RQ3 在模块化编译栈中，领域特定优化在多大程度上能提升推理和训练效率？
• RQ4 在类似 Swift 的系统语言中，分阶段 DSL 是否能同时提供低级代码的安全性和性能，同时保持高级 DSL 的开发效率？
• RQ5 在实践中，DLVM 的模块化与通用性与单体式深度学习框架相比表现如何？

主要发现

• DLVM 的线性代数 IR 支持高阶、可组合且可优化的深度学习模型表示，具备强类型安全性。
• 伴随代码生成确保了梯度计算的正确性与高效性，避免了源到源或操作符重载 AD 的常见陷阱。
• 如计算融合和布局感知转换等领域的特定优化，显著提升了 GPU 内核的性能。
• Swift 中的分阶段 DSL 证明了在不牺牲表达性的情况下，可实现高阶、安全且高性能的深度学习框架。
• 通过重用 LLVM 的优化和代码生成管道，DLVM 实现了可移植且高效的 GPU 代码生成，性能表现具有竞争力。