[论文解读] Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Keno Fischer, Elliot Saba|arXiv (Cornell University)|Oct 23, 2018

Distributed and Parallel Computing Systems参考文献 4被引用 9

一句话总结

本文提出一种基于编译器的方法，通过XLA编译器自动将完整的Julia程序（包括VGG19等神经网络模型及其梯度）卸载到Google Cloud TPU。该方法利用Julia的静态分析和类型推断能力，将整个前向和反向传播过程融合为单个优化的TPU内核，实现100张图像批量处理时227倍的加速（0.23秒 vs. 52.4秒），代码量少于1000行，且无需修改Julia编译器或核心包。

ABSTRACT

Google's Cloud TPUs are a promising new hardware architecture for machine learning workloads. They have powered many of Google's milestone machine learning achievements in recent years. Google has now made TPUs available for general use on their cloud platform and as of very recently has opened them up further to allow use by non-TensorFlow frontends. We describe a method and implementation for offloading suitable sections of Julia programs to TPUs via this new API and the Google XLA compiler. Our method is able to completely fuse the forward pass of a VGG19 model expressed as a Julia program into a single TPU executable to be offloaded to the device. Our method composes well with existing compiler-based automatic differentiation techniques on Julia code, and we are thus able to also automatically obtain the VGG19 backwards pass and similarly offload it to the TPU. Targeting TPUs using our compiler, we are able to evaluate the VGG19 forward pass on a batch of 100 images in 0.23s which compares favorably to the 52.4s required for the original model on the CPU. Our implementation is less than 1000 lines of Julia, with no TPU specific changes made to the core Julia compiler or any other Julia packages.

研究动机与目标

• 实现通用Julia程序（包括机器学习模型）的全自动编译，无需修改Julia编译器或核心包，即可在云端TPU上执行。
• 通过与Zygote.jl集成，支持模型的端到端自动微分，实现前向和反向传播的TPU卸载。
• 证明Julia的高级语言特性（如多分派、高阶函数和泛型编程）可被有效编译为TPU可执行的XLA中间表示（IR）。
• 通过将整个程序区域编译为单个内核，实现复杂工作负载（如训练循环）在TPU上的高效融合执行。
• 提供可扩展、可维护且用户友好的Julia代码到TPU的编译接口，克服基于追踪或框架专用方法的局限性。

提出的方法

• 利用Julia编译器的静态分析和类型推断能力，直接将Julia代码编译为XLA中间表示（IR），避免运行时追踪。
• 利用Julia的多分派和泛型编程机制，表达并优化从Julia IR到XLA IR的转换，实现类型泛化的编译。
• 通过分析数据依赖关系并应用静态优化，在卸载前将控制流结构（如循环、条件语句）融合为XLA操作。
• 通过与Zygote.jl的源到源自动微分机制集成，生成并卸载前向和反向传播，作为单个融合的TPU内核。
• 使用XLA编译器的通用中间表示作为目标，支持TPU执行，实现对非TensorFlow工作负载和以线性代数为主的非ML应用的支持。
• 采用infeed/outfeed操作管理CPU与TPU之间的数据传输，并通过StructsOfArrays（SoA到AoS转换）支持XLA的类型约束。

实验结果

研究问题

• RQ1 是否可以无需运行时追踪，自动将包含控制流和高级抽象的完整Julia程序编译为TPU可执行的XLA内核？
• RQ2 Julia的多分派和泛型编程特性在多大程度上可被用于生成高效、类型特定的XLA内核？
• RQ3 Julia的编译器基础设施与Zygote的自动微分相结合，是否能够实现深度学习模型前向和反向传播的端到端TPU卸载？
• RQ4 完全卸载的VGG19模型在性能上与CPU执行相比如何？在云端TPU上可实现多大的加速？
• RQ5 通过XLA将通用Julia代码编译到TPU时，调试困难、类型推断失败和分布式执行支持不足等关键限制是什么？

主要发现

• 该方法成功将VGG19模型的完整前向传播编译为单个融合的TPU内核，在100张图像批量处理中实现227倍于CPU执行的加速（0.23秒 vs. 52.4秒）。
• 通过结合Zygote.jl的自动微分和XLA的融合能力，该方法支持前向和反向传播的端到端卸载，包括训练循环。
• 实现代码少于1000行Julia代码，且无需修改核心Julia编译器或外部包。
• 该方法实现了稳定的性能测量，TPU计算时间（FluXLA TPU）显著比总时钟时间更一致。
• 该方法可推广至非ML的Julia代码，只要工作负载以线性代数运算为主，表明其在深度学习之外具有更广泛的应用潜力。
• 当前实现面临类型推断失败时的可调试性挑战，以及对XLA分布式原原子支持有限的问题，这些被识别为未来改进的关键方向。