99行代码的《冰雪奇缘》

•

——闲聊物理引擎，可微编程，SIGGRAPH生产力难题，与Taichi编程语言

目录

  1. 序
  2. Material Point Method（物质点法）
  3. Moving Least Squares Material Point Method（移动最小二乘物质点法）
  4. Differentiable MLS-MPM (ChainQueen可微物理引擎)
  5. “硬核”的计算机图形学
  6. Taichi编程语言
  7. Differentiable Programming (可微编程) 与 Taichi
  8. import taichi as ti：假装成Python的Taichi语言
  9. 总结

序

终于下决心一口气把四颗智齿拔了，这两天伴随着牙疼只能吃土豆泥和蒸蛋，写编译器之类的体力活是干不了了。登录知乎发现自己居然曾今开了个专栏，无奈一直沉迷于写代码，一篇文章都没写，实在有些惭愧。

最近正好看了《冰雪奇缘2》，就顺便聊聊动画电影里面连续介质，比如雪的模拟吧。实际上，由于最近图形技术的发展，只需要99行代码就可以写一个简单的连续介质模拟器，模拟三种相互作用的不同材料（水，果冻，雪），并且在你的笔记本上就能运行：

14万个水，果冻，和雪“粒子”https://www.zhihu.com/video/1195682534812659712

代码：

import taichi as ti
quality = 1 # Use a larger value for higher-res simulations
n_particles, n_grid = 9000 * quality ** 2, 128 * quality
dx, inv_dx = 1 / n_grid, float(n_grid)
dt = 1e-4 / quality
p_vol, p_rho = (dx * 0.5)**2, 1
p_mass = p_vol * p_rho
E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters

x = ti.Vector(2, dt=ti.f32, shape=n_particles) # position
v = ti.Vector(2, dt=ti.f32, shape=n_particles) # velocity
C = ti.Matrix(2, 2, dt=ti.f32, shape=n_particles) # affine velocity field
F = ti.Matrix(2, 2, dt=ti.f32, shape=n_particles) # deformation gradient
material = ti.var(dt=ti.i32, shape=n_particles) # material id
Jp = ti.var(dt=ti.f32, shape=n_particles) # plastic deformation
grid_v = ti.Vector(2, dt=ti.f32, shape=(n_grid, n_grid)) # grid node momemtum/velocity
grid_m = ti.var(dt=ti.f32, shape=(n_grid, n_grid)) # grid node mass
ti.cfg.arch = ti.cuda # Try to run on GPU

@ti.kernel
def substep():
  for i, j in ti.ndrange(n_grid, n_grid):
    grid_v[i, j] = [0, 0]
    grid_m[i, j] = 0
  for p in range(n_particles): # Particle state update and scatter to grid (P2G)
    base = (x[p] * inv_dx - 0.5).cast(int)
    fx = x[p] * inv_dx - base.cast(float)
    # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
    w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1), 0.5 * ti.sqr(fx - 0.5)]
    F[p] = (ti.Matrix.identity(ti.f32, 2) + dt * C[p]) @ F[p] # deformation gradient update
    h = ti.exp(10 * (1.0 - Jp[p])) # Hardening coefficient: snow gets harder when compressed
    if material[p] == 1: # jelly, make it softer
      h = 0.3
    mu, la = mu_0 * h, lambda_0 * h
    if material[p] == 0: # liquid
      mu = 0.0
    U, sig, V = ti.svd(F[p])
    J = 1.0
    for d in ti.static(range(2)):
      new_sig = sig[d, d]
      if material[p] == 2:  # Snow
        new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)  # Plasticity
      Jp[p] *= sig[d, d] / new_sig
      sig[d, d] = new_sig
      J *= new_sig
    if material[p] == 0:  # Reset deformation gradient to avoid numerical instability
      F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)
    elif material[p] == 2:
      F[p] = U @ sig @ V.T() # Reconstruct elastic deformation gradient after plasticity
    stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T() + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)
    stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
    affine = stress + p_mass * C[p]
    for i, j in ti.static(ti.ndrange(3, 3)): # Loop over 3x3 grid node neighborhood
      offset = ti.Vector([i, j])
      dpos = (offset.cast(float) - fx) * dx
      weight = w[i][0] * w[j][1]
      grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
      grid_m[base + offset] += weight * p_mass
  for i, j in ti.ndrange(n_grid, n_grid):
    if grid_m[i, j] > 0: # No need for epsilon here
      grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
      grid_v[i, j][1] -= dt * 50 # gravity
      if i < 3 and grid_v[i, j][0] < 0:          grid_v[i, j][0] = 0 # Boundary conditions
      if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
      if j < 3 and grid_v[i, j][1] < 0:          grid_v[i, j][1] = 0
      if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
  for p in range(n_particles): # grid to particle (G2P)
    base = (x[p] * inv_dx - 0.5).cast(int)
    fx = x[p] * inv_dx - base.cast(float)
    w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0), 0.5 * ti.sqr(fx - 0.5)]
    new_v = ti.Vector.zero(ti.f32, 2)
    new_C = ti.Matrix.zero(ti.f32, 2, 2)
    for i, j in ti.static(ti.ndrange(3, 3)): # loop over 3x3 grid node neighborhood
      dpos = ti.Vector([i, j]).cast(float) - fx
      g_v = grid_v[base + ti.Vector([i, j])]
      weight = w[i][0] * w[j][1]
      new_v += weight * g_v
      new_C += 4 * inv_dx * weight * ti.outer_product(g_v, dpos)
    v[p], C[p] = new_v, new_C
    x[p] += dt * v[p] # advection

import random
group_size = n_particles // 3
for i in range(n_particles):
  x[i] = [random.random() * 0.2 + 0.3 + 0.10 * (i // group_size), random.random() * 0.2 + 0.05 + 0.32 * (i // group_size)]
  material[i] = i // group_size # 0: fluid 1: jelly 2: snow
  v[i] = [0, 0]
  F[i] = [[1, 0], [0, 1]]
  Jp[i] = 1

import numpy as np
gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
for frame in range(20000):
  for s in range(int(2e-3 // dt)):
    substep()
  colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
  gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
  gui.show() # Change to gui.show(f'{frame:06d}.png') to write images to disk

这段代码用Python 3运行。运行前要根据操作系统和CUDA版本(如果有CUDA的话)安装taichi：

# CPU 版本 (支持Linux, OS X and Windows)
python3 -m pip install taichi-nightly

# GPU (CUDA 10.0) (只支持Linux)
python3 -m pip install taichi-nightly-cuda-10-0

# GPU (CUDA 10.1) (只支持Linux)
python3 -m pip install taichi-nightly-cuda-10-1

这99行代码虽然很短，其背后的故事却很长。虽然语法看起来是Python，其计算部分却会被一整套编译系统接管，最后输出可执行的x86_64或者PTX instructions，能够在CPU/GPU上高效运行。要只用99行代码实现一个这样一个模拟器，这其中凝聚了几代researcher的努力。我很幸运能够参与到这个接力赛中，并跑完最后一棒。接下来，我会提到其中用到的一系列技术，分享这里面横跨20多年的故事。

Material Point Method（物质点法）

Material Point Method (MPM）是一种模拟连续介质的方法，最早被Sulsky等人在1995年发明[Application of a particle-in-cell method to solid mechanics]。由于在计算中同时使用了Lagrangian particle和Eulerian grid，他被归类到Hybrid Lagrangian-Eulerian Simulation Method。这类方法最早可以追溯到Particle-in-Cell (PIC, The particle-in-cell method for numerical solution of problems in fluid dynamics, 1963) 和Fluid Implicit Particles (FLIP, A method for adaptively zoned, particle-in-cell calculations of fluid flows in two dimensions，1986)。熟悉有限元的同学可以把MPM里面的grid nodes对应到FEM里面的DOFs, 把MPM里面的particles对应到FEM里面的quadrature points。和FEM里使用显式网格测策略不同，作为一种Element-Free Galerkin（EFG）法，MPM里面并没有显式的Elements和Lagrangian grid，只有能够随意移动的粒子作为quadrature points。这样的特性使得它非常适合处理大形变，而其背景网格带来的自动碰撞处理、多材料耦合。由于离散化出自weak form，这使得MPM的physical accuracy有了保证。

这些优良属性让它在影视特效领域获得青睐。2013年UCLA数学系教授Joseph Teran第一次把MPM引入到计算机图形学社区[A material point method for snow simulation，SIGGRAPH 2013]。这项技术第一次被用在迪士尼的动画电影《冰雪奇缘》中。迪士尼动画工作室的物理引擎之一Matterhorn就是一个MPM solver。而其他Studios如梦工厂、Weta Digital也有自己的MPM solvers。

Disney还专门做了一个视频介绍自己的连续介质引擎：https://www.youtube.com/watch?v=z4e9i0sCX7c 几年前我在NOI冬令营给中学生介绍物理模拟，还放了这个视频，大家都看的津津有味。

在计算机图形学社区里，研究MPM的代表性实验室有UCLA数学系的Joseph Teran，和宾夕法尼亚大学(UPenn)计算机系的蒋陈凡夫老师。Columbia的Raymond (Yun) Fei，Yonghao Yue， Peter Chen也做了不少MPM的工作。这篇文章里面提到的MPM都是非常基础的内容，关心state of the art的读者可以去看看他们的主页。蒋老师有一个非常全面的MPM paper list: MPM in Computer Graphics 和 MPM教程。

蒋老师等人的MPM教程封面

Moving Least Squares Material Point Method（移动最小二乘物质点法）

前面提到2013年以来，MPM在影视特效中逐渐获得了不小的关注。但是，其性能和复杂程度，则是近两年才有所改进。早期MPM是运行速度是非常慢的：Disney的工程师Alexy Stomackin（Joseph Teran的学生）就提到，《冰雪奇缘》里面Elsa跨过雪地的镜头，在集群上跑了整整一个星期。

如果我没猜错的话，Alexy说的应该是这个场景。

2017年大四暑假，我在宾夕法尼亚大学visit蒋老师。我们设计了一个新算法，Moving Least Squares MPM (MLS-MPM)。MLS-MPM不仅性能比之前的state of the art高两倍，而且代码短了很多，非常容易实现。为了加速我们的实验，让MPM跑的快一点，我花了点时间手写SIMD intrinsics，把它加速了6倍。在这个基础上，MLS-MPM在算法上的改进还能进一步提高两倍的性能。

在费城一个烟雨蒙蒙的晚上，我们几个作者（Yu Fang, Ziheng Ge, Ziyin Qu）集中在蒋老师的实验室，经过一系列讨论，我突发奇想：既然APIC的affine velocity field和MPM里面的deformation gradient update需要的都是velocity gradients，那能不能用moving least squares统一两种离散化？我快速的实现了一下，发现确实可行，但是问题是这个新离散化理论上是否有原来的MPM的精度？那天晚上我们去吃了顿中餐，一路上一直在讨论这个问题。吃饱了饭，晚上我们一直工作到凌晨3点，当时我已经困得不行了，只记得蒋老师突然兴奋的说了一句“weak form推出来了！算法是weak-form consistent的！” 我看了一眼蒋老师草稿纸上的推导，平均每个符号有3个上下标(\alpha, \beta, i, j, k，下划线，双下划线，等等)，似懂非懂，只知道有了这张草稿，算法的物理正确性算是有了保证。后来我花了两整天的时间才弄懂蒋老师那晚上推出来的公式。（蒋老师是中科大少年班毕业，只比我大四五岁，却已经当了一两年教授了。相比之下我那时候还没开始读PhD…）

除了新的离散化，我们还设计了一个新的把MPM和rigid body双向耦合在一起的算法，并且能够实现切割的模拟。最后我们整合这些结果做了一个（我认为）质量挺好的论文，发表在SIGGRAPH 2018上。直观的演示在下面这个视频里。

[SIGGRAPH 2018] 移动最小二乘物质点法与兼容性粒子网格法（双向流体/固体耦合）_哔哩哔哩 (゜-゜)つロ 干杯~-bilibili

文章发表以后，为了进一步证明MLS-MPM的简易性，我又设法用88行C++代码实现了一个self-contained MLS-MPM demo。为了能够让各个平台的玩家都能编译、运行我的代码，我还专门写了个parser来分析"#include"，用来聚合(amalgamate) taichi library，得到一个两万行的单文件taichi.h来支持88行的MPM solver。我还在MIT给了个不太正式的talk专门讲这个single header和跨平台GUI是怎么实现的。很多人吐槽我#include "taichi.h" 以后声称代码只有88行有点cheating的意味，但其实taichi.h里面只是提供了可视化、向量运算之类的基础设施，核心算法还是88行实现的，也非常容易学习。

吐槽归吐槽，真正喜欢graphics的人还是会把这份代码好好研究一遍的。之前很多研究者都“谈MPM色变”，因为他们觉得这个算法实在是臭名昭著地难以实现：既有particle又有grid，每个particle要维护好多物理量（形变梯度、Young's Modulus, Poisson's ratio之类的），早期的实现也没有开源，而且有很多物理参数需要设置。这个88行版本简单易懂，而且跨平台，受到很多想实现MPM的同学欢迎，甚至有热心网友实现了Javascript、Unity、iOS版本。有了这个88行版本，在各路physically based animation算法中，MPM一下子从最难实现的算法变成了最简单的算法之一。现在那个88行版本几乎成了入门MPM的必备参考实现，连我自己在重新实现MLS-MPM的时候也以它作为reference。

再后来，GPU-MPM“三剑客” 高明、王鑫磊、武奎把MLS-MPM高效地在GPU上进行了实现，又把性能提到了一个数量级。他们的工作为我后来的几个项目铺平了道路。最近蒋老师组的王鑫磊和李旻辰等大神又做了一个基于h-multigrid的implicit MPM solver，能够无视材料硬度的限制做到CFL-rate的time stepping，很值得关注: https://arxiv.org/abs/1911.07913。

对高性能MPM感兴趣的同学欢迎check out我们( @张心欣、高明、 @蒋陈凡夫 )在SIGGRAPH 2019的课程(http://taichi.graphics/mpm_course2019/)。为了简单，这里的99行实现没有使用GPU shared memory、particle sorting之类的优化。

我们在SIGGRAPH 2019的MPM教程的封面

Differentiable MLS-MPM (ChainQueen可微物理引擎)

从88行的C++版本，到支持CPU多线程、GPU、多材料模拟的99行的Taichi版本，又过去了快两年的时间。接下来我就分享一下这两年间发生的故事。

在蒋老师那里愉快的暑假结束后，我去了MIT开始读Ph.D.。我又做了几个项目，其中就包含基于MLS-MPM的ChainQueen。这个项目的动机非常单纯：既然MLS-MPM只要88行就能实现，那么求出它的导数应该也不会太复杂。（这里导数指经过若干时间步以后粒子状态关于初始粒子状态或者controller参数的Jacobian）。有了导数，这样我们就能只用梯度下降优化neural network controller了。

有了simulator的梯度，就可以用暴力梯度下降优化controller。这里演示了梯度下降的不同迭代中，模拟结束时软体机器人的位置，目标是尽量向右移动。58次梯度下降就能优化完成，比reinforcement learning快若干数量级。

问题是技术上怎么求导数？当时常见的自动微分工具比如TensorFlow针对deep learning设计，用它做simulation几乎没有利用任何producer-consumer locality，本来能暂存在registers里面的数据只好甩到main memory，导致和mega-kernel的设计比起来，arithematic intensity极低，而且kernel launches太多。在小规模数据上，可能95%的时间都在kernel launching。再加上simulation一般得上千时间步，如果把每一步simulation里面的几百个操作都看成neural network的层，相当于写了个一个一百万层的神经网络。TensorFlow光是编译计算图就要好几分钟时间。JAX/XLA也没有彻底解决这些问题。为了性能，我去手写了一些CUDA。后面会提到一个叫DiffTaichi的项目，解决了物理模拟里面的自动微分问题。

虽然MLS-MPM已经被简化到极致了，手动把它的导数推出来还是挺麻烦的，最后文章补充材料里面有120行公式，实际上就是我求导数的草稿纸。由于这个求导过程太丧心病狂了，需要不断使用chain rule，我决定让这个项目的名字里面带上“chain”，来纪念我在推导数过程中受的罪。于是我把项目叫做ChainKun，谐音”乾坤“，后来合作者佳俊觉得“Kun”太random了，就选了ChainQueen。其实我们还考虑了“ChainKing”，后来觉得“ChainQueen”不但发音更像中文的“乾坤”，而且更有“diversity”。不过后来国内的媒体们似乎没有理解这个名字的意义，把这个颇有意味的名字暴力逐字汉化，成了“链后”。（后面会提到DiffTaichi解决了手动求导的问题。）

公式推对了以后把这一坨用CUDA写出来，然后debug，是一个更加让人头疼的过程。对了，对于这种gradient-based optimization，即使算出来的gradient是错的，有的时候optimization仍然能够表现良好。直观想一下，只要算出来的graident和真正的gradient的dot product大于0，那么gradient descent也能在一定程度里面make progress。这也意味着光看loss curve很难推断gradient是不是有bug，再一次说明了一个不会出错的自动微分系统的重要性。

在给力的姚班小学弟刘剑成，学长Andy和佳俊的大力支持下，整个项目前后只花了30天，最后赶上了ICRA 2019的deadline，并很顺利地被接收了。其实最后正文deadline的时候我们还没有靠谱的3D可视化工具，只好暂时用Houdini里面的renderer勉强渲染了一下。好在视频上传日期在正文deadline后一周，所以我们在视频里面加了一些结果。这短短的一周里，为了能够及时渲染出来，Houdini的CPU渲染是来不及的。

我只好花了一天用CUDA手写了个GPU ray tracing renderer, 专门用来渲染软体机器人，这下基本上结果都能在几分钟之内被渲染出来了。后来这个renderer到了Andy手上，他又接着做了一个工作投到了NeurIPS 2020，那个工作还被MIT主页报道了。我一天内糙快猛写出来的renderer渲染出来的结果就这么登上了MIT主页。一些可视化的结果：

[ICRA 2019] ChainQueen（乾坤）可微软体物理引擎_哔哩哔哩 (゜-゜)つロ 干杯~-bilibili

当然我自己也知道ChainQueen这样的文章是够不上SIGGRAPH的标准的。我还记得有一次我回MSRA给talk，讲到ChainQueen，一个SIGGRAPH大佬问我：“你这工作不就是求了个导么？” 我一下子被问住了，我觉得他说的很对，只好尴尬地回了一句“你说的对啊！但是我也承担了‘求完了导以后导数没用’的风险，用实验验证了几千的timestep的simlation导数还是有用的，所以还算有那么一丁点贡献吧！ 🙂 ”

虽然算不上最得意的工作，由于完成时间合适，这篇论文后来成了我的master thesis。在我24岁生日那天，开始读PhD 13个月后，我把master thesis给前老板签了字，换了个组继续我的PhD生涯。（MIT EECS PhD的前半部分可以拿一个master学位。）

在此期间另外一个值得一提的工作是和Dartmouth College的朱博教授、UWisconsin-Madison的Haixiang Liu和Eftychios Sifakis做的Giga-Voxel topology optimization（1,040,875,347个体素的拓扑优化）。后来这篇文章发到了SIGGRAPH Asia 2018，我们成功证明了通过end-to-end codesign出来的solver，使得一台机器能够发挥100台机器的效果。经过这样的一个项目，我从朱博、Eftychios还有Haixiang身上学到很多。这其中的故事值得有空的时候单独写一篇文章，这里就不多说了。

由于种种原因，我的Ph.D.第一年很难说过得有多愉快，不过还是发了6篇paper。在“一年内发出的ACM Transactions on Graphics的数量”这个指标上，我感觉可能很难再刷新自己的记录了。“既然已经发了足够的SIGGRAPH，不需要再追求数量，现在也没有太大毕业压力，不如用剩下来的时间和科研自由做点真正有意义的事情，比如说提高整个社区的生产力。” 我这么想。

“硬核”的计算机图形学

众所周知，发一篇SIGGRAPH的工作量实在是太大了。大部分SIGGRAPH项目耗时1-2年，在北美有2-3篇SIGGRAPH一作一般就能PhD毕业了。功利地说，最近几年发SIGGRAPH的性价比是比较低的：工作量大，要求高，发出来引用量还比较少。尽管如此，我还是很佩服在这种模式下做出的高质量文章，比如朱教授PhD期间的codimensional流体模拟系列。要么不发，发了就要保证高质量、问心无愧，做出突破性的工作，这也是我的目标。

在这样高标准、严要求下，做一篇SIGGRAPH需要大量的脑力、人力、物力，做出来以后实际deploy到生产环境中需要的工程量就更大了。久而久之，计算机图形学就成了大家口中的“硬核”学科。问题是，计算机图形学的科研是否应该是像大家通常认为的那样“硬核”？

在相对不那么硬核的AI领域，这几年有了deep learning，中学生会一点Python和TensorFlow/PyTorch以后都能发CVPR之类的顶会。很多时候大家从负面看这个事（投稿数目太多、质量下降等），但是这里我从正面来看。Deep learning现在这么流行，好用的基础设施功不可没。如果工具易用到连中学生都能灵活使用，那博士生、工程师岂不是更能好好利用工具实现自己的想法？

回到图形学领域，graphics研究者对计算机底层的了解，让他们能够为了性能写出非常low-level的代码，但是low-level engineering本身就会消耗大量的精力，不管是新手还是老手。这些底层知识反倒成了一种“诅咒”，限制了这门学科自身的发展。这就好像工业革命发生以后，手工作坊的“硬核”工人嘲笑机器大批量生产出来的产品“没有灵魂”。后来发生了什么我们都知道。落后的生产方式终究还是要被淘汰的。

当然这并不意味着掌握low-level knowledge就没有意义。即使有了更好的工具，理解工具内部如何运行也是能让用户更好的使用工具的。但问题是，在SIGGRAPH社区里，造工具的和用工具的，是一帮人。或者说，好用的工具，很大程度上并没有分离出来，大部分还停留在“一次性”的阶段，所有人都要把底层的知识重新学一遍。

比如说，计算机图形学的算法大多用C++实现。一些组里的代码库出现了编译一把要一个多小时的情况。由于大量使用templates (比如2D/3D simulator用template一起写，再对float/double实例化)，新来的同学都表示上手非常困难，很多时候需要debug compilation，因为C++模板的compilation error太难懂了。甚至有同学表示“通过编译就是成功的一半”。而就是同样这一批同学，可能毕业的时候已经成了C++高手，满口“TMP“， ”SFINAE”，代码必须要C++17才能编译，编译的时候compiler疼的嗷嗷叫，甚至出现g++ out of memory等等。新来的同学又要继续要承受modern C++和各种over design的折磨。这种形式的知识继承，是相对比较低效的。

我就曾今接手过一个MPM代码库，有一次在一台“只”有32GB内存的机器上开4个线程编译。由于改了一个底层的header，估计要一个小时才能编译完，我决定先去吃午饭了。去楼下中东餐车了排半天，买了个$7还送饮料的lamb over rice（某种“羊肉盖浇饭”？），吃完回实验室已经快一个小时以后了。我本以为代码应该编译好了可以运行了，结果看到屏幕的那一刻我内心是沮丧的：编译了20分钟的时候，机器out of memory了，只好把memory swap到hard disk，现在OS已经卡的没法操作了。这是一个宝贵的教训：要编译那份代码的时候，如果我只有32GB的内存，那我只敢开3个线程编译，即使CPU有8个cores。

Taichi编程语言[项目主页]

为了解决计算机图形学研究对性能的追求以及生产力低下的问题，我决定重新设计编程语言。这是一个大胆的决定，因为虽然我之前花了很多时间做low-level performance engineering，但是真让我写个编译器我还从没干过。从2019年1月起，我一直在做Taichi programming language (这里非常感谢我现在的两位老板Fredo和Bill给我这样的自由) ，这个结果发表在SIGGRAPH Asia 2019上。

毫无疑问，Taichi编程语言的工作量是非常大的。这是一个全新的系统，项目早期要做的设计决策非常多。这带来一正一反两个结果：

  1. 设计决策可以只在脑子里进行，而不需要盯着屏幕。这意味着这个项目总体还是比较“健康”的：代码写累了就可以闭眼甚至躺床上想想别的地方怎么设计。可以说，Taichi programming language的1/3是在床上、公交车上完成的。这样的时间分配一定程度上避免了Taichi编译器的开发损伤我的颈椎和视力。
  2. 早期设计决策需要随机应变，决策的交流是非常低效的。人与人之间同语言交流的带宽比脑内带宽低好几个数量级，就像network bandwidth和L1-d$ bandwidth的差异一样，于是我只好一路从底层的代码生成(codegen)，到中间表示(IR)和优化(opt)，到前端语法，最后到图标设计(???)，都自己干了。

神奇的是，即使有的时候感觉工作量太大让我“累觉不爱”，睡一觉起来又会回复斗志。表面上看这样是低效甚至“独裁”的，但是最后结果却不错。现在回顾整个开发过程，大部分设计虽然面临很纠结的trade offs，但最后的做出来的决策总体还是较优的。早期编译器的开发本身是一个非常严密的过程，两个人一起写只会引入更多的bugs，我一个人写反而能够降低写出bug的概率，反而比多人协作快。由于我知道我自己想达到什么样的目标，我一个人做决定，决策效率其实非常高。而且自己实现自己的想法，并且设计决策的（短期）正确性，我还是比较开心的。

这个项目本身还算挺enjoyable的，让我学到了不少新东西。一开始做这个项目的时候刚换组，新的组还没有电脑用，我就work from home了一个月。我还记得有一次我去了一趟中国城买菜以后连续在家写代码5天没出门，最后实在受不了没人说话去实验室找同学聊天，学长韬哥都说我身上“长蘑菇”了。

最后事实证明我低估了写编译器的工作量，一个月写出一个我理想中的compiler最后被证明还是有难度的。中间因为第一次写，缺乏设计compiler的经验，我还把IR设计错了，最后只好砍掉重来，才有了现在Hierarchical SSA的结构，使得很多优化成为可能。本以为能赶上一个月后的SIGGRAPH 2019，虽然我已经非常努力了，最后SIGGRAPH还是没赶上，做饭的水平倒是提高了不少。

蒸蛋里鸡蛋和水的最佳比例好像是1:1.75?

错过deadline以后，我只好默默的接受了这个事实，淡定地玩了几天《刺客信条：奥德赛》，用我那块买来跑CUDA的GTX 1080 Ti体验了一下尖端real-time rendering技术。春节回了趟家，稍作休整。虽然嘴上自我安慰“反正也不缺这一篇paper”，在deadline的驱使下身体还是诚实的，回MIT以后我又“码”不停蹄几个月，总算赶上了5月份SIGGRAPH Asia的deadline。

刨去春节回国休假，和最后做benchmark、写paper (这里要多谢Tzu-Mao Li和Luke Anderson在benchmark和writing方面的帮助)，Taichi的语言设计和编译器实现本身大概只花了我4个月的样子。我想这很大程度上得益于我和女友平时有12-13个小时时差，她睡觉的时间我可以专心写代码。

虽然reviewer都吐槽文章有点难懂，但最后还是给出了了非常建设性的意见，和一致通过。

“几百个…….通过”: 那年SIGGRAPH评分系统有6挡: strong rej, reject, weak rej, weak acc, accept, strong acc。这几年规则一直在变，不知道今年有几档…

不得不说SIGGRAPH的reviewer真是我投过的所有的会里面最负责、bar最高的reviewer，在技术上的专业性是非常给力的。不过，他们的建议里让我最印象深刻的，还是其中有一个reviewer鸡蛋里面挑骨头，喷我的video做的太差：

Comments about Video: The video would be a great place to provide an overview and a big picture of the entire framework, and could also help the reader understand details that were not adequately explained in the paper. Unfortunately, this was not done, the video is extremely short, and only shows some examples.

实际上提交的时候的video是我在deadline前三个小时手忙脚乱做的，总时长只有1分钟，只演示了几个编译出来的程序运行的视觉效果。最后关头我都在想，“这video这么烂（和之前做过的simulation paper比起来），要不干脆别交了？”从事后诸葛亮的角度看，可能当时确实不交比较好。

我在知乎上看到很多人说SIGGRAPH论文得花一个星期做video，其实根据我的经验那只是少数，而且我相信大部分作者在deadline之前是不太可能还有一周时间来做video的。Video大多数情况下对于reviewer的作用也只是一个印象分，这个批评我video太差的reviewer最后也给了accept。“SIGGRAPH论文靠fancy的video吸引眼球”其实也是一个误解，首先大部分SIGGRAPH论文其实是没有video的(这里有一个幸存者偏差)，大多数时候video的质量也并不起到决定性作用，除非你的工作目的就是demonstrate视觉效果。

不过video做的好的工作往往论文本身质量也很高：如果对工作本身不喜欢，谁还有动力去花大力气做个video呢？当然，在提交paper的同时就提交一个高质量的video，让reviewer感受到你的诚意，他们自然也会认真review你的工作，提出更多的建议。Review一篇SIGGRAPH submission是极其费时费力的，以至于很多教授劝自己学生把项目多做一段时间再投，给出的委婉的原因都是“这样的文章投了可能会浪费reviewer的时间”。

当然最后我觉得那个一分钟的video实在拿不出手，等paper中了以后我又重新做了个中规中矩的video：

[SIGGRAPH Asia 2019] Taichi编程语言：高性能稀疏视觉计算_哔哩哔哩 (゜-゜)つロ 干杯~-bilibili

只实现一个好用的编程语言，是够不上SIGGRAPH的标准的，还必须得有非常明显的novelty才可以。我们的文章里面真正的贡献是实现algorithm和sparse data structure的解耦(decoupling)，以及对稀疏数据结构的支持。结果还是让人比较满意的，在我们的4个benchmark case上，只用1/10的代码，就能达到4.55x的性能。有了Taichi，我可以83分钟用300行代码写出一个GPU multigrid preconditioned conjugate gradients Poisson solver，这在之前用C++/CUDA是几乎不可能做到的。我把文章的摘要翻译过来放在结尾了，感兴趣的同学可以去围观。希望下次有机会（???）能够分享一下真正在论文里面的内容。

Differentiable Programming (可微编程) 与 DiffTaichi[项目主页]

今年暑假我在Adobe又花了一个月给Taichi重新设计前端语法使得其可以被嵌入到Python中（后面会介绍），并且加入了Differentiable programming的功能，再也不用手动求导数了。Differentiable programming本身也不是什么新东西，有几十年历史了。只是最近由于deep learning火起来了，大家就突然又对这一套求梯度的奇技淫巧燃起了兴趣。由于我自己对simulation很感兴趣，就在想怎么设计一个在simulation的computational pattern下能够高效运行的可微编程语言，主要针对机器学习、视觉、机器人方面的应用。

在Adobe我的manager Qi Sun给了我无微不至的关怀：不但为我安排了一个靠窗的座位写代码，还把暑假要用到的硬件全都提前准备好了，实习期间还批准我回国看了一下家人。他的大力支持使得研究能够非常顺利地开展。这期间还见到了传说中的Li-Yi Wei和Andrew Adams (Halide的主要作者之一)，相谈甚欢。

实际上在Taichi上面实现differentiable programming并不困难，核心的部分就是加一个source code transform用来做reverse-mode automatic differentiation，我专心写了一周基本上就弄完了。当然，从性能和编译器实现复杂性的角度考虑，并不是任意Taichi程序都能被autodiff，已有的simulator需要进行一些很小的改动为自动微分铺路。

由于暑假基本不用参与别的项目了，也终于有了一些空余时间。在加州的住的地方我还买了个Nintendo Switch，和住一起的史亮同学一起通关了《塞尔达传说：荒野之息》。我们的游戏策略是，他主要进行high-level planning，我负责low-level control，作为一个hierarchical control system，我们俩配合严丝合缝而又充满了欢乐。

离开加州最后一天晚上赶deadline打通关了，见到了塞尔达公主…

暑假结束以后我回到MIT，觉得DiffTaichi这个项目要投到SIGGRAPH工作量有点大，就决定偷个懒，投了15天以后就deadline的ICLR 2020。为了证明DiffTaichi的高生产力，我一口气写了8个不同的differentiable physical simulator，合作者Luke写了2个，这样就用10个differentiable physical simulator作为文章里面的案例。一些机器人演示：

[ICLR 2020] DiffTaichi：基于Taichi的可微编程和可微物理引擎_哔哩哔哩 (゜-゜)つロ 干杯~-bilibili

写完paper发现ICLR只能交8~10页，我只好把大量内容移到Appendix，最后8页正文12页Appendix。这期间还和reviewer关于哪些内容应该在正文、哪些应该在Appendix扯皮了一趟，并没有我预期的那样顺利。

几个月过去，现在回想起来这是个正确的决定。一方面当然是因为paper中了（虽然去埃塞俄比亚开会好像有点折腾的样子…），另外一方面是很多用户急切的需要这样的一个工具，我及时把它做出来并且开源他们就能用上，不用自己手动在CUDA里面写导数了。

当然，投ICLR而不是SIGGRAPH，还有个非常、非常重要的考虑，那就是，早点把这个项目结束掉，我就能安心回国拔智齿了…

import taichi as ti（Taichi与Python的无缝衔接）

DiffTaichi完成到现在，我一直在做一些语法、标准库（主要是linear algebra）方面的设计，编译器实现方面的工程工作。MIT内部有3个项目在使用Taichi开发，校外还有两三个组在试用Taichi。我现在每天的工作基本上就是为他们提供支持、开发新的feature方便这些用户。

这意味着我得暂缓自己的新项目，全天待命回答用户的问题，但是我很enjoy这样的生活：和“开会出差找秘书报销”、"为了招intern和做事龟速的HR打交道“、“自行车系统性爆胎维修”、“拔智齿”和”给金主大爷做报告”等等让人头疼的琐事比起来，纯粹地写代码真是人生第一大乐事。

实际上，我开始写编译器以后，每天对于Taichi编译器相关问题的思考，大大丰富了我的精神世界，把人世间的种种噪声都隔绝了。这让我生活质量有了很大的提高，用“快活似神仙”形容也不为过。

虽然一系列论文都发了，要让大家能够方便的把这个工具用起来，还是挺有挑战性的。CI、测试、文档是最基本的。早期的时候Taichi被嵌入在C++17里，用了大量的macro、TMP、lambda之类的技巧，要入门还是有一定门槛的。而且C++编译速度本身就很慢，体验不是很好。后来我实在受不了了，决定重新构建语言的前端。这里有两个选择，一是重新设计一个新的前端语法，二是把语言嵌入到Python里。

俗话说"喜欢是放纵，爱却是克制"，我最终还是成功抑制住了自己重新造一套语法的冲动，写了一套运行在Python AST上面的frontend compiler把Python AST编译到Taichi AST（这是得益于Python作为动态语言的特性，在C++里是很难获得AST的）。之所以这样做，主要是考虑到以下几个优点

  1. 由于Python本身是一门很容易学、普及程度很高的语言，伪装成Python以后的Taichi语言很容易学习；
  2. Python自带包管理系统，这使得作为编程语言的Taichi可以伪装成一个普通的Python package；
  3. 已有的很多Python IDE可以直接拿来当做Taichi的IDE；
  4. 这样设计以后很多常用的Python包，如numpy、matplotlib、PyTorch都可以和Taichi无缝衔接。

当然缺点也不是没有，比如说Taichi“寄人篱下”以后有了一些限制，首先是语法必须parsable by the Python parser，不然根本拿不到AST；其次是有些时候性能会受到Python的限制；另外就是作为编译性、静态类型、lexically scoped的Taichi语言和Python本身有一定差别，可能会给习惯了Python思维方式的新用户一定的误导，使得一些bug难以被发现。

不过这些问题后来都得到了还算不错的解决，用户们也可以很容易的“import taichi as ti”了，就像文章开头的99行代码一样。回过头来看，总体上“Taichi假装成Python”这个决定还是利大于弊的。

另一个比较大的工程就是把原来的source-to-source的backend compiler换成LLVM。在赶SIGGRAPH Asia的时候我用让Taichi直接吐出C++14（其实是各种SSE/AVX intrinsics封装的模板库）和CUDA。和我当时不熟悉的LLVM IR相比，这套我对C++和CUDA至少还是能熟练使用的，并且这两个语言里面的high-level constructs用起来还是很顺手的。当然代价就是由于用了挺多templates，随便编译一个两行的Taichi程序都要好几秒。我很快就受不了了：这几秒的延迟开个小差，context switch一下足以把我脑子里cache住的东西刷没，非常影响开发效率。

在几个deadline都赶完了以后，我下决心换到LLVM。当然这样做的代价是挺高的：我在LLVM backend上加起来已经花了两个月的full time，到现在为止也没有实现以前source-to-source的全部功能和性能。不过如果我一开始就用LLVM，估计当时就赶不上SIGGRAPH Asia deadline了。LLVM的大部分自动生成的文档，说实话对我帮助不大，很多时候还是得靠看源代码或者看别的编译器怎么使用LLVM。比如，Taichi的NVPTX的backend就参考了Halide的实现，因为那坨API除了靠已有的example，实在是难以了解如何使用。LLVM IR作为low-level IR，也没有像C++这样的高级语言的模板系统，这些都得我自己去在compile-time simulate然后forceinline使得LLVM有优化的余地。当然，虽然过程有点酸爽，切换到LLVM总体还是利大于弊的，因为

  1. 11倍的编译速度提升（没有了C++前端的parsing/template resolution之类的overhead）
  2. LLVM对SIMD的良好支持（大部分SIMD intrinsics不用手写了）
  3. 基于LLVM开发的系统的高portability（ship一个包含了statically-compiled LLVM的shared object比要求用户装gcc/clang-7/8容易多了）。

这里面Portability是非常重要的一点。经过一系列麻烦的工程，用户终于可以一键pip install taichi-nightly了。这使得文章开头的程序，几乎在每一台有Python 3的机器上都能很容易的跑起来。有经验的同学都知道，以前的图形学项目，基本上要换一台机器跑都是很头疼的，往往伴随着“灵魂10问”：

  1. 你用啥操作系统？
  2. glibc版本是啥？“ld –version”输出给我看看？
  3. CPU哪一年的？AVX2支持吗？
  4. GPU啥型号？compute capability多少？OpenGL 4.1支持吗？
  5. GLFW，GLEW装了吗？
  6. gcc/clang版本对不？C++14支持吗？
  7. …

说多了都是泪。希望很快这些问题都会简化为一个问题：“Taichi版本是多少？”

总结：

计算机图形学，和其他科学一样，应该把"简单性"作为追求之一。如果MLS-MPM不能被88行代码实现，求它的导数可能也就不能“只”用120个公式完成，可能也就不会有ChainQueen（可微物理引擎），更不会有ChainQueen启发的工作（比如DiffTaichi）。化繁为简本身就是科研的意义之一。

计算机图形学研究比较小众，和其相对较高的门槛是有一定关系的。我觉得，归根结底，"门槛高"还是个生产力工具的问题。人脑能handle的复杂程度是有限的，在没有好用的工具的情况下，要求一个researcher既要懂点连续介质力学，又要懂数值计算，最后还得写高性能并行代码，这是很困难的。

你可能要问，可不可以找三个人来分工解决一个科研项目？很遗憾，要做出一篇传统意义上的好paper并被SIGGRAPH接收，结果往往是end-to-end codesign出来的，最后呈现给reviewer的得是一个浑然一体的设计，这就要求leading author（在国外通常是第一作者或者最后作者）或多或少对每个方面都有所了解。在这方面登峰造极的是U-Wisconsin Madison的Eftychios Sifakis，他的好多项目都是从discretization (FEM, FDM)到numerical algorithm (multigrid, Krylov subspace solver, Schur complement)到高性能solver实现(SIMD, ILP, mixed-precision，virtual memory，heterogeneous computing)每个细节都优化到丧心病狂的程度，这里限于篇幅，没法继续深入了。他对我的影响非常大，希望以后有机会能够写一写我从他身上学到的东西。类似的，如果我不懂numerical solver和Monte Carlo renderer的computational pattern，不懂编译器怎么写，或者不懂点高性能计算、体系结构，很难想象Taichi语言会被设计成什么样、有多大实用性。

计算机图形学本来是非常吸引人的学科：只用最基本的加减乘除，每个人都可以通过编程创造一个虚拟世界。只是硬件的复杂、计算量的巨大、工具的缺失，让太多人对这个美妙的学科望而生畏。而Taichi的使命，恰恰就是让以后的graphics爱好者、研究者们能够从low-level engineering中解放出来，能够专注于建模、数学方面的创新。

我非常感谢韬哥、武奎、Andy、Srinivas、KLoze以及朱教授的学生们（特别是Xingzhe He， Yijun Li等）愿意做第一批Taichi语言用户（i.e. “吃螃蟹的人”）。除了他们之外，Luke在Taichi语言非常早期还不成熟的时候，写了相当复杂的程序：他居然在没有发邮件塞满我的邮箱的情况下几乎独立写了一个体积渲染器。如果我自己用那样不稳定的编程语言，估计早就把键盘砸了。

为了尽可能减少世界上键盘的非正常损坏，目前我的指导方针是“用户就是我大爷”。对于两个月前的他们，用一门像婴儿一样“一天一个样”的编程语言，还有个啰嗦的语言设计者整天缠着问哪里用的不舒服，想必也是一种特殊的体验吧 🙂

好在经过几个月的工程，现在语言和编译器已经渐渐稳定下来，从“婴儿期”进入了“青春期”：虽然不会动不动尿裤子了，经常有点“叛逆行为”还是挺常见的。毕竟我一个人再努力，能够做的事也是微小的。即使远远算不上完善，Taichi好歹也基本实现了用99行代码就能写一个GPU 多材料MPM solver的目标。

逐渐的也有了一些热心网友加入了开发，我的那种”个人英雄主义“的开发模式应该很快就可以告一段落了。希望Taichi在不久的将来（或许明年底？），能够真正解决SIGGRAPH的生产力难题。

不知不觉扯得有点远了。本来只是陪女友看了《冰雪奇缘2》以后想写篇短文分享一下电影里面雪的模拟算法，最后居然写成了我的PhD前两年的回忆录。更没想到的是，文章写了一半，拔牙的伤口居然不疼了，于是我又回去集中注意力写编译器了。而这篇闲聊一拖就是一个月，好在今天得空，在也算赶在2019年写完了，希望对大家有帮助吧。

从1995年被MPM用于固体模拟，2013年Joseph Teran组第一次用它来模拟雪等弹塑性体并在《冰雪奇缘》中被用来做VFX，2018年MLS-MPM极大简化MPM使得其可以用88行C++代码在CPU上运行，到今天有了Taichi编程语言后MPM能够用99行代码就在GPU上实现，或许MPM本身20多年的发展也可以算得上一部《MPM奇缘》了。感谢一路上给力的合作者们，也衷心祝愿圣诞、元旦不放假还在赶SIGGRAPH 2020的小伙伴们，能够一切顺利。

附录：

Taichi论文摘要：(SIGGRAPH Asia 2019) 三维体积数据通常具有空间稀疏性。为了利用这种性质，计算机图形学社区开发了层级体素稀疏数据结构，如SPGrid、VDB和八叉树等。但是，由于其内在复杂性和额外开销，开发、应用这些高性能数据结构有很多挑战。我们提出Taichi，一个新的面向（稀疏）数据的编程语言，大大降低了空间稀疏数据结构的开发、使用成本。由于Taichi实现了算法和数据结构的解耦，使用者可以快速尝试不同数据结构，以在特定问题和体系结构上找到最优数据结构。语言前端提供给用户易用的接口，使得用户可以以访问稠密数据结构的方式访问稀疏数据结构，大大提高了代码可读性和生产力。Taichi编译器使用对数据结构的语义和下标分析来优化程序的局部性，移除多余数据结构遍历，以及进行自动内存管理和并行化、向量化。在x86_64和CUDA体系结构上，只需要1/10的代码，Taichi程序就能比手动优化的稀疏计算基准程序快4.55倍，测试程序包括物质点法、有限元模拟、多重网格泊松方程求解，真实感渲染，和3D稀疏卷积神经网络等。项目主页和代码：https://github.com/yuanming-hu/taichi

DiffTaichi论文摘要：(ICLR 2020) 基于Taichi，我们提出可微编程语言DiffTaichi，用于构建端到端可微程序。和目前常用的可微编程工具如TensorFlow、PyTorch相比，DiffTaichi更适合构建比常用操作（如卷积、BN等）更不规则的可微运算符，比如可微物理引擎中的粒子网格交互，网格采样等等。DiffTachi的自动微分系统使用“两个尺度”设计：底层通过源代码变换保持并行性和算术强度(arithmetic intensity)，上层通过一个轻量级的磁带(Tape)来记录大内核(Megakernel)的启动。由于省去了枯燥的手动求导过程，DiffTaichi程序比CUDA短4.2倍并具有相同的性能；同时由于其Megakernel的设计，在编写复杂可微程序时，DiffTaichi比TensorFlow快188倍、比PyTorch快13.4倍。我们用DiffTaichi实现了10个不同的可微物理引擎。将神经网络嵌入其中，控制器优化可以在几十个梯度下降迭代中完成，这比增强学习(RL)收敛速度快若干数量级。项目主页和代码：https://github.com/yuanming-hu/difftaichi

(标题图来自互联网)

来源：知乎 www.zhihu.com

作者：知乎用户（登录查看详情）

【知乎日报】千万用户的选择，做朋友圈里的新鲜事分享大牛。
点击下载