当前位置: 首页 > 科技观察

使用Rust将numpy、scikit和pandas加速100倍!开源焊接技术揭秘

时间:2023-03-15 08:54:47 科技观察

在使用Python和R数据科学堆栈几周后,我开始问自己:是否有可能有一个可以在多种语言中使用的通用中间表示(类似于CUDA)?同时生效?现在我要用不同的语言重新实现和优化现有的方法,不是更高效吗?另外,我希望通过通用操作来对程序进行整体优化,而不是像现在这样只针对功能进行优化。经过几天的研究和测试,我发现了Weld。令我惊讶的是,Weld的创建者之一正是Spark的发明者MateiZaharia。为此,我联系并采访了Weld项目的主要贡献者ShoumikPalkar。Shoumik目前是斯坦福大学计算机科学系的博士生,在MateiZaharia的推荐下加入了Weld项目。Weld还远未准备好投入生产,但其开发前景一片光明。如果你对数据科学和Rust的未来感兴趣,相信这次访谈会给你带来很多启发。开发Weld项目的动机是什么,它解决了什么问题?该项目背后的动机是为依赖现有高级API(例如NumPy和Pandas)的应用程序提供裸机性能。它主要解决的问题是实现跨函数、跨库的优化,这是目前其他库做不到的。具体来说,目前有很多常用的库提供了基于个别函数的算法的最新实现(比如Pandas中用C语言实现的快速连接算法,或者NumPy语言实现的快速矩阵乘法),但仍然有没有可用的工具来优化这些函数(例如在执行矩阵乘法和聚合时防止不必要的内存扫描)。Weld希望提供一个通用的runtime,让各个库都可以在通用的IR中进行计算和表达;接下来,开发者还可以使用编译器优化程序对IR进行优化,再通过循环融合、向量化等优化方式对IR进行优化。JIT与并行本机代码接口。Weld的IR本身是可并行化的,因此开发人员可以随时例行并行化其中表示的程序。我们还构建了一个名为splitannotations的新项目,它将与Weld集成,旨在降低对现有库实施优化的进入门槛。Numpy、Pandas和Scikit优化难吗?速度能提高多少?Weld跨单个函数优化了这些库,库的整体优化也使得单个函数调用更快。事实上,很多这样的数据库已经基于功能层面进行了高度优化,但由于没有充分利用并行性或内存结构,性能仍然低于现代硬件的极限水平。例如,已经有许多可以用C实现的Numpyndarray函数,但调用每个函数都需要对每个输入进行完整扫描。如果这些数组放不下CPU缓存,那么大部分执行时间都花在了从主存加载数据上,计算的执行就成了次要的耗时因素。Weld可以查看单个函数调用并执行优化(例如循环融合)以将数据保存在CPU缓存或寄存器中。通过更好的可扩展性,这些类型的优化可以在多核系统上将性能提高一个数量级或更多。在与Weld原型解决方案集成后,Spark(左上)、NumPy(右上)和TensorFlow(左下)的执行速度比原始框架快30倍,而无需对用户的应用程序代码进行任何调整。此外,Pandas和NumPy之间的跨库优化(右下)将性能提高了两个数量级。巴鲁是什么?Baloo是一个库,负责使用Weld实现PandasAPI的一个子集。Baloo由阿姆斯特丹CWI研究所的硕士生RaduJica开发。它的目标是在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设计的,属于SparkSQL的可扩展优化器。但最终放弃了Scala,因为我们发现很难将基于JVM的语言嵌入到其他运行时和语言中。如果Weld主要用于CPU和GPU,那么它与特定于GPU的Python数据科学实现库(如RAPIDS)有何不同?ProjectWeld与RAPIDS等其他系统的最大区别在于,它专注于优化以JIT编译代码编写的跨多个内核的应用程序,而不是为每个单独的功能提供优化。例如,Weld的GPU后端可以为端到端应用程序JIT编译优化的单个CUDA内核,而不是像其他系统那样直接调用现有的独立内核。此外,Weld的IR旨在独立于硬件,因此它可以同时针对GPU和其他定制硬件,例如CPU或矢量加速器。当然,Weld与其他系统也有很多交集,也受到了RAPIDS等其他类似项目的影响和启发。Bohrium(一个惰性求值的NumPy)和Numba(一组支持数字代码JIT编译的Python库)与Weld具有相同的高级目标,而SparkSQL等优化器系统直接影响Weld的优化器设计思路。Weld项目除了数据科学库优化还有其他应用吗?Weld的IR最令人兴奋的优势之一是其对数据并行性的原生支持。这意味着以WeldIR表示的循环始终可以安全并行化。这使得Weld成为适用于新型硬件的极具吸引力的IR解决方案。例如,NEC的合作者演示了如何使用Weld在定制的高内存带宽矢量加速器上运行Python工作负载——只需向现有的WeldIR添加一个新的后端。此外,IR还可用于在数据库内实现物理执行层,我们还计划添加一系列新功能,允许开发人员将Python的子集编译为Weld代码。这些库准备好用于实际项目了吗?如果没有,您认为什么时候可以准备好?我们测试的大多数库示例和基准都是从实际工作负载中提取的。因此,如果用户想在他们自己的应用程序中试用我们发布的当前版本,提供反馈甚至编写开源补丁,我们将不胜感激。换句话说,我们不期望所有功能都可以立即用于现实世界的应用程序。在接下来的几个月里,我们将发布一系列专门关注Python库的可用性和健壮性的版本。我们的目标是努力提高这些库的质量,逐步将它们引入到实际项目中,让用户在无法支持的情况下顺利回退到非Weld库版本。正如我在第一个问题的回答中提到的,简化整个过程的方法之一是将注释分解为重要的相关项。拆分注解是一个允许用户对现有代码进行注解以定义其特定的拆分、路由和并行化方法的系统。根据我们的观察,系统能够发现Weld有望产生最佳结果的优化方向(例如在函数调用之间将大量数据保存在CPU缓存中,而不是扫描完整数据集);更重要的是,由于其Reusesexistinglibrarycode而不是依赖编译器IR,因此集成难度远低于Weld。正因为如此,拆分注释可以通过使优化目标库更易于维护和调试来提高稳健性水平。当无法支持Weld时,这些库还可以回退到拆分注释,从而使我们能够根据用户的反馈逐步添加Weld支持,同时继续尽可能进行优化。