使用CuPy.nanargmax处理GPU数组时触发编译异常的求助

问题原因 #

你遇到的编译错误是因为CuPy调用的CUDA Cooperative Groups API（cudaCGGetIntrinsicHandle等函数）在当前CUDA版本中未定义，本质是CuPy版本与CUDA版本不兼容——较新的CuPy版本依赖CUDA 11.0+的新API，而你的CUDA版本过旧，无法识别这些函数。

解决方案 #

1. 匹配CuPy与CUDA版本 #

不同CuPy版本对应不同的CUDA支持范围：

• 若你的CUDA版本≤10.2，需安装CuPy v9.x及以下版本（例如执行pip install cupy-cuda102==9.6.0）
• 若你的CUDA版本≥11.0，可安装对应版本的CuPy（例如CUDA 11.7对应pip install cupy-cuda117）

2. 临时替代实现（无需版本变更） #

如果暂时无法调整版本，可通过替换NaN为对应数据类型的最小值后调用cp.argmax实现等效功能：

import cupy as cp
import numpy as np

# 创建测试数组
test_array_gpu = cp.array(np.random.randn(1250, 552))
test_array_gpu[200, 100] = cp.nan
test_array_gpu[500, 200] = cp.nan
test_array_gpu[1000, 300] = cp.nan

# 替换NaN为当前数据类型的最小值，避免干扰argmax结果
dtype_min = cp.finfo(test_array_gpu.dtype).min
array_no_nan = cp.where(cp.isnan(test_array_gpu), dtype_min, test_array_gpu)

# 调用argmax获取每行最大值索引（效果等同于nanargmax）
best_fit_indices_gpu = cp.argmax(array_no_nan, axis=1)

print("Best fit indices (indices of max values, ignoring NaNs):")
print(best_fit_indices_gpu)

这个方法绕过了需要JIT编译的cp.nanargmax，直接用原生CuPy函数实现，不会触发编译错误，且性能接近原生nanargmax。

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