Learning Taichi

我都好久没学过新语言了。大三的时候我就知道胡同学在设计一门图形学编程语言“太极”（当时好像还只是加速计算的库），当年他把早期版本刚刚打包上传到PyPI的时候，我还跑了一次pip install看看能不能顺利装上。然而之后我就再没碰过，只是见证了他逐渐成名，也越来越多人注意到这门语言，“99行代码的《冰雪奇缘》”那篇文章火遍大江南北。听说他开了公开课，我自然也去凑热闹，然而终究是纸上谈兵，学会了语法跑通了样例，之后就淡忘了。

然而这一天，数学家陈同学向我抱怨，自己的Matlab代码要跑很久，想问我在学校集群上跑代码能不能比自己笔记本上快两个数量级。我很警惕，决定先看看现有的代码还有没有优化空间。他拿着一台刚买回来的又有中高端CPU又有独立GPU的笔记本，理应还能压榨出很多性能才对。

他于是丢给我一个这样的计算问题：给定一个N=100行，M=400*400列的矩阵，计算所有行两两之间的L1距离。他的Matlab代码要算11秒。用最原始的python代码，速度在同一个数量级：

for i,j in itertools.product(range(n),range(n)):	
	ret[i,j]=np.sum(np.abs(mat[i,:]-mat[j,:]))

我先跟他解释了一下CPU体系架构的基础知识（指令集并行，一二三级缓存），然后强迫他看完了Halide语言的介绍视频（胡同学当年让我看过）作为入门，大概说明了为啥重排循环能大幅提升代码性能。然后，我指出GPU算力比CPU多很多很多倍，我们应该设法用GPU来算。最后，我向他大力推销太极语言，因为性能优化、适配具体计算硬件的工作应该由编译器进行。我决定现场学习，然后当场用太极写出算法。

安装一如既往的顺利，和多年前一样。看着胡同学引以为傲的由社区维护的中文文档，我磕磕碰碰的写出了第一版代码。上GPU！然而刚跑起来，就碰壁了，屏幕直接冻结，必须强制重启。改成用CPU跑，等了一会之后却能跑通，也不知道发生了什么。在CPU上debug了一会，调整了一下循环顺序，速度快了一些；这时同样的代码再去GPU上跑也能跑通了，然而比CPU慢了几十倍。不过至少，CPU版比Matlab快一些；陈同学拒绝用太极重写自己前后所有代码，于是他拿着太极编译器打出的中间表示，试着去Matlab里实现类似的循环展开逻辑。

可是GPU上跑怎么会这么慢？我想了想，可能是因为我不知怎么正确的写出向量化。然而时间有限，我也只能不求甚解，第一天写太极就这样草草收场。

我于是向作者胡同学发了些牢骚，几周后收到回复，最内层循环的向量化应该单独提出来写，以及还有许多别的trick可以写出更快的太极代码，还得熟能生巧。他还指我去看他在SIGGRAPH上给的tutorial，啊，我当时怎么就没找到呢。我知道太极编译器只对最外层循环做并行优化，当时只想到把最长的循环写在最外层，而正解是把最内层循环抽出来变成函数，一下就快很多了。

@ti.func
def L1dist(i,j):
	s=0.0
	for k in range(m):
		s+=ti.abs(mat[i,k]-mat[j,k])
	return s

@ti.kernel
def calc():
	for i,j in ret:
		ret[i,j]=L1dist(i,j)

2020 年 08 月 22 日

matlab, python, taichi

技术派