gputools.py 文件源码

def laplace_stack_gpu(y_gpu, mode='valid'):
  """
  This funtion computes the Laplacian of each slice of a stack of images
  """
  shape = np.array(y_gpu.shape).astype(np.uint32)
  dtype = y_gpu.dtype
  block_size = (6,int(np.floor(512./6./float(shape[0]))),int(shape[0]))
  grid_size = (int(np.ceil(float(shape[1])/block_size[0])),
               int(np.ceil(float(shape[0])/block_size[1])))
  shared_size = int((2+block_size[0])*(2+block_size[1])*(2+block_size[2])
                    *dtype.itemsize)

  preproc = _generate_preproc(dtype, (shape[1],shape[2]))
  mod = SourceModule(preproc + kernel_code, keep=True)

  laplace_fun_gpu = mod.get_function("laplace_stack_same")
  laplace_gpu = cua.empty((y_gpu.shape[0], y_gpu.shape[1], y_gpu.shape[2]),
                          y_gpu.dtype)

  laplace_fun_gpu(laplace_gpu.gpudata, y_gpu.gpudata,
                  block=block_size, grid=grid_size, shared=shared_size)

  return laplace_gpu