AOS和SOA布局的运行速度

根据文档,在需要频繁访问两个field时,使用AOS布局,内存离得比较近,性能应该比SOA更好,但是我测试中发现好像不太一样。代码如下:

import taichi as ti
from time import time
ti.init()

N = 10000000

@ti.kernel
def post(x:ti.template(), y:ti.template() ):
    for i in range(N):
        x[i] += 2*y[i]

x = ti.field(ti.i32)
y = ti.field(ti.i32)
ti.root.dense(ti.i, N).place(x)
ti.root.dense(ti.i, N).place(y)
z = ti.field(ti.i32)
w = ti.field(ti.i32)
ti.root.dense(ti.i, N).place(z, w)
start = time()
post(x, y)
end = time()
print(f'time costs SOA with JIT: {end-start} s.')
repeat = 30
start = time()
for i in range(repeat):
    post(x, y)
end = time()
print(f'time costs SOA real time: {(end-start)/repeat} s.')
start = time()
post(z, w)
end = time()
print(f'time costs AOS with JIT: {end-start} s.')
start = time()
for i in range(repeat):
    post(z, w)
end = time()
print(f'time costs AOS real time: {(end-start)/repeat} s.')

输出如下:

[Taichi] version 1.6.0, llvm 15.0.1, commit f1c6fbbd, win, python 3.8.0
[Taichi] Starting on arch=x64
time costs SOA with JIT: 0.13309955596923828 s.
time costs SOA real time: 0.006166577339172363 s.
time costs AOS with JIT: 0.029431819915771484 s.
time costs AOS real time: 0.00807642141977946 s.

AOS布局在初次运行时候比SOA快,但是之后的速度反而更慢,这是为什么呢?还是我哪里理解不准确?

AOS 与 SOA 布局在不同场景下的性能会有差异,并不是说 AOS 就一定比 SOA 好,文档中似乎并没有讲清楚这件事情。内存排布带来的性能差距主要是来源于不同后端内存读写的粒度大小与 cache 策略,而这两点在 cpu 与 gpu 上有很大差距,甚至不同型号的 cpu 之间也会有不同,实际的运行效率会受到许多因素的影响,内存排布的优越性比较也离不开具体的机器和任务。

回到你的代码和运行结果,我觉得有两点值得说明:

  1. 我不确定 python 的计时工具是否可靠,可以考虑用 taichi 自带的一些 profiler。想要得到准确的计时最好还是运行多次后取平均结果。
  2. 实际上运行效率的差距只有 24% 左右,看起来并不算很显著,一些精心设计的 soa/aos 任务比较一般可以看到一些更显著的性能差距。

你好,关于第一点,我修改为重复运行30次的平均值,结果似乎是差不多的,主要是这个测试结果和文档里面描述的结果看起来是相反的,所以我会比较疑惑。

基本来说我和 @BillXu2000 的观点是一致的。如果想要AOS - SOA之间有巨大的性能差异,需要:

  1. 构建一个强Memory Bound的程序
  2. 在Memory Access以外,不能有其他的大overhead操作

如果你试一下Taichi doc里的例子,你会发现虽然AOS比SOA整体快那么一点点,但是并不显著。这都是因为Compute Bound > Memory Bound的原因。所以实际使用的时候其实能够提前预测AOS v.s. SOA的情况可能并不常见,大部分实践的时候是比如先用AOS写一般代码,需要优化的时候会尝试一下SOA看性能有没有提高这样。

下面是改造成Taichi docs:Fields (advanced) | Taichi Docs 的例子

import taichi as ti
from time import time
ti.init(arch=ti.cpu)

N = 200000
pos = ti.field(ti.f32)
vel = ti.field(ti.f32)

k = 10.0
dt = 2.0

# AoS placement
ti.root.dense(ti.i, N).place(pos, vel)

@ti.kernel
def step_aos():
    for i in pos:
        pos[i] += vel[i] * dt
        vel[i] += -k * pos[i] * dt


pos1 = ti.field(ti.f32)
vel1 = ti.field(ti.f32)
# SoA placement
ti.root.dense(ti.i, N).place(pos1)
ti.root.dense(ti.i, N).place(vel1)

@ti.kernel
def step_soa():
    for i in pos1:
        pos1[i] += vel1[i] * dt
        vel1[i] += -k * pos1[i] * dt


start = time()
step_soa()
end = time()
print(f'time costs SOA with JIT: {end-start} s.')
repeat = 30000
start = time()
for i in range(repeat):
    step_soa()
end = time()
print(f'time costs SOA real time: {(end-start)/repeat} s.')


start = time()
step_aos()
end = time()
print(f'time costs AOS with JIT: {end-start} s.')
start = time()
for i in range(repeat):
    step_aos()
end = time()
print(f'time costs AOS real time: {(end-start)/repeat} s.')

差距并不大,并且还会随系统状态波动:

time costs SOA with JIT: 0.06007862091064453 s.
time costs SOA real time: 7.249840100606282e-05 s.
time costs AOS with JIT: 0.04664349555969238 s.
time costs AOS real time: 7.22206989924113e-05 s.