开源Weld：用Rust让numpy、scikit和pandas加速100倍

在利用Python与R的数据科学堆栈工作了几个礼拜之后，我开始问自己：有没有可能存在一种通用的中间表示形式（类似于CUDA），能够在多种语言之间同时起效？现在我必须得在不同语言中重新实现并优化已经存在的方法，难道就不能更高效一点？除此之外，我希望通过通用的运行对程序进行整体优化，而不是像现在这样只能针对函数进行优化。经过几天的研究和测试，我发现了 Weld。令我意外的是，Weld的缔造者之一正是Matei Zaharia，也就是Spark的发明者。 正因为如此，我联系并采访了Weld项目的主要贡献者Shoumik Palkar。Shoumik目前是斯坦福大学计算机科学系的博士生，在Matei Zaharia的推荐下加入了Weld项目。 Weld还远没有达到生产就绪级别，但其发展愿景却是一片光明。如果大家对数据科学的未来以及Rust抱有兴趣，相信这次访谈会给你带来不少启示。

开发Weld项目的动机是什么，它又能够解决哪些问题？

项目的开发动机，在于为依赖现有高级API（例如NumPy与Pandas）的应用程序提供裸机级别的性能表现。其解决的主要问题，是实现跨函数与跨库层面的优化，而这正是目前其他库所无法实现的。具体来讲，当下有不少常用库以各个函数为基础提供算法的最新实现（例如以C语言在Pandas中实现的快速连接算法，或者以NumPy语言实现的快速矩阵乘法），但市面上却还不存在能够跨越这些函数实现优化的可用工具（例如在执行矩阵乘法后进行聚合时，防止不必要的内存扫描）。Weld希望提供一个可通用的运行时，使各个库能够在通用IR当中进行计算表达；接下来，开发人员还可以利用编译器优化程序对IR进行优化，而后通过循环融合、矢量化等优化手段将JIT与并行本机代码加以对接。Weld的IR具备本机并行特性，因此开发人员可以随时对其中表示的程序进行常规并行化处理。

我们还建立了一个名为拆分注释（split annotations）的新项目，该项目将与Weld相集成，旨在降低对现有库进行优化的实施门槛。

对Numpy、Pandas以及Scikit进行优化困难吗？速度方面能有多大提升？

Weld跨越各个函数对这些库进行优化，而对库的整体优化也能让单一函数的调用速度更快。实际上，不少此类数据库都已经立足函数层面进行了高度优化，只是由于其没有充分利用并行性或者内存结构，所以性能表现仍然低于现代硬件的极限水平。举例来说，目前已经有不少Numpy ndarray函数可以在C语言中实现，但调用每一项函数都需要对每一条输入进行完整扫描。如果这些数组不适用于CPU缓存，那么大多数执行时间就都被耗费在了从主内存中加载数据上，计算的执行反而成了次要耗时因素。Weld能够查看各个函数调用并执行优化，例如循环融合，从而将数据保留在CPU缓存或者寄存器当中。凭借着更好的扩展能力，这些类型的优化能够在多核心系统上将性能提升一个数量级甚至更高。

与Weld原型方案进行集成之后，Spark（左上）、NumPy（右上）以及TensorFlow（左下）的性能相较于原始框架提高了30倍，且无需对用户的应用程序代码做出任何调整。此外，Pandas与NumPy之间的跨库优化（右下）则可将性能提升达两个数量级。

Baloo是什么？

Baloo是一个库，负责利用Weld实现Pandas API中的一个子集。Baloo由阿姆斯特丹CWI研究所硕士生Radu Jica开发完成，其目标是在Pandas当中完成之前提到的各项优化，从而提高其单线程性能、减少内存占用量并引入并行性。

Weld/Baloo是否支持利用核外计算（例如Dask）处理不适合驻留在内存中的数据？

Weld与Baloo目前都还不支持核外计算，但我们欢迎更多开源贡献者能够实现这一功能！

你们为什么选择利用Rust与LLVM来实现Weld？Rust一直是你们的第一选择吗？

我们之所以选择Rust，是因为：

• Rust的运行时很小（基本上只负责对数组进行边界检查），而且易于嵌入到其他语言（例如Java与Python）当中。
• Rust包含函数式编程范式，例如模式匹配，能够使代码编写（例如模式匹配编译器优化）更轻松。
• Rust拥有强大的技术社区与高质量的工具包（在Rust中被称为“crates”），这使我们能够更容易地进行系统开发。

在另一方面，之所以选择LLVM，是因为它拥有一套用途广泛且广受支持的开源编译器框架。我们能够直接生成LLVM以替代C/C++，因此不必依赖于C编译器，从而缩短编译时间（过程中不需要解析C/C++代码）。

Rust并不是Weld采用的第一款实现语言；最初，我们选择的是Scala，当时主要考虑到它拥有代数数据类型以及强大的模式匹配机制。这能够极大简化编译器核心部分优化器的编写过程。我们最初的优化器基于Catalyst设计，该设计器属于Spark SQL的可扩展优化器。但最终之所以放弃了Scala，是因为我们发现很难将基于JVM的语言嵌入至其他运行时与语言当中。

如果Weld主要针对CPU与GPU，那么它与RAPIDS这类专门面向GPU的Python数据科学实现库有什么不同？

Weld项目与RAPIDS等其它系统的最大区别，在于它专注于跨越多种由JIT编译代码编写的内核实现应用程序优化，而非针对每一项独立函数提供优化效果。举例来说，Weld的GPU后端能够为端到端应用程序JIT编译出经过优化的单一CUDA内核，而不像其它系统那样直接调用现有独立内核。此外，Weld的IR意欲设计为硬件无关的，使其既能够面向GPU，也可以面向CPU或者矢量加速器等其它定制化硬件。当然，Weld也与其它系统有着不少交集，同时也受到了RAPIDS等其他同类项目的影响与启发。Bohrium（一种懒惰评估式的NumPy）以及Numba（一套可实现数字代码JIT编译的Python库）就与Weld有着相同的高级目标，而Spark SQL等优化器系统则直接影响着Weld的优化器设计思路。

除了数据科学库优化之外，Weld项目是否还有其他应用？

Weld的IR最令人兴奋的优势之一，在于其自身就支持数据并行性。这意味着在Weld IR当中表达的循环将始终可以安全实现并行化。这使Weld成为一种对新型硬件极具吸引力的IR解决方案。举例来说，NEC公司的合作者们演示了如何利用Weld在自定义高内存带宽矢量加速器上运行Python工作负载——过程非常简单，只需要向现有Weld IR添加新的后端即可。此外，IR还可用于在数据库内实现物理执行层，我们也在计划添加一系列新的功能，使开发人员能够将Python的部分子集编译为Weld代码。

这些库是否已经准备好应用于实际项目？如果还没有，你认为大概什么时候能够做好准备？

我们测试过的大部分库示例与基准都提取自真实工作负载。因此，如果用户愿意在自己的应用程序当中试用我们发布的当前版本，提供反馈甚至编写开源补丁，我们将不胜感激。换句话说，我们并不指望所有功能都能在现实应用程序当中立即使用。在接下来的几个月内，我们将发布一系列发行版，专门关注Python库的可用性与健壮性。我们的目标是努力提升这些库的质量，将其逐步引入现实项目当中，并在无法支持的情况下允许用户顺利回退至非Weld库版本。

正如我在回答第一个问题时所提到，简化整个流程的途径之一，在于拆分注释这个重要的相关项目。拆分注释是一套系统，允许用户对现有代码进行注释以定义其具体拆分、传递与并行化方法。根据我们的观察，该系统能够发现Weld有望带来最佳效果的优化方向（例如在函数调用之间将大量数据保留在CPU缓存中，而非扫描完整数据集）；更重要的是，由于其重用现有库代码，而非依赖于编译器IR，因此其集成难度要比Weld低得多。正因为如此，拆分注释能够使作为优化目标的库更易于维护与调试，从而提高健壮性水平。当无法支持Weld时，这些库也可退回至拆分注释，让我们根据用户提出的反馈以增量方式添加Weld支持，同时继续采取可行的优化操作。

原文链接：

https://notamonadtutorial.com/weld-accelerating-numpy-scikit-and-pandas-as-much-as-100x-with-rust-and-llvm-12ec1c630a1