为什么用 Taichi 做 CFD?只看这一篇就够了~

:books: 简介

合集

Taichi 是一门开源的,嵌入在 Python 中的高性能并行计算语言。相较于一些传统的数值计算用语言(Fortran / C++),Taichi 提供了如下的优势:

  • 简单易用:Taichi 的语法和 Python 极为类似,最大限度降低了并行程序学习和编写的门槛

  • 数据结构:Taichi 提供了专门为数值计算设计的 ti.field ,实现了计算代码和数据结构的解耦,用户可以在不修改计算逻辑的情况下灵活调试数据的行列主序

同时,相较于数值计算用户熟悉的 Matlab / Numpy 等工具,Taichi 则在性能上提供了大幅强化:

  • 超高性能:Taichi 使用独自的 JIT 即时编译技术将用户代码编译成原生的 CPU 或 GPU 机器码,相较于解释型语言的性能提升可达数个量级

  • 可移植性:Taichi 可运行于不同的后端(CPU,CUDA,Metal,OpenGL 等等)且无需用户修改原代码

Taichi 的这些特性使得其在数值计算算法实现中体现出其他语言不可比拟的优势,有越来越多的学者/研究小组将 Taichi 作为算法实现的工具语言:

不仅如此,由于 Taichi 易于安装,几乎不存在附加依赖,使得用户复现他人代码变得前所未有的轻松,因此我们相信 Taichi 将极大地促进数值计算用户间的学术科研交流。

:wrench: 安装

Taichi 作为一个 Python 包的形式出现,用户安装 Taichi 和安装 Python 包一样容易,完全省去了折腾软件依赖的麻烦。首先确认你已经有 pip 并更新到了最新版本:

pip3 install pip --upgrade

假设你的 Python 版本是 3.0 以上,你仅需要这一条指令来安装 Taichi:

pip3 install -U taichi

然后你就可以尝试运行 Taichi 程序了!我们提供了丰富的 example 示例代码

ti example

你可以在命令行中选择自己感兴趣的 example,输入编号即可运行。

我们近期也添加了更加便于交互的 gallery 功能,便于用户用 GUI 界面来浏览选择自己喜欢的示例

ti gallery

:star: 用 Taichi 编写数值计算代码的优势

Taichi 创始的初衷,就是解决数值计算用户在 编程效率 和 运行性能 间的冲突;为了让用户可以用更少的代码,写出更快更强性能的数值计算程序,Taichi 提供了如下的特性:

易用的高性能并行计算

传统的并行计算语言要求用户对于计算机硬件有比较多的了解,在代码编写过程中甚至需要手动申请,释放内存(CUDA),手动分配调度线程块(Numba / CUDA)等。在 Taichi 中,并行计算代码的编写几乎和串行代码一样直白容易,比如用有限差分计算一个 1 维变量的 1 阶和 2 阶导数在 Taichi 中可以这样书写:

f = ti.field(dtype=ti.f64, shape=(128,))
dfdx = ti.field(dtype=ti.f64, shape=(128,))
d2fdx = ti.field(dtype=ti.f64, shape=(128,))
dx = 1e-3

@ti.kernel
def diff():
    for i in range(1,128):
        dfdx[i] = (f[i] - f[i-1]) / dx
    for i in range(1,127):
        d2fdx[i] = (f[i-1] - 2*f[i] + f[i+1]) / dx**2

不仅如此,在 Taichi 中,用户只需要修改 Taichi 初始化时的一条指令,就可以让自己的程序运行在多核 CPU(如 x64, Arm)和 GPU(如 CUDA) 等多种后端上,完全无需顾虑底层硬件的变化:

ti.init(arch=ti.cpu)  # Run on CPU backend
ti.init(arch=ti.cuda) # Run on CUDA backend

这一点在很多的现有工具中都是几乎不可能的。例如在 Python 高性能计算用户中比较流行的 Numba 中,需要用户针对 CPU 和 CUDA 后端编写不同的代码,其中 CUDA 后端的代码编写和 CUDA 的原生接口极为类似,对于数值计算用户来说学习负担较大,代码维护也很不方便。

关于 Taichi 和 Numba / CUDA 的易用性及性能 benchmark 对比,我们有发布一篇技术博客以及一个完整的仓库可供参考。

灵活的数据排布

在数值计算任务中,多维数组的内存排布对于计算性能是至关重要的,而 Taichi 中基础的数据容器 ti.field 是专门为数值计算用户设计的数据结构,其创建和访问很类似于一般语言中的多维数组,或者是 Numpy 中的 Ndarray:

f = ti.field(dtype=ti.f64, shape=(128,128)) # Creates a 2D 128x128 field 

@ti.kernel
def init():
    for i,j in f:
        f[i,j] = 0.0.   # Access field elements

在此基础之上,ti.field 将传统语言中非常头疼的数据排布问题彻底简化:

a = ti.field(dtype=ti.i32, shape=(n, m), order='ij') # explicit row major
b = ti.field(dtype=ti.i32, shape=(n, m), order='ji') # explicit column major

@ti.kernel
def access():
    for i,j in a:
        a[i,j] += 1.0
    for i,j in b:
        b[i,j] += 1.0

注意到 ab 的访问形式是完全相同的,用户仅需要在声明时指定 order 即可实现不同排布。

相比之下,在传统的 C/C++ 类型语言中,你需要显式地调整数据的排布和访问顺序:

int a[3][2];  // row-major
int b[2][3];  // column-major

for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 2; j++) {
        do_something(a[i][j]);
        do_something(b[j][i]);
    }
}

这个调整的难度在 3 维甚至以上的维度时会变得极其繁复且难以检查。关于 Taichi 中的数据排布,可以参考官方文档

:cat: 用 Taichi 编写的流体仿真示例

在 taichi-fluid 中收录了现在我们掌握的来自社区的 CFD 方向的项目,以下列举一些:

:chart_with_upwards_trend: Benchmark

我们在 taichi_benchmark 中给出了大量 Taichi 与其他常见语言实现的性能对比数据,并且我们公开了完整的测试代码,可供用户在自己的硬件环境中测试。(我们的测试硬件平台一般为 i9-11900K / RTX 3080)

我们测试的项目包括了基本运算模式:比如规约求和 (reduce_sum),矩阵元素相乘求和运算 (SAXPY)等:

同时也有实际的数值计算案例,比如 N-body 问题求解,MAC 法中的显式速度场更新,可微不可压缩流体仿真等:

在性能对比中 Taichi 往往可以用极其简洁的代码实现接近甚至小幅超越同等语义 CUDA 代码的计算性能,并在一些算例中大幅超过 Numpy / Numba 等工具的性能。

:ocean: 数值计算兴趣小组 (Numerical Computation SIG) 简介

为了促进数值计算用户的互相学习和交流,进一步降低用户解答疑惑时的负担与门槛,我们设立了 Taichi 数值计算 SIG (Special Interest Group),群组成员可以通过微信群组进行技术讨论,并在每周五组织高质量的数值计算线上分享活动,广泛讨论数值计算中大家关心的问题。目前为止的活动中我们的议题包括:

  • 数值计算数学基础

  • 流体数值计算基础(有限体积法 / SIMPLE 算法)

  • 高性能并行计算基础

  • PBD 软体仿真 Taichi 实现实战

活动的录像在 B 站进行了公开:

:wink: 了解更多

Taichi CFD 合集:GitHub - houkensjtu/taichi-fluid: A collection of CFD related resources for Taichi developers.

Taichi 语言官网:https://taichi-lang.cn/

Taichi 语言开源代码:GitHub - taichi-dev/taichi: Productive & portable high-performance programming in Python.

Taichi 数值计算 SIG 小组:关于数值计算兴趣小组,内含分享排期(板块必读~)

太极图形 b 站数值计算合集:如何用 Taichi 实现数值计算?(持续更新,可收藏后反复观看)_哔哩哔哩_bilibili

2 个赞

太极用sparse做CG与scipy用sparse的加速效果在大网格数下加速效果怎么样?现在CG这个函数支持做多大网格的加速了呢?

现在 Taichi 里有 MatrixFreeCG 和 SparseCG 两种迭代求解器,前者是 Taichi 内的无矩阵求解方式,后者是用 Sparse 方式存储矩阵后进行求解方式。
规模较小的情况下是 Scipy 约等于 SparseCG,略快于 MatrixFreeCG,规模上去了以后则是 MatrixFreeCG 优于 Scipy 和 SparseCG。
关于规模的上限问题,MatrixFreeCG 由于不需要存储矩阵所以可做的规模最大,SparseCG 和 Scipy 都需要存储矩阵,所以规模上限比较低一些,具体情况都要根据硬件和代码来分析。