能否减少外部python循环带来的性能损失？

综合讨论区
#1

核心代码如下：

@ti.kernel
def fd2d(it:int,ww:ti.template(),uu:ti.template(),
            fux:ti.template(),fuz:ti.template(),bwx:ti.template(),bwz:ti.template(),
           xx:ti.template(),zz:ti.template(),xz:ti.template()):
    for x, y in ti.ndrange((1, nzbc - 1), (1, nxbc - 1)):
#         print(x,y)
        uu[x, y] = temp[x, y] * uu[x, y] + b[x, y] * (xx[x, y + 1] - xx[x, y] + xz[x, y] - xz[x - 1, y]) 
        ww[x, y] = temp[x, y] * ww[x, y] + b[x, y] * (xz[x, y] - xz[x, y - 1] + zz[x + 1, y] - zz[x, y])

    ww[2,48]=s[it,0]  

    for x, y in ti.ndrange((1, nzbc - 1), (1, nxbc - 1)):
        fux[x, y] = uu[x, y] - uu[x, y - 1]
        fuz[x, y] = uu[x + 1, y] - uu[x, y]
        bwx[x, y] = ww[x, y + 1] - ww[x, y]
        bwz[x, y] = ww[x, y] - ww[x - 1, y]  

#     for x, y in ti.ndrange((1, nzbc - 1), (1, nxbc - 1)):
        xx[x, y] = temp[x, y] * xx[x, y] + (ca[x, y] * fux[x, y] + cl[x, y] * bwz[x, y])*dtx
        zz[x, y] = temp[x, y] * zz[x, y] + (ca[x, y] * bwz[x, y] + cl[x, y] * fux[x, y])*dtx
        xz[x, y] = temp[x, y] * xz[x, y] + (cm1[x, y] * (fuz[x, y] + bwx[x, y]))*dtx

    for y in ti.ndrange((0, nxbc)):
        zz[1,y]=0
    for x in ti.ndrange((0, ng)):
        seismo_w[it, x] = ww[2,40 + x * 2]
        seismo_u[it, x] = uu[2,40 + x * 2]
ti.profiler.clear_kernel_profiler_info()
T1 = time.time()
for it in range(1000):
#     print(it)
    fd2d(it,ww,uu,fux,fuz,bwx,bwz,xx,zz,xz)
#     if it%500==0:
#         imagesc(ww.to_torch())
T2 = time.time()
print(T2-T1)
ti.profiler.print_kernel_profiler_info()

0.4355292320251465

Kernel Profiler(count, default) @ CUDA on NVIDIA GeForce GTX 1060

[ % total count | min avg max ] Kernel name

[ 30.70% 0.009 s 1000x | 0.008 0.009 0.040 ms] fd2d_c80_0_kernel_4_range_for
[ 25.89% 0.008 s 1000x | 0.006 0.008 0.026 ms] fd2d_c80_0_kernel_0_range_for
[ 14.28% 0.004 s 1000x | 0.003 0.004 0.022 ms] fd2d_c80_0_kernel_3_range_for
[ 14.16% 0.004 s 1000x | 0.003 0.004 0.022 ms] fd2d_c80_0_kernel_1_serial
[ 12.86% 0.004 s 1000x | 0.003 0.004 0.022 ms] fd2d_c80_0_kernel_2_range_for
[ 0.67% 0.000 s 40x | 0.004 0.005 0.016 ms] jit_evaluator_2_kernel_0_serial
[ 0.53% 0.000 s 34x | 0.004 0.005 0.008 ms] jit_evaluator_3_kernel_0_serial
[ 0.13% 0.000 s 2x | 0.020 0.020 0.021 ms] jit_evaluator_7_kernel_0_serial
[ 0.12% 0.000 s 1x | 0.036 0.036 0.036 ms] jit_evaluator_14_kernel_0_serial
[ 0.11% 0.000 s 3x | 0.005 0.011 0.017 ms] jit_evaluator_20_kernel_0_serial
[ 0.10% 0.000 s 3x | 0.005 0.010 0.016 ms] jit_evaluator_19_kernel_0_serial
[ 0.08% 0.000 s 3x | 0.005 0.008 0.011 ms] jit_evaluator_21_kernel_0_serial
[ 0.08% 0.000 s 4x | 0.004 0.006 0.008 ms] jit_evaluator_9_kernel_0_serial
[ 0.07% 0.000 s 4x | 0.004 0.005 0.007 ms] jit_evaluator_1_kernel_0_serial
[ 0.05% 0.000 s 4x | 0.003 0.004 0.005 ms] jit_evaluator_0_kernel_0_serial
[ 0.03% 0.000 s 2x | 0.004 0.005 0.005 ms] jit_evaluator_4_kernel_0_serial
[ 0.03% 0.000 s 1x | 0.009 0.009 0.009 ms] jit_evaluator_12_kernel_0_serial
[ 0.03% 0.000 s 1x | 0.008 0.008 0.008 ms] jit_evaluator_17_kernel_0_serial
[ 0.03% 0.000 s 1x | 0.008 0.008 0.008 ms] jit_evaluator_16_kernel_0_serial
[ 0.02% 0.000 s 1x | 0.007 0.007 0.007 ms] jit_evaluator_15_kernel_0_serial
[ 0.02% 0.000 s 1x | 0.007 0.007 0.007 ms] jit_evaluator_13_kernel_0_serial

[100.00%] Total execution time: 0.030 s number of results: 21

以上代码做的是弹性波模拟：ti.kernel中用for循环自动并行了波场应力等计算；用外部python循环按时间序列模拟弹性波的传播。
可以看到python time统计的时间比内核时间多很多，随着循环次数的增加二者时间差还会进一步增加。
有没有办法能加快外部调用ti.kernel时的循环速度？
总而言之需求外部循环顺序执行，内部循环并行执行高效运行方式。

是否能在串行的循环中执行并行的循环？
#2

把循环包装在kernel中可行吗
计算波场应力用func

#4

不行，func和写一个大kernal没区别--------

#5

能方便贴一下完整的可以运行的代码吗
我用一些小函数测试，效果还是挺明显的

#6

您可以看看，关键是最外层循环要顺序执行，里层循环并行执行，非常感谢

#7

呃 你这个代码太长了把…不是很方便…

#8

能否分享一下您的小函数测试结果

#9

关注下这个话题。

#10

你可以试一下把最外层循环写到 kernel 里面，然后设置 loop_config 让它串行执行，再在下一层并行？

#11

loop_config后下一层的for循环能够并行执行啊？

#12

更正下，想达到这个效果应该是用 loop unroll。参考下面的例子：

import taichi as ti

ti.init(arch=ti.gpu)

@ti.kernel
def f():
    for i in ti.static(range(10)):
        for j in range(100):
            print(i, j)


f()

运行可以看到最外层串行，内层并行。因为最外层被 unroll 了之后，内层变成了实际上的最外层。

#13 
import taichi as ti
import time
ti.init(arch=ti.gpu)

@ti.kernel
def f1():
    for i in ti.static(range(100)):
        for j in range(100):
             print(i, j)
@ti.kernel
def subs():
    for j in range(100):
        print(j)
        
def f2():
    for i in range(100):
        subs()
        

T1=time.time()
f1()
T2=time.time()
Tf1 = T2-T1
print(Tf1)

T1=time.time()
f2()
T2=time.time()
Tf2 = T2-T1
print(Tf2)

loop unroll 确实是对的，但是如果unroll 100 1000个循环还是没优势
Tf1= 0.6472702026367188
Tf2= 0.06530094146728516

#14

unroll 之后编译时间会比较久，但是有了 cache 之后应该会快的。

#15

我这边的测试结果：

[Taichi] version 1.6.0, llvm 15.0.4, commit f1c6fbbd, linux, python 3.10.6
[Taichi] Starting on arch=cuda
0.33597517013549805 ← loop unroll 第一次运行（包含编译）
0.013126373291015625 ← loop unroll 第二次运行
0.023073911666870117 ← 外部 loop