在 numba 函数中，如果不在 CUDA 中，请将 cuda.popc() 替换为 CPU 等效项

Question

我正在编写通用代码，支持 CPU 上的 numba-jitting 和 GPU 上的 numba.cuda-jitting。

一切都运行良好，除了在公共代码的深处，我想使用一个内部指令来计算整数中的位数。 它是 CUDA 路径的

cuda.popc()

，以及 CPU 路径的辅助函数

cpu_popc()

。 不幸的是，

cuda.popc

仅在numba.cuda-jitted GPU内核中有效，

cpu_popc

仅在numba-jitted CPU函数中有效。

有没有办法同时实现

cpu_compute

和

gpu_compute

——无需复制整个公共代码？

这是一个简单的框架来测试这一点：

# CPU equivalent ctpop, from https://stackoverflow.com/a/77103233
@numba.extending.intrinsic
def popc_helper(typing_context, src):
  def codegen(context, builder, signature, args):
    return numba.cpython.mathimpl.call_fp_intrinsic(builder, "llvm.ctpop.i64", args)
  return numba.uint64(numba.uint64), codegen

@numba.njit(numba.uint64(numba.uint64))
def cpu_popc(x):
  """Return the (population) count of set bits in an integer."""
  return popc_helper(x)

@numba.njit
def common_function(x):
  # ...
  # some_long_code_that_should_not_get_duplicated.
  # ...
  # return cpu_popc(x)  # This works on the CPU path.
  return cuda.popc(x)  # This works on the GPU path.

@numba.njit
def cpu_compute(n=5):
  array_in = np.arange(n)
  array_out = np.empty_like(array_in)
  for i, value in enumerate(array_in):
    array_out[i] = common_function(value)
  return array_out

@cuda.jit
def gpu_kernel(array_in, array_out):
  thread_index = cuda.grid(1)
  if thread_index < len(array_in):
    array_out[thread_index] = common_function(array_in[thread_index])

def gpu_compute(n=5):
  array_in = np.arange(n)
  array_out = cuda.device_array_like(array_in)
  gpu_kernel[1, len(array_in)](cuda.to_device(array_in), array_out)
  return array_out.copy_to_host()

# print(cpu_compute())
print(gpu_compute())

Answer 1

我可以通过使用如下所示的工厂函数来使其工作（并将

cuda.popc

包装到 jitted 函数中）。

尽管如此，这并不理想。 有更好的解决方案吗？

# Earlier code plus...

@numba.njit
def gpu_popc(x):
  return cuda.popc(x)
  
def make_common_function(popc):

  def common_function(x):
    # ...
    # some_long_code_that_should_not_get_duplicated.
    # ...
    return popc(x)  # Works on both CPU and GPU path.

  return common_function

common_function_numba = numba.njit(make_common_function(cpu_popc))
common_function_cuda = numba.njit(make_common_function(gpu_popc))

@numba.njit
def cpu_compute(n=5):
  array_in = np.arange(n)
  array_out = np.empty_like(array_in)
  for i, value in enumerate(array_in):
    array_out[i] = common_function_numba(value)
  return array_out

@cuda.jit
def gpu_kernel(array_in, array_out):
  thread_index = cuda.grid(1)
  if thread_index < len(array_in):
    array_out[thread_index] = common_function_cuda(array_in[thread_index])

def gpu_compute(n=5):
  array_in = np.arange(n)
  array_out = cuda.device_array_like(array_in)
  gpu_kernel[1, len(array_in)](cuda.to_device(array_in), array_out)
  return array_out.copy_to_host()

print(cpu_compute())
print(gpu_compute())