给Taichi矩阵运算测速时的疑问

大家好!

为了更深入地理解使用Taichi的正确姿势,我尝试做了一个简单的矩阵访问测速,意在比较Taichi和numpy的速度区别:

# Taichi version
import taichi as ti
import time

ti.init()
nx, ny = 10000,10000
p = ti.field(dtype=ti.f32, shape=(nx,ny))
@ti.kernel
def main():
    for i,j in p:
        p[i,j] = p[i,j] + 0.1

start = time.time()
main()
print(time.time()-start)

得到的结果是约1.3 sec。

然后写了一个numpy版本的同样的运算,如下:

# numpy version
import numpy as np
import time

nx, ny = 10000,10000
pnp = np.zeros((nx,ny))
def main():
    for i in range(nx):
        for j in range(ny):
            pnp[i,j] = pnp[i,j] + 0.1

start = time.time()
main()
print(time.time()-start)

结果发现大约是34秒,和Taichi比相差了近30倍。。

这里想问两个问题:

  • 我的CPU只是一个i7-6600U,为何同样的运算会相差30多倍 我发现,当矩阵越小的时候结果就会越接近,哪怕是1000 x 1000的矩阵,两者就会降到几乎相同的数量级上。这是合理的现象吗?如果是的话,为什么?
  • 测速的时候,有没有比用time()更推荐更准确的做法?

谢谢各位大佬!

import numpy as np
from timeit import Timer
import time

nx = 10000
ny = nx

pnp = np.zeros((nx,ny))
def main():
    global pnp
    pnp += 0.1

start = time.time()
main()
print(time.time()-start)

写成这样的 numpy 版本跑出来的速度跟你的 taichi 版本速度差不多。
你的 numpy 版本的速度瓶颈可能在于 python 语言本身。

1 Like

谢谢点拨!我试了一下你说的没错,而且是numpy的版本更快了。这样的话我还是有几个问题:
1.改成第二种以后numpy比taichi更快,而taichi应该是并行处理的,numpy比并行更快是因为什么?
2.也就是说用range-loop来循环numpy矩阵是很不合理的,那正确的遍历应该是要怎么写呢?你给出的例子是没有具体去操作每个元素的。

谢谢大佬

我 python 学的不好,这边我只能写一些自己的推测,没有什么正确性保证,如果有错误欢迎指出。

  1. 有可能 numpy 本身就支持并行;也有可能这个代码的速度瓶颈在于内存读取而不是 cpu 运算,因此是否并行不会影响速度。至于 numpy 为啥比 taichi 稍微快一点我也不知道,可能 numpy 实现得比较好吧。

  2. 正确的遍历要用别的语言写或者调用一些 python 库函数。taichi 应该算是套着 python 皮的另一种语言。

1 Like

numpy的“正确”用法是尽可能少用for loop,尽可能多用向量化操作(整个数组加减之类)和numpy自带的函数。因为numpy是在调用这些函数和数组操作的时候去用C或者fortran后端,python解释器还是会一遍一遍去跑循环里的每一圈的,而不是像taichi一样用即时编译处理整个for loop,和taichi比较像的数值计算包可以参考numba和jax。numpy如果向量化运算写的好也是可能非常快的,不过有时候不太好想,比如把矢量运算打包成矩阵运算之类。

numpy似乎对指令集层面的SIMD优化得也不错,甚至支持avx512,所以CPU上的“并行”效率应该不输taichi。对楼主的程序来说,访问一遍大矩阵的所有元素主要还是memory bound的,所以运行的时间会主要取决于内存带宽。另外taichi会在第一遍调用kernel的时候才执行编译过程,所以运行两遍main()取第二遍的用时可能更准确。

2 Likes

这个我也意识到了,而在Taichi里去跑for-loop也可以获得和numpy向量化以后差不多的速度,应该已经是非常厉害的了。

这个我会试一试,谢谢指点!

这个有办法用什么Taichi的调试工具看出来吗?

另外taichi会在第一遍调用kernel的时候才执行编译过程,所以运行两遍main()取第二遍的用时可能更准确。

除此之外,其实还有一个缓存预热(cache warmup)的效应。比如我发现下面numpy这个测试如果没有warmup速度就会慢将近两倍。
还有一点,其实taichi的默认精度是32位,而numpy是64位。

# Taichi version
import taichi as ti
import time

ti.init()
nx, ny = 10000,10000
p = ti.field(dtype=ti.f64, shape=(nx,ny))
@ti.kernel
def main():
    for i,j in p:
        p[i,j] = p[i,j] + 0.1

main()  # JIT compile
start = time.time()
main()
print(time.time()-start)

# numpy version
import numpy as np
import time

nx, ny = 10000,10000
pnp = np.zeros((nx,ny))

_ = pnp + 1  # cache warm up
start = time.time()
pnp = pnp + 0.1
print(time.time()-start)

# numba version
import numpy as np
import numba  # 一个看似和Taichi很相似的包,不过它是针对更广泛的数据类型的,比如numpy array甚至是Python自带的list
import time

nx, ny = 10000,10000
pnp = np.zeros((nx,ny))
@numba.jit(numba.void(numba.f8[:]))
def main(pnp):
    for i in range(nx):
        for j in range(ny):
            pnp[i,j] = pnp[i,j] + 0.1

main(np.zeros((nx, ny)))  # JIT compile
start = time.time()
main(pnp)
print(time.time()-start)

在我的电脑上结果分别是:
Taichi: 0.2829749584197998
NumPy: 0.12112188339233398
Numba: 0.5453431606292725

1 Like