如何让Python CUDA中的atomicAdd支持long int类型

问题根源 #

1. Windows平台下C/C++的long是32位类型，而主机端使用的np.int64是64位，类型不匹配导致编译错误
2. CUDA原生atomicAdd对64位整数的支持针对long long（等价于int64_t）类型，而非long
3. 不能对内核的传值参数（如InitianCount）调用atomicAdd，每个线程会持有该参数的独立副本，原子操作必须作用在全局内存指针上

修正后的代码 #

import os

_path = r"C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Tools\MSVC\14.33.31629\bin\Hostx64\x64"

if os.system("cl.exe"):
    os.environ['PATH'] += ';' + _path
if os.system("cl.exe"):
    raise RuntimeError("cl.exe still not found, path probably incorrect")

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np

# 初始化64位计数器
result_count = np.zeros(1, dtype=np.int64)
result_count_gpu = cuda.mem_alloc(result_count.nbytes)
cuda.memcpy_htod(result_count_gpu, result_count)

record_result = np.zeros((100000000, 4)).astype(np.float32) 
record_result_gpu = cuda.mem_alloc(record_result.nbytes)

print('result_count.nbytes is ' + str(result_count.nbytes))

mod = SourceModule("""
  #include <stdint.h>  // 使用标准64位整数类型

  __global__ void test_cuda_LongIntArray(int64_t *result_count, float *record_result)
  {
      // 对全局内存的64位计数器执行原子加操作
      int64_t result_index = atomicAdd(result_count, 1LL);
      // 可根据result_index写入record_result，示例暂不实现具体逻辑
  }
  """, no_extern_c=True)  # 避免CUDA内核函数的C链接冲突
      
func = mod.get_function("test_cuda_LongIntArray")
# 传递参数时确保类型与内核定义匹配
func(result_count_gpu, record_result_gpu, block=(4,16,16), grid=(1,1))

# 从设备端拷贝回最终计数器值
cuda.memcpy_dtoh(result_count, result_count_gpu)
print("最终计数器值:", result_count[0])

# 释放设备内存
record_result_gpu.free()
result_count_gpu.free()

关键修改说明 #

• 用int64_t（CUDA标准64位整数类型）替代long，确保跨平台类型一致，与主机端np.int64严格匹配
• 内核仅保留全局内存指针参数int64_t *result_count，原子操作必须作用在全局内存上，舍弃传值参数InitianCount
• 调用atomicAdd时用1LL作为增量，明确指定为64位整数常量
• 给SourceModule添加no_extern_c=True参数，解决CUDA内核函数的C链接兼容性问题

内容的提问来源于stack exchange，提问作者Henry