gputools.py 文件源码

def comp_ola_sdeconv(gx_gpu, gy_gpu, xx_gpu, xy_gpu, Ftpy_gpu, f_gpu, L_gpu, alpha, beta, gamma=0):
    """
    Computes the division in Fourier space needed for sparse deconvolution
    """

    sfft = xx_gpu.shape
    block_size = (16,16,1)   
    grid_size = (int(np.ceil(np.float32(sfft[0]*sfft[1])/block_size[0])),
                 int(np.ceil(np.float32(sfft[2])/block_size[1])))

    mod = cu.module_from_buffer(cubin)
    comp_ola_sdeconv_Kernel = mod.get_function("comp_ola_sdeconv_Kernel")

    z_gpu = cua.zeros(sfft, np.complex64)

    comp_ola_sdeconv_Kernel(z_gpu.gpudata,
                            np.int32(sfft[0]), np.int32(sfft[1]), np.int32(sfft[2]),
                            gx_gpu.gpudata, gy_gpu.gpudata,
                            xx_gpu.gpudata, xy_gpu.gpudata, 
                            Ftpy_gpu.gpudata, f_gpu.gpudata, L_gpu.gpudata,
                            np.float32(alpha), np.float32(beta),
                            np.float32(gamma),
                            block=block_size, grid=grid_size)

    return z_gpu