[论文解读] FusionStitching: Deep Fusion and Code Generation for Tensorflow Computations on GPUs

Guoping Long, Jun Yang+2 more|arXiv (Cornell University)|Nov 13, 2018

Advanced Neural Network Applications参考文献 11被引用 6

一句话总结

FusionStitching 是一个针对 TensorFlow 的深度融合与代码生成系统，通过块组合和共享内存优化实现高效的 GPU 内核融合，与 XLA 相比将 GPU 内核启动次数减少了 55%（几何平均值），并在延迟敏感工作负载上将端到端性能提升了最高 20%。

ABSTRACT

In recent years, there is a surge on machine learning applications in industry. Many of them are based on popular AI frameworks like Tensorflow, Torch, Caffe, or MxNet, etc, and are enpowered by accelerator platforms such as GPUs. One important challenge of running Tensorflow computations on GPUs is the fine granularity problem, namely, FLOPS of individual ops are far from enough to fully exploit the computing power of underlying accelerators. The XLA framework provides a solid foundation to explore this problem further. In this paper, we propose FusionStitching, a novel, comprehensive Op fusion and code generation system to stitch computations into large GPU kernels. Experimental results on four public models and two of our large inhouse applications show another 55% (geometric mean) reduction of GPU kernel launches, compared to the XLA fusion baseline. This increases the E2E performance of both of our latency critical inhouse applications up to 20%.

研究动机与目标

• 解决 TensorFlow 中细粒度算子问题，即单个操作的 FLOPS 不足，无法充分利用 GPU 硬件资源。
• 克服 XLA 静态融合规则的局限性，这些规则在处理复杂情况（如高成本的逐元素操作、批量矩阵乘法和布局转换）时失效。
• 通过将计算规范与底层调度细节解耦，实现非 MatMul/Convolution 算子的可扩展、自动融合与代码生成。
• 提出一种新颖的代码生成策略，利用块组合与共享内存，组合具有不同并行循环发射器的内核。
• 通过最小化内核启动开销，在多种模型上实现显著性能提升，特别是在延迟敏感的工业应用场景中。

提出的方法

• 提出关键路径分析以指导融合决策，优先选择数据依赖链中最长路径上的操作。
• 引入领域特定的调度规范语言，抽象低层调优参数，同时支持自动优化。
• 设计新型中间表示发射器 IrEmitterStitched，通过允许多个算子使用独立的并行循环发射器并以共享内存作为中间媒介，支持块组合。
• 利用片上共享内存作为暂存区，实现具有不同调度模式的融合操作之间的高效数据传输。
• 当使用量超过 20KB 阈值时，实现共享内存收缩机制，动态减少内存占用，以维持内核启动的可行性。
• 使用受依赖分析启发的分层跨度图，对复杂计算图中的融合机会进行建模与优化。

实验结果

研究问题

• RQ1 当标准融合启发式方法因复杂依赖关系或内存访问模式而失效时，如何有效融合 TensorFlow 中的细粒度 GPU 操作？
• RQ2 共享内存组合与块级调度对深度学习工作负载中内核融合的可扩展性与性能有何影响？
• RQ3 基于编译器驱动、用户透明的融合与代码生成系统是否能在减少内核启动次数和端到端延迟方面超越现有的 XLA 融合策略？
• RQ4 在多个融合操作之间，共享内存的复用程度在多大程度上可降低内存带宽压力并提升吞吐量？
• RQ5 所提出的融合策略在多样化工作负载（包括公开模型与大规模工业应用）上的可扩展性如何？

主要发现

• 在四个公开模型和两个内部应用中，FusionStitching 相较于 XLA 融合基线，将 GPU 内核启动次数减少了 55%（几何平均值）。
• 在两个延迟敏感的工业应用中，端到端性能最高提升了 20%，证明了其在真实场景中的实际影响。
• 该系统成功处理了 XLA 静态融合规则难以应对的复杂融合案例，如批量矩阵乘法、高成本的逐元素操作和内存布局转换。
• 在语音工作负载中，共享内存使用量最高达到 10,432 字节，在 NMT 模型中，17% 的分配空间被多个算子共享，表明结果复用效果良好。
• 在 100 多个内核中，仅有 3 个触发了收缩机制，表明在大多数工作负载中内存压力可控制在 20KB 限制内。
• LR、W2V、RNN 和 BiRNN 等工作负载的共享内存使用量极低且未触发收缩，表明其融合模式较简单且性能稳定。