[论文解读] AsyncTaichi: Whole-Program Optimizations for Megakernel Sparse Computation and Differentiable Programming.

Yuanming Hu, Mingkuan Xu+2 more|arXiv (Cornell University)|Dec 15, 2020

Parallel Computing and Optimization Techniques参考文献 16被引用 1

一句话总结

本文提出 AsyncTaichi，一个针对 Taichi 的全程序优化框架，可减少稀疏计算与可微分计算中不必要的列表生成。通过在运行时分析整个计算图并跨内核优化，该框架在无需修改代码的情况下，使稀疏网格模拟与可微分编程工作负载的内核启动次数减少 3.07–3.90 倍，性能提升 1.73–2.76 倍。

ABSTRACT

We present a whole-program optimization framework for the Taichi programming language. As an imperative language tailored for sparse and differentiable computation, Taichi's unique computational patterns lead to attractive optimization opportunities that do not present in other compiler or runtime systems. For example, to support iteration over sparse voxel grids, excessive list generation tasks are often inserted. By analyzing sparse computation programs at a higher level, our optimizer is able to remove the majority of unnecessary list generation tasks. To provide maximum programming flexibility, our optimization system conducts on-the-fly optimization of the whole computational graph consisting of Taichi kernels. The optimized Taichi kernels are then just-in-time compiled in parallel, and dispatched to parallel devices such as multithreaded CPU and massively parallel GPUs. Without any code modification on Taichi programs, our new system leads to $3.07 - 3.90 imes$ fewer kernel launches and $1.73 - 2.76 imes$ speed up on our benchmarks including sparse-grid physical simulation and differentiable programming.

研究动机与目标

• 解决由于 Taichi 中过度生成列表而导致的稀疏计算与可微分编程性能瓶颈。
• 在无需修改代码的前提下，实现跨多个 Taichi 内核的全程序优化。
• 减少内核启动开销，并提升在 CPU 和 GPU 等并行设备上的执行效率。
• 通过运行时优化，支持灵活且高性能的稀疏计算工作负载执行。

提出的方法

• 在高层分析 Taichi 程序的整个计算图，以识别冗余的列表生成任务。
• 在执行过程中进行运行时的全程序优化，对 Taichi 内核进行动态优化。
• 通过即时编译将优化后的内核并行编译，以便在多线程 CPU 和 GPU 上执行。
• 通过跨内核分析稀疏网格迭代模式，消除冗余的列表生成。
• 动态优化执行图，以最小化内核启动次数并提升数据局部性。

实验结果

研究问题

• RQ1 全程序优化如何减少稀疏 Taichi 程序中的冗余列表生成？
• RQ2 在不修改代码的前提下，跨多个内核优化能带来多大的性能提升？
• RQ3 运行时优化在多大程度上能减少稀疏与可微分工作负载中的内核启动开销？
• RQ4 该优化框架在物理模拟与可微分编程等多样化工作负载中的有效性如何？

主要发现

• 该框架在基准工作负载中使内核启动次数减少了 3.07 至 3.90 倍。
• 在稀疏网格物理模拟与可微分编程任务中，实现了 1.73 至 2.76 倍的性能提升。
• 优化过程透明执行，无需对现有 Taichi 程序进行任何修改。
• 该系统可在多线程 CPU 和大规模并行 GPU 上实现高效执行。