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)


发表评论

电子邮件地址不会被公开。 必填项已用*标注

Powered by WordPress. Design: Supermodne.