是否能在串行的循环中执行并行的循环?

能否减少外部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]
1 个赞

@StellarWarp 如果想在串行的循环中执行并行循环,目前最简单的办法确实是把串行写在python,然后把并行放进taichi kernel, 如果需要复用代码的话可以放进ti.func.

如果python的调用overhead是一个比较大的concern的话,可以试试cgraph, taichi/mpm88_graph.py at master · taichi-dev/taichi · GitHub 能显著的降低python overhead

感谢大佬的解答