关于kernel参数的问题

import taichi as ti
import numpy as np
ti.init()
@ti.func
def fff(z,p):
    z[0]=z[0]+p[0]
    z[1]=z[1]+p[1]
    return p[0]+p[1]+p[2]+z[0]+z[1]
def f(p):
    global z
    return fff(z,p)

x = ti.Vector.field(3, ti.f32, 4)
z=[1,2]
y=np.array([[1,2,3],[2,3,4],[3,4,5],[4,5,6]])
x.from_numpy(y)
result = ti.field(ti.f32, 4)

@ti.kernel
def ff():
    for i in range(4):
        #result[i] = f(y[i]) #y为ndarray,报错Indexerror
        result[i] = f(x[i]) #x为vector , 正常运行
        #result[i] = f([1,2,3]) #正常运行

ff()
print(result)

各位老师:
在运行上述代码时,ff函数中我想将它并行化,但直接使用y[i]时会报Indexerror,使用ti.vector则不会,但直接使用一个列表[1,2,3]作为参数也是可以运行的,这是为什么呢?除了使用from_numpy外,还有什么方法可以让我使用到y吗?

这里报indexerror的原因是ff()这个kernel内是一个Taichi range for loop,循环拿到的i不是一个int,是一个taichi expr,所以numpy array报了indexerror;要使用y的话可以这样写:

@ti.kernel
def ff(y:ti.any_arr()):
    for i in range(4):
        result[i] = fff([1,2], [y[i, j] for j in range(3)])

另外, 在taichi scope里(也就是kernel里)尽量调用ti.func这种taichi 的函数

感谢!

1 个赞