从能否减少外部python循环带来的性能损失?继续讨论:
这是一个简单的并行前缀和算法的代码,算法示意图如上
@ti.kernel
def prefix_sum_interal(t:ti.u32,h:ti.u32,x:ti.template(),y:ti.template()):
for i in range(h):
src = ((i>>t)<<(t+1)) + (1<<t) - 1
dst = src + 1 + (i & ((1<<t) - 1))
y[dst] += y[src]
def prefix_sum(x:ti.template(),y:ti.template()):
n = x.shape[0]
y.copy_from(x)
h = n // 2
total_step = q_log2(n)
for t in range(total_step):
prefix_sum_interal(t,h,x,y)
因为每一步计算需要一层执行结束后才能进入下一层。但是taichi目前应该还不支持如下的写法:串行的循环中指定并行循环,这样做会导致变成完全串行执行。最终我只能将串行循环写在python层
@ti.kernel
def prefix_sum(x:ti.template(),y:ti.template()):
n = x.shape[0]
h = n // 2
total_step = q_log2(n)
for i in range(x):
y[i] = x[i]
ti.loop_config(serialize=True)
for t in range(total_step):
ti.loop_config(serialize=False)
prefix_sum_interal(t,h,x,y)
在【能否减少外部python循环带来的性能损失?】中也提到了使用循环展开的方式来使得串行的循环中的并行循环变成最外层循环,但是对于大多数场景来说,还是需要动态的串行循环
并且如果只能在python层串行地调度kernel执行并行循环,代码无法在其它kernel中使用
当然,调换循环的嵌套关系是可以的,下面这段代码是在并行循环中执行串行循环,但是这会带来同步问题,当数组长度大于最大可分配的线程数时,下面这段代码会给出错误结果
@ti.kernel
def prefix_sum(x:ti.template(),y:ti.template()):
n = x.shape[0]
h = n // 2
total_step = q_log2(n)
for i in range(x):
y[i] = x[i]
for i in range(h):
for t in range(total_step):
src = ((i>>t)<<(t+1)) + (1<<t) - 1
dst = src + 1 + (i & ((1<<t) - 1))
y[dst] += y[src]