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

[houkensjtu]

大家好！

为了更深入地理解使用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()更推荐更准确的做法？

谢谢各位大佬！

[BillXu2000]

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 语言本身。

[houkensjtu]

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

谢谢大佬

[BillXu2000]

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

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

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

[victoriacity]

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()取第二遍的用时可能更准确。

[houkensjtu]

@victoriacity

numpy的“正确”用法是尽可能少用for loop，尽可能多用向量化操作

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

@victoriacity

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

这个我会试一试，谢谢指点！

@victoriacity

访问一遍大矩阵的所有元素主要还是memory bound的，所以运行的时间会主要取决于内存带宽。

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

[archibate]

另外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

[Msure]

想请问一下，taichi的cpu并行到底是什么逻辑呀？没办法像openmpi一样指定核数运行吗？

[archibate]

和openmp一样的逻辑，针对顶层for循环，每个线程分配到若干个循环体(iteration body)。试试看ti.init(cpu_max_threads=2)来指定线程数？

[Msure]

好的，谢谢，没很搞清楚taichi的cpu并行到底是怎样。那这么说的话 在大规模集群服务器上，taichi的并行并不如基于opemmpi编写的c或c++代码了吧？也不知道taichi的cuda并行和其他采用openmpi的并行代码效率相比怎样。我不是学计算机的，所以很迷惑。官方文档没看到详细的对并行的解释，是不是只能看taichi源码了呀？

[archibate]

我来解释一下，并行基本上就是把一个大任务分拆成许多个小任务，在不同的CPU核心上运行，从而达到减少计算时间的目的。
比如我的大任务是计算从1加到100的值。按照传统的做法是这样：

sum = 0
for i from 1 to 100:
  sum += i
print sum

但是这样的一个for循环里的每一个sum = sum + i是顺序执行的，只利用了一个CPU核心，因此效率较低，不能充分发挥多核心的作用。
注意到1加到100这个任务，每一次相加和前一次相加的结果其实没有依赖性，我们完全不一定要一个个串联在一个线程里执行。

那么并行的做法就是首先把1加到100分成四个小任务：

1. 1加到25
2. 26加到50
3. 51加到75
4. 76加到100

然后在4个CPU核心上同时执行，因为每个小任务的时长是原来的1/4，四个CPU因为同时执行所以总共的时长是原来的1/4，理论上会快4倍：

sum = 0

start thread 1:
  for i from 1 to 25:
    sum += i

start thread 2:
  for i from 26 to 50:
    sum += i

start thread 3:
  for i from 51 to 75:
    sum += i

start thread 4:
  for i from 76 to 100:
    sum += i

wait threads done
print sum

但是要自己操纵线程的启停来实现并行，不如原来单独一个for循环直观了，牺牲了代码的可读性。
因此太极采用了自动for循环并行化：你只需要写出一个普通的for循环，而不需要管多线程具体的实现细节，他会自动帮你启停多个线程来利用多核心CPU。

@ti.kernel
def calc() -> int:
  sum = 0
  for i in range(1, 101):  # 自动分拆成小任务在多个CPU核心执行
    sum += 1
  return sum

print(calc())  # 应得到 5050

不过，并行有好处也有坏处，那就是他可能导致数据的竞争 (data race)，不同的线程如果同时读写同一个变量，有时会导致结果不正确。不过不用担心，太极已经帮你考虑到了，只要你用的是 sum += 1 的写法，这个累加会被自动处理成原子性的 (atomic)。

@ti.kernel
def calc() -> int:
  sum = 0
  for i in range(1, 101):
    sum = sum + 1  # 非原子性的写法
  return sum

print(calc())  # 多线程数据竞争！结果可能会小于 5050

此外，CUDA等GPU后端也是同理，不过GPU的线程（相当于CPU的多核）数目更多，虽然他们单个核心执行的速度更慢（高延迟），但是他们的核心数目多呀（大吞吐量），因此更擅长处理高度并行的任务。
希望以上内容对你有所帮助

[archibate]

关于openmpi：据我所知mpi是一个进程间通信的协议？而openmp是一个cpp里和太极一样能自动并行for循环的框架？
如果你是说mpi，我想太极的设计目标和他不一样，应用场景也有所不同。太极的并行主要是针对单台计算机上的多个CPU核心，调度单位是线程，优势是for循环易于编写而直观，但无法扩展到多台计算机。而mpi可以针对多台计算机组成的系统，提供之间通信的桥梁，优势是可以在许多计算机之间通信，但效率不如单台之间高效。

[Msure]

是的 感谢 你的解答很详细 那这么说的话MPI确实更适合在超算集群上运行 taichi在单台电脑上做到了很高的自动并行效率 但是上亿网格那普通电脑还是不行了 采用MPI上超算集群可能是更好的解决方法 但是基于我目前的工程应用需求我就一般几百万或上千万网格用taichi应该是性价比最高的 我到时候应该会比较一下速度 感谢感谢！