关于如何将numpy传入taichi.kernel

综合讨论区
#1

将numpy数组作为函数参数ti.types.ndarray()似乎会很花费时间，例如：

import numpy as np
import taichi as ti
import time

ti.init(arch=ti.gpu)

@ti.kernel
def sum_thousand(x: ti.types.ndarray()) -> ti.f32:
    s = 0.0
    for i in range(1000):
        s += x[i]
    return s

xs1 = np.arange(0, 10, 0.001)

start = time.time()
for i in range(1000):
    sum_thousand(xs1)
print("duration(s): ", time.time() - start)

xs2 = np.arange(0, 100, 0.001)
start = time.time()
for i in range(1000):
    sum_thousand(xs2)
print("duration(s): ", time.time() - start)

执行输出为：

[Taichi] Starting on arch=cuda
duration(s):  0.659170389175415
duration(s):  1.0771942138671875

可以看出两次计算的计算量完全一致，但是由于传入数组的长度不同，计算时间有显著差异。

请问有什么更好的传入数据的方法吗？

谢谢！

#2

在测试performance的时候，第一次还包含编译时间，所以一般都先调用一次kernel,然后再测时间。最后的结果看起来时间差不多。

import numpy as np
import taichi as ti
import time

ti.init(arch=ti.gpu, default_fp=ti.f32)

@ti.kernel
def sum_thousand(x: ti.types.ndarray()) -> ti.f32:
    s = 0.0
    for i in range(1000):
        s += x[i]
    return s

xs1 = np.arange(0, 10, 0.001,dtype=np.float32)

sum_thousand(xs1)
start = time.time()
for i in range(1000):
    sum_thousand(xs1)
print("duration(s): ", time.time() - start)

xs2 = np.arange(0, 100, 0.001,dtype=np.float32)
sum_thousand(xs2)
start = time.time()
for i in range(1000):
    sum_thousand(xs2)
print("duration(s): ", time.time() - start)

image
image1236×476 75.7 KB

#3

原来同样的函数对于不同数组的调用也是要重新编译的呀！

但是在cuda上测试了一下，似乎还是有不小的差距。

另外，我观察到cuda上写上dtype=np.float32可以显著地优化速度，请问这是为什么？np.float32会比float好一些吗？

import numpy as np
import taichi as ti
import time

ti.init(arch=ti.gpu, default_fp=ti.f32)

@ti.kernel
def sum1000(x: ti.types.ndarray()) -> ti.f32:
    s = 0.0
    for i in range(1000):
        s += x[i]
    return s

def test(x):
    sum1000(x)
    start = time.time()
    for i in range(1000):
        sum1000(x)
    print("duration(s): ", time.time() - start)
    
test(np.arange(0, 10, 0.001, dtype=np.float32))
test(np.arange(0, 100, 0.001, dtype=np.float32))
test(np.arange(0, 100, 0.001))

输出为

[Taichi] Starting on arch=cuda
duration(s):  0.8674345016479492
duration(s):  0.9991717338562012
duration(s):  1.772221565246582
1 个赞
#4

Hi @qiao-bo @haidonglan，你们有什么看法么？

#5

python的float和numpy array的默认dtype是FP64，所以numpy的array一定要指定np.float32。GPU上很少有具备完整的FP64的计算能力的，即使有也是FP32的1/2或者1/3

#6

numpy数组是动态的话，这种对JIT模式会很麻烦。有没有可能先assume一个合适的数组大小，不够用的时候再做reallocation？

#7

另外，你可以在ti.init这里加一个kernel_profile=True，然后调用ti.profiler.print_kernel_profiler_info()来看一下时间的分布 如果是示例代码的话能看到大部分开销都在编译上

image
image1874×546 44 KB