对于numba中出现这种问题numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend) Cannot,就拿下面这个代码来说:

x = np.arange(100).reshape(10, 10)
@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a):
    trace = 1
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace
print(go_fast(x))

这里是类型错误,就是trace(int)的类型和a(float)的类型不匹配,所以只需要把trace=1改成trace=1.0或者trace=float(1)或者trace=np.float32(1)等等:

x = np.arange(100).reshape(10, 10)
@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a):
    trace = 1.0#trace=float(1)#改这里
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace
print(go_fast(x))

还有这种问题TypeError: No matching definition for argument type(s) array(int32, 2d, C),原因是参数没用匹配到或者匹配不正确。如下面代码:

@jit("float32(float32)",nopython=True) 
def go_sum(a):
    result=a+1
    return result
x = np.arange(100).reshape(10, 10)
print(go_sum(x))

这里x是一维数组(就是向量),而我们定义的参数是一个浮点型的参数,所以不能将数组带进去,这里这样改:

@jit("int32(int32)",nopython=True) 
def go_fast(a):
    result=a+1
    return result
cc=1#cc=1或者1.0都行,因为我们已经在前面的jit里面定义好类型了,所以cc可以是浮点型或者整型
print(go_fast(cc))

接下来是关于numba中@cuda.jit的使用:从官网英文找了一个自定义的模块如下:中文文档关于AP部分I

from numba import vectorize, cuda
# define a device function
@cuda.jit('float32(float32, float32, float32)', device=True, inline=True)
def cu_device_fn(x, y, z):
    return x ** y / z

但是这个函数如何调用,官网并没用讲。根据官网对numba的定义,有如下方法可以调用。
第一种就是先定义一个装饰器为cuda.jit的主函数,并在内部定义输入和输出变量:

@cuda.jit()#jit这里有无括号都行
def main_func(ipx,ipy,ipz,out_xx):
    #申请线程索引
    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    #申请块索引
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y
    #申请块内线程数
    bw = cuda.blockDim.x
    bh = cuda.blockDim.y
    #计算数据下标
    idx = tx + ty * bw + bx * bw * bh
    XX = ipx[idx]
    YY = ipy[idx]
    ZZ = ipz[idx]
    #调用设备函数
    out_xx[idx]=cu_device_fn(XX,YY,ZZ)
if __name__=="__main__":
    ipx = np.array([1, 2, 3],dtype=np.float32)
    ipy = np.array([4, 5, 6],dtype=np.float32)
    ipz = np.array([7, 8, 9],dtype=np.float32)
    out_xx =cuda.device_array_like(ipy)
    main_func[64,2](ipx,ipy,ipz,out_xx)#这里的[64,2]不是用的很严谨

第二种也是需要先编写一个在GPU上运行的函数,然后用该函数调用设备函数:

@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')#也可以写成@cuda.jit()或者@cuda.jit
def cu_host_fn(x,y,z,out):
    i=cuda.grid
    if i<x.size:
        out[i]=cu_device_fn(x[i],y[i],z[i])#调用设备函数
    '''这里cuda.grid(1)语句用于获取线程在ID线程块中的位置,以便可以将其映射到输入数组的相应元素上。
    如果线程索引i小于x的大小,则计算设备函数的结果,并将其存储在out[i]中。'''
if __name__=='__main__':
    ipx = np.array([1, 2, 3],dtype=np.float32)
    ipy = np.array([4, 5, 6],dtype=np.float32)
    ipz = np.array([7, 8, 9],dtype=np.float32)
    out_xx =cuda.device_array_like(ipy)
    thread_per_block=64
    block_per_grid=(x.size+(thread_per_block-1))//thread_per_block
    #调用主机函数
    cu_host_fn[thread_per_block,block_per_grid](ipx,ipy,ipz,out_xx)
    #打印结果
    print(out_xx.copy_to_host())#copy_to_host是复制设备->主机

第三中呢就是更改设备函数cu_device_fn:

from numba import vectorize, cuda
# define a device function
@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')#也可以写成@cuda,jit
def cu_device_fn(x, y, z,out_xx):
    tid=cuda.grid(1)
    if tid<len(out_xx):
        for i in range(len(x)):
            out_xx[i]=x[i] ** y[i] / z[i]

上面的这个for循环其实是多余的,应该去掉,然后写成如下:

from numba import vectorize, cuda
# define a device function
@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')
def cu_device_fn(x, y, z,out_xx):
    tid=cuda.grid(1)
    if tid<len(out_xx):
        out_xx[idx]=x[idx] ** y[idx] / z[idx]
if __name__=='__main__':
    ipx = cuda.to_device(np.array([1, 2, 3], dtype=np.float32))
    ipy = cuda.to_device(np.array([4, 5, 6], dtype=np.float32))
    ipz = cuda.to_device(np.array([7, 8, 9], dtype=np.float32))
    out_xx = cuda.device_array_like(ipy)
    '''
    to_device()可用于创建阵列的设备端副本,即分配一个 numpy ndarray 或结构化标量并将其传输到设备。
    array_like()是使用数组中的信息调用。而device_array_like()创建一个与现有数组具有相同形状和类型的未初始化数组
    '''
    cu_device_fn[64, 2](ipx, ipy, ipz, out_xx)
    print(out_xx.copy_to_host())#打印结果
Logo

欢迎来到由智源人工智能研究院发起的Triton中文社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐