在本系列第一篇博客《初探CFD：Taichi能用来做计算流体力学吗？》中，我们了解到Taichi作为一门内嵌在Python中的高性能计算语言，不仅可以作为图形学，渲染器的开发工具，也同样非常适用于需要对2维，3维数组进行大量运算的数值计算，包括计算流体力学。

　　那么用Taichi来做数值计算性能如何？在本篇博客中，我们将考察数值计算中几个常见的运算模式，并分别比较Taichi和它们的计算性能以及易用性。

　　01入门：求和运算中的性能比较实验

　　所谓求和运算就是给定一个数组，求出所有元素之和，是数值计算中很常用的一个运算模式。本节我们将通过一个求和运算的实例，比较Taichi与CUB、Thrust、CuPy和Numba的计算性能。

　　CUB和Thrust均由Nvidia官方主导开发，提供了常用的并行计算函数。CUB和Thrust虽然在用法上有相似之处，但CUB专门为Nvidia GPU开发且更接近CUDA的底层指令，而Thrust则为了兼顾多种平台而建立在比较高的抽象层级，因此CUB的性能往往要高于Thrust。

　　Python用户常见的GPU加速解决方案有CuPy和Numba。其中CuPy提供了和Numpy非常类似的接口，用户可以像调用Numpy一样调用CuPy。同时，CuPy可以选择使用加速接口，这里我们选择了前述性能最佳的CUB作为加速引擎。Numba的规约求和实现则采用了其官方文档中的实现方式。

　　⬇Taichi的求和实现非常简单:

　　ti.kernel

　　def reduce_sum_kernel():

　　sum=0.0

　　for i in f:

　　sum+=f<i>

　　我们将上述几种实现在不同的数据规模下进行性能评估，得到了如下的性能图表，柱形越高表示越接近硬件性能上限。

　　测试硬件i9-11900K+RTX 3080，Ubuntu 20.04，测试的度量为实际达到的内存带宽，越高越好。

　　可以看到在数据规模比较小的时候，各个实现的计算速度都不理想。一方面因为小数组喂不饱GPU，另一方面因为调用本身存在着不同程度的固定开销。而随着数据规模增加，固定开销所占的比例开始下降，所有实现的性能开始大幅上升。其中，CUB和CuPy的CUB后端版本都取得了相当不错的成绩（>90%硬件峰值带宽），而Taichi也和这两种高度优化的实现版本不相上下，并在所有的数据规模上远远超过了Thrust版本的性能。

　　事实上，为了达到高性能，Taichi编译器在底层做了很多工作。在前述的求和代码中，对sum的累加事实上是一个原子操作。原子操作无法并行，运算效率不高。向量求和最常见的并行计算优化手段是规约(Reduction)，这也是入门并行计算的基本功之一。在Benchmark中，我们发现「利用编译器自动实现规约优化的Taichi，达到了与手工实现的CUB类似的性能，远远超过Numba」。

　　在简单的运算模式中，Taichi用非常直观的代码实现了和高度优化的CUDA/CUB相近的计算性能，并大幅超越了Thrust、Numba。而在实际的数值计算任务中，运算公式往往更为复杂，Taichi又有怎么样的表现呢？下面我们来看一下流体计算中的进阶实例。

　　02进阶：速度场运算中的性能比较实验

　　在复杂计算任务中，除了计算性能，代码编写的难易程度直接决定了用户的工作效率。我们以在Marker-And-Cell（MAC）法中的速度场计算为例，来比较Taichi在性能、易用性上的表现。MAC法需要将两个二维速度场沿着时间轴进行推演，计算流体速度场在对流和粘性力的作用下如何变化。这个运算用数学公式来表示是这样的：

　　动图封面

　　这个公式超出了CUB库的API范畴，没办法高效地编写。Thrust作为模板库或许可以编写，但模板编程的debug充满困难。CuPy对于这种复杂的运算则只能通过编写CUDA代码来实现。在这种没有现成的运算库的情况下，我们需要真正的高性能编程语言来实现这个公式。因此，接下来的比较限定在Numba，CUDA和Taichi之间。

　　首先，Taichi和Numba都是嵌入在Python中的计算语言，可以简单明了地用Python语法构造算法。然而，Python的写法在Numba中仅能调用CPU算力。如果想要使用GPU，则要求用户了解CUDA编程模型，自己做好对线程的调度。相比之下，Taichi只需要在ti.init中指定CUDA后端就可以将代码运行在GPU上，不需要任何跟CUDA相关的前置经验。

　　Numba和Taichi实现的对比：Taichi只需要一份代码就可以在CPU和GPU上执行，Numba则需要分别对CPU和GPU编写计算函数

　　至于原生的CUDA实现，笔者半小时写完核心函数，但是折腾了差不多2小时才把数值完全对齐。CUDA到Python还需要准备编译、接口封装、内存分配释放、CPU-GPU数据拷贝等等额外工具代码。虽然我们可以用各种库来简化这个流程，但编程上仍是个很大的负担。

　　那么，在提供了这么多的编程上的便利之后，Taichi的性能如何呢？我们将每个实现连续运行1000次，取每次调用的平均耗时作为度量，得到如下图中的对比结果，柱形越低表示用时越短，性能越好。

　　测试硬件i9-11900K+RTX 3080，Ubuntu 20.04，测试的度量为每次调用耗时，越低越好。计时采用“墙钟”(Wall clock)外部计时，因此包含了函数启动、设备同步等开销。具体细节请参阅我们的代码库。

　　可以看到三个GPU实现的性能均大幅超过了Numba CPU实现的性能。而在边长超过2048的速度场计算中，Taichi的耗时不到Numba CUDA的1/3，甚至比原生的CUDA还要快。这是因为Taichi编译器会在自动使用一些硬件特性来加速访存，而CUDA也可以通过各种编程优化手段来达到和Taichi一致的性能，但需要对底层硬件有深入了解才能实现。

　　03小结

　　我们测试了Taichi在向量加和和流场计算这两个个运算上相对于其他常见实现的性能。通过比较我们得到如下的初步结论：

　　在简单的求和运算中，Taichi取得了和高度优化的主流GPU库实现几乎相同的运算性能

　　在复杂的流场计算中，Taichi取得了Numba(CUDA)约3-4倍的性能，同时略领先于同等语义的手写CUDA的版本

　　在代码编写难度上，Taichi相比Numba（CUDA）以及CUDA大幅度简化，且实现了多种后端的无缝转换，无需用户针对后端修改代码

　　读到这里，相信读者们对Taichi的运算性能和代码编写易用性都有了一定的认识，开始迫不及待地要用Taichi来实现一些自己的数值计算程序了。

　　在本系列的下一篇博客中，我们就要一起来看看，如何用taichi编程语言在99行代码之内来编写一个流体仿真程序。taichi https://taichi-lang.cn/